-- | This module implements the translation from the multi-tick
-- calculus to the single tick calculus.

{-# LANGUAGE CPP #-}

module AsyncRattus.Plugin.SingleTick
  (toSingleTick) where

#if __GLASGOW_HASKELL__ >= 900
import GHC.Plugins
#else
import GhcPlugins
#endif

  
import AsyncRattus.Plugin.Utils
import Prelude hiding ((<>))
import Control.Monad.Trans.Writer.Strict
import Control.Monad.Trans.Class
import Data.List

-- | Transform the given expression from the multi-tick calculus into
-- the single tick calculus form.
toSingleTick :: CoreExpr -> CoreM CoreExpr
toSingleTick :: CoreExpr -> CoreM CoreExpr
toSingleTick (Let (Rec [(Id, CoreExpr)]
bs) CoreExpr
e) = do
  CoreExpr
e' <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
  [(Id, CoreExpr)]
bs' <- ((Id, CoreExpr) -> CoreM (Id, CoreExpr))
-> [(Id, CoreExpr)] -> CoreM [(Id, CoreExpr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((CoreExpr -> CoreM CoreExpr)
-> (Id, CoreExpr) -> CoreM (Id, CoreExpr)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> (Id, a) -> m (Id, b)
mapM CoreExpr -> CoreM CoreExpr
toSingleTick) [(Id, CoreExpr)]
bs
  CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bind Id -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let ([(Id, CoreExpr)] -> Bind Id
forall b. [(b, Expr b)] -> Bind b
Rec [(Id, CoreExpr)]
bs') CoreExpr
e')
toSingleTick (Let (NonRec Id
b CoreExpr
e1) CoreExpr
e2) = do
  CoreExpr
e1' <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e1
  CoreExpr
e2' <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e2
  CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bind Id -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let (Id -> CoreExpr -> Bind Id
forall b. b -> Expr b -> Bind b
NonRec Id
b CoreExpr
e1') CoreExpr
e2')
toSingleTick (Case CoreExpr
e Id
b Type
ty [Alt Id]
alts) = do
  CoreExpr
e' <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
  [Alt Id]
alts' <- (Alt Id -> CoreM (Alt Id)) -> [Alt Id] -> CoreM [Alt Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((\ (AltCon
c,[Id]
bs,CoreExpr
f) -> (CoreExpr -> Alt Id) -> CoreM CoreExpr -> CoreM (Alt Id)
forall a b. (a -> b) -> CoreM a -> CoreM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ CoreExpr
x -> AltCon -> [Id] -> CoreExpr -> Alt Id
forall {b}. AltCon -> [b] -> Expr b -> Alt b
mkAlt AltCon
c [Id]
bs CoreExpr
x) (CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
f)) ((AltCon, [Id], CoreExpr) -> CoreM (Alt Id))
-> (Alt Id -> (AltCon, [Id], CoreExpr)) -> Alt Id -> CoreM (Alt Id)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alt Id -> (AltCon, [Id], CoreExpr)
forall {b}. Alt b -> (AltCon, [b], Expr b)
getAlt) [Alt Id]
alts
  CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case CoreExpr
e' Id
b Type
ty [Alt Id]
alts')
toSingleTick (Cast CoreExpr
e CoercionR
c) = do
  CoreExpr
e' <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
  CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoercionR -> CoreExpr
forall b. Expr b -> CoercionR -> Expr b
Cast CoreExpr
e' CoercionR
c)
toSingleTick (Tick CoreTickish
t CoreExpr
e) = do
  CoreExpr
e' <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
  CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreTickish -> CoreExpr -> CoreExpr
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t CoreExpr
e')
toSingleTick (Lam Id
x CoreExpr
e) = do
  (CoreExpr
e', [(Id, CoreExpr, CoreExpr)]
advs) <- WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
-> CoreM (CoreExpr, [(Id, CoreExpr, CoreExpr)])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e)
  [(Id, CoreExpr, CoreExpr)]
advs' <- ((Id, CoreExpr, CoreExpr) -> CoreM (Id, CoreExpr, CoreExpr))
-> [(Id, CoreExpr, CoreExpr)] -> CoreM [(Id, CoreExpr, CoreExpr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\ (Id
x,CoreExpr
a,CoreExpr
b) -> (CoreExpr -> (Id, CoreExpr, CoreExpr))
-> CoreM CoreExpr -> CoreM (Id, CoreExpr, CoreExpr)
forall a b. (a -> b) -> CoreM a -> CoreM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\CoreExpr
b' -> (Id
x,CoreExpr
a,CoreExpr
b')) (CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
b)) [(Id, CoreExpr, CoreExpr)]
advs
  CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets' [(Id, CoreExpr, CoreExpr)]
advs' (Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
x CoreExpr
e'))
toSingleTick (App CoreExpr
e1 CoreExpr
e2)
  | CoreExpr -> Bool
isDelayApp CoreExpr
e1 = do
      (CoreExpr
e2', [(Id, CoreExpr)]
advs) <- WriterT [(Id, CoreExpr)] CoreM CoreExpr
-> CoreM (CoreExpr, [(Id, CoreExpr)])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e2)
      [(Id, CoreExpr)]
advs' <- ((Id, CoreExpr) -> CoreM (Id, CoreExpr))
-> [(Id, CoreExpr)] -> CoreM [(Id, CoreExpr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((CoreExpr -> CoreM CoreExpr)
-> (Id, CoreExpr) -> CoreM (Id, CoreExpr)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> (Id, a) -> m (Id, b)
mapM CoreExpr -> CoreM CoreExpr
toSingleTick) [(Id, CoreExpr)]
advs
      CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
advs' (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 CoreExpr
e2'))
  | Bool
otherwise = do
      CoreExpr
e1' <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e1
      CoreExpr
e2' <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e2
      CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1' CoreExpr
e2')

toSingleTick e :: CoreExpr
e@Type{} = CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
toSingleTick e :: CoreExpr
e@Var{} = CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
toSingleTick e :: CoreExpr
e@Lit{} = CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
toSingleTick e :: CoreExpr
e@Coercion{} = CoreExpr -> CoreM CoreExpr
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e

foldLets :: [(Id,CoreExpr)] -> CoreExpr -> CoreExpr
foldLets :: [(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
ls CoreExpr
e = (CoreExpr -> (Id, CoreExpr) -> CoreExpr)
-> CoreExpr -> [(Id, CoreExpr)] -> CoreExpr
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\CoreExpr
e' (Id
x,CoreExpr
b) -> Bind Id -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let (Id -> CoreExpr -> Bind Id
forall b. b -> Expr b -> Bind b
NonRec Id
x CoreExpr
b) CoreExpr
e') CoreExpr
e [(Id, CoreExpr)]
ls

foldLets' :: [(Id,CoreExpr,CoreExpr)] -> CoreExpr -> CoreExpr
foldLets' :: [(Id, CoreExpr, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets' [(Id, CoreExpr, CoreExpr)]
ls CoreExpr
e = (CoreExpr -> (Id, CoreExpr, CoreExpr) -> CoreExpr)
-> CoreExpr -> [(Id, CoreExpr, CoreExpr)] -> CoreExpr
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\CoreExpr
e' (Id
x,CoreExpr
a,CoreExpr
b) -> Bind Id -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let (Id -> CoreExpr -> Bind Id
forall b. b -> Expr b -> Bind b
NonRec Id
x (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
a CoreExpr
b)) CoreExpr
e') CoreExpr
e [(Id, CoreExpr, CoreExpr)]
ls

extractAdvApp :: CoreExpr -> CoreExpr -> WriterT [(Id,CoreExpr)] CoreM CoreExpr
extractAdvApp :: CoreExpr -> CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdvApp CoreExpr
e1 CoreExpr
e2
  | CoreExpr -> Bool
isVar CoreExpr
e2 = CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 CoreExpr
e2)
  | Bool
otherwise = do
  Id
x <- CoreM Id -> WriterT [(Id, CoreExpr)] CoreM Id
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (FastString -> CoreExpr -> CoreM Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> CoreExpr -> m Id
mkSysLocalFromExpr (String -> FastString
fsLit String
"adv") CoreExpr
e2)
  [(Id, CoreExpr)] -> WriterT [(Id, CoreExpr)] CoreM ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Id
x,CoreExpr
e2)]
  CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
x))

-- removes casts and ticks from a tree
filterTree :: CoreExpr -> CoreExpr
filterTree :: CoreExpr -> CoreExpr
filterTree (Cast CoreExpr
e CoercionR
_) = CoreExpr -> CoreExpr
filterTree CoreExpr
e
filterTree (Tick CoreTickish
_ CoreExpr
e) = CoreExpr -> CoreExpr
filterTree CoreExpr
e
filterTree CoreExpr
e = CoreExpr
e


extractSelectApp :: CoreExpr -> CoreExpr -> WriterT [(Id,CoreExpr)] CoreM CoreExpr
extractSelectApp :: CoreExpr -> CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractSelectApp CoreExpr
e1 CoreExpr
e2
  | CoreExpr -> Bool
isVar CoreExpr
e' Bool -> Bool -> Bool
&& CoreExpr -> Bool
isVar CoreExpr
e2 = CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 CoreExpr
e2)
  | CoreExpr -> Bool
isVar CoreExpr
e2 = do
    Id
x <- CoreM Id -> WriterT [(Id, CoreExpr)] CoreM Id
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (FastString -> CoreExpr -> CoreM Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> CoreExpr -> m Id
mkSysLocalFromExpr (String -> FastString
fsLit String
"selectFreshVar") CoreExpr
e')
    [(Id, CoreExpr)] -> WriterT [(Id, CoreExpr)] CoreM ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Id
x, CoreExpr
e')]
    CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
x)) CoreExpr
e2)
  | CoreExpr -> Bool
isVar CoreExpr
e' = do
    Id
x <- CoreM Id -> WriterT [(Id, CoreExpr)] CoreM Id
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (FastString -> CoreExpr -> CoreM Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> CoreExpr -> m Id
mkSysLocalFromExpr (String -> FastString
fsLit String
"selectFreshVar") CoreExpr
e2)
    [(Id, CoreExpr)] -> WriterT [(Id, CoreExpr)] CoreM ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Id
x, CoreExpr
e2)]
    CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
x))
  | Bool
otherwise = do
    Id
x <- CoreM Id -> WriterT [(Id, CoreExpr)] CoreM Id
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (FastString -> CoreExpr -> CoreM Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> CoreExpr -> m Id
mkSysLocalFromExpr (String -> FastString
fsLit String
"selectFreshVar") CoreExpr
e')
    Id
y <- CoreM Id -> WriterT [(Id, CoreExpr)] CoreM Id
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (FastString -> CoreExpr -> CoreM Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> CoreExpr -> m Id
mkSysLocalFromExpr (String -> FastString
fsLit String
"selectFreshVar") CoreExpr
e2)
    [(Id, CoreExpr)] -> WriterT [(Id, CoreExpr)] CoreM ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Id
x, CoreExpr
e')]
    [(Id, CoreExpr)] -> WriterT [(Id, CoreExpr)] CoreM ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Id
y, CoreExpr
e2)]
    CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
x)) (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
y))
  where (App CoreExpr
e CoreExpr
e') = CoreExpr -> CoreExpr
filterTree CoreExpr
e1


-- This is used to pull adv out of delayed terms. The writer monad
-- returns mappings from fresh variables to terms that occur as
-- argument of adv.
-- 
-- That is, occurrences of @adv t@ are replaced with @adv x@ (for some
-- fresh variable @x@) and the pair @(x,t)@ is returned in the
-- writer monad.
extractAdv :: CoreExpr -> WriterT [(Id,CoreExpr)] CoreM CoreExpr
extractAdv :: CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv (App expr :: CoreExpr
expr@(App CoreExpr
e CoreExpr
_) CoreExpr
e2) | CoreExpr -> Bool
isSelectApp CoreExpr
e = CoreExpr -> CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractSelectApp CoreExpr
expr CoreExpr
e2
extractAdv e :: CoreExpr
e@(App CoreExpr
e1 CoreExpr
e2)
  | CoreExpr -> Bool
isAdvApp CoreExpr
e1 = CoreExpr -> CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdvApp CoreExpr
e1 CoreExpr
e2
  | CoreExpr -> Bool
isSelectApp CoreExpr
e1 = CoreExpr -> CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractSelectApp CoreExpr
e1 CoreExpr
e2
  | CoreExpr -> Bool
isDelayApp CoreExpr
e1 = do
      (CoreExpr
e2', [(Id, CoreExpr)]
advs) <- CoreM (CoreExpr, [(Id, CoreExpr)])
-> WriterT [(Id, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr)])
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (CoreM (CoreExpr, [(Id, CoreExpr)])
 -> WriterT [(Id, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr)]))
-> CoreM (CoreExpr, [(Id, CoreExpr)])
-> WriterT [(Id, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr)])
forall a b. (a -> b) -> a -> b
$ WriterT [(Id, CoreExpr)] CoreM CoreExpr
-> CoreM (CoreExpr, [(Id, CoreExpr)])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e2)
      [(Id, CoreExpr)]
advs' <- ((Id, CoreExpr) -> WriterT [(Id, CoreExpr)] CoreM (Id, CoreExpr))
-> [(Id, CoreExpr)]
-> WriterT [(Id, CoreExpr)] CoreM [(Id, CoreExpr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr)
-> (Id, CoreExpr) -> WriterT [(Id, CoreExpr)] CoreM (Id, CoreExpr)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> (Id, a) -> m (Id, b)
mapM CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv) [(Id, CoreExpr)]
advs
      CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
advs' (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 CoreExpr
e2'))
  | CoreExpr -> Bool
isBoxApp CoreExpr
e1 = CoreM CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (CoreM CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr)
-> CoreM CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a b. (a -> b) -> a -> b
$ CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
  | Bool
otherwise = do
      CoreExpr
e1' <- CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e1
      CoreExpr
e2' <- CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e2
      CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1' CoreExpr
e2')
extractAdv (Lam Id
x CoreExpr
e) = do
  (CoreExpr
e', [(Id, CoreExpr, CoreExpr)]
advs) <- CoreM (CoreExpr, [(Id, CoreExpr, CoreExpr)])
-> WriterT
     [(Id, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr, CoreExpr)])
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (CoreM (CoreExpr, [(Id, CoreExpr, CoreExpr)])
 -> WriterT
      [(Id, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr, CoreExpr)]))
-> CoreM (CoreExpr, [(Id, CoreExpr, CoreExpr)])
-> WriterT
     [(Id, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr, CoreExpr)])
forall a b. (a -> b) -> a -> b
$ WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
-> CoreM (CoreExpr, [(Id, CoreExpr, CoreExpr)])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e)
  [(Id, CoreExpr)]
advs' <- ((Id, CoreExpr, CoreExpr)
 -> WriterT [(Id, CoreExpr)] CoreM (Id, CoreExpr))
-> [(Id, CoreExpr, CoreExpr)]
-> WriterT [(Id, CoreExpr)] CoreM [(Id, CoreExpr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\ (Id
x,CoreExpr
a,CoreExpr
b) -> (CoreExpr -> (Id, CoreExpr))
-> WriterT [(Id, CoreExpr)] CoreM CoreExpr
-> WriterT [(Id, CoreExpr)] CoreM (Id, CoreExpr)
forall a b.
(a -> b)
-> WriterT [(Id, CoreExpr)] CoreM a
-> WriterT [(Id, CoreExpr)] CoreM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\CoreExpr
b' -> (Id
x,CoreExpr
b')) (CoreExpr -> CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdvApp CoreExpr
a CoreExpr
b)) [(Id, CoreExpr, CoreExpr)]
advs
  CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
advs' (Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
x CoreExpr
e'))
extractAdv (Case CoreExpr
e Id
b Type
ty [Alt Id]
alts) = do
  CoreExpr
e' <- CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e
  [Alt Id]
alts' <- (Alt Id -> WriterT [(Id, CoreExpr)] CoreM (Alt Id))
-> [Alt Id] -> WriterT [(Id, CoreExpr)] CoreM [Alt Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((\ (AltCon
c,[Id]
bs,CoreExpr
f) -> (CoreExpr -> Alt Id)
-> WriterT [(Id, CoreExpr)] CoreM CoreExpr
-> WriterT [(Id, CoreExpr)] CoreM (Alt Id)
forall a b.
(a -> b)
-> WriterT [(Id, CoreExpr)] CoreM a
-> WriterT [(Id, CoreExpr)] CoreM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ CoreExpr
x -> AltCon -> [Id] -> CoreExpr -> Alt Id
forall {b}. AltCon -> [b] -> Expr b -> Alt b
mkAlt AltCon
c [Id]
bs CoreExpr
x) (CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
f)) ((AltCon, [Id], CoreExpr)
 -> WriterT [(Id, CoreExpr)] CoreM (Alt Id))
-> (Alt Id -> (AltCon, [Id], CoreExpr))
-> Alt Id
-> WriterT [(Id, CoreExpr)] CoreM (Alt Id)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alt Id -> (AltCon, [Id], CoreExpr)
forall {b}. Alt b -> (AltCon, [b], Expr b)
getAlt) [Alt Id]
alts
  CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case CoreExpr
e' Id
b Type
ty [Alt Id]
alts')
extractAdv (Cast CoreExpr
e CoercionR
c) = do
  CoreExpr
e' <- CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e
  CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoercionR -> CoreExpr
forall b. Expr b -> CoercionR -> Expr b
Cast CoreExpr
e' CoercionR
c)
extractAdv (Tick CoreTickish
t CoreExpr
e) = do
  CoreExpr
e' <- CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e
  CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreTickish -> CoreExpr -> CoreExpr
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t CoreExpr
e')
extractAdv e :: CoreExpr
e@(Let Rec{} CoreExpr
_) = CoreM CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (CoreM CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr)
-> CoreM CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a b. (a -> b) -> a -> b
$ CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
extractAdv (Let (NonRec Id
b CoreExpr
e1) CoreExpr
e2) = do
  CoreExpr
e1' <- CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e1
  CoreExpr
e2' <- CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e2
  CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bind Id -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let (Id -> CoreExpr -> Bind Id
forall b. b -> Expr b -> Bind b
NonRec Id
b CoreExpr
e1') CoreExpr
e2')
extractAdv e :: CoreExpr
e@Type{} = CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv e :: CoreExpr
e@Var{} = CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv e :: CoreExpr
e@Lit{} = CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv e :: CoreExpr
e@Coercion{} = CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e

-- This is used to pull adv out of lambdas. The writer monad returns
-- mappings from fresh variables to occurrences of adv and the term it
-- is applied to.
-- 
-- That is occurrences of @adv t@ are replaced with a fresh variable
-- @x@ and the triple @(x,adv,t)@ is returned in the writer monad.
-- For select a b, the triple @(x, select a, b) is returned in the writer monad.
extractAdv' :: CoreExpr -> WriterT [(Id,CoreExpr,CoreExpr)] CoreM CoreExpr
extractAdv' :: CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' e :: CoreExpr
e@(App CoreExpr
e1 CoreExpr
e2)
  | CoreExpr -> Bool
isAdvApp CoreExpr
e1 = do
       Id
x <- CoreM Id -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM Id
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (FastString -> CoreExpr -> CoreM Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> CoreExpr -> m Id
mkSysLocalFromExpr (String -> FastString
fsLit String
"adv") CoreExpr
e)
       [(Id, CoreExpr, CoreExpr)]
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Id
x,CoreExpr
e1,CoreExpr
e2)]
       CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
x)
  | CoreExpr -> Bool
isSelectApp CoreExpr
e1 = do
      Id
x <- CoreM Id -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM Id
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (FastString -> CoreExpr -> CoreM Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> CoreExpr -> m Id
mkSysLocalFromExpr (String -> FastString
fsLit String
"select") CoreExpr
e)
      [(Id, CoreExpr, CoreExpr)]
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Id
x,CoreExpr
e1,CoreExpr
e2)]
      CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
x)
  | CoreExpr -> Bool
isDelayApp CoreExpr
e1 = do
      (CoreExpr
e2', [(Id, CoreExpr)]
advs) <- CoreM (CoreExpr, [(Id, CoreExpr)])
-> WriterT
     [(Id, CoreExpr, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr)])
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (CoreM (CoreExpr, [(Id, CoreExpr)])
 -> WriterT
      [(Id, CoreExpr, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr)]))
-> CoreM (CoreExpr, [(Id, CoreExpr)])
-> WriterT
     [(Id, CoreExpr, CoreExpr)] CoreM (CoreExpr, [(Id, CoreExpr)])
forall a b. (a -> b) -> a -> b
$ WriterT [(Id, CoreExpr)] CoreM CoreExpr
-> CoreM (CoreExpr, [(Id, CoreExpr)])
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e2)
      [(Id, CoreExpr)]
advs' <- ((Id, CoreExpr)
 -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM (Id, CoreExpr))
-> [(Id, CoreExpr)]
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM [(Id, CoreExpr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr)
-> (Id, CoreExpr)
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM (Id, CoreExpr)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> (Id, a) -> m (Id, b)
mapM CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv') [(Id, CoreExpr)]
advs
      CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
advs' (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 CoreExpr
e2'))
  | CoreExpr -> Bool
isBoxApp CoreExpr
e1 = CoreM CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (CoreM CoreExpr
 -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr)
-> CoreM CoreExpr
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a b. (a -> b) -> a -> b
$ CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
  | Bool
otherwise = do
      CoreExpr
e1' <- CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e1
      CoreExpr
e2' <- CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e2
      CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1' CoreExpr
e2')
extractAdv' (Lam Id
x CoreExpr
e) = do
  CoreExpr
e' <- CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e
  CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
x CoreExpr
e')
extractAdv' (Case CoreExpr
e Id
b Type
ty [Alt Id]
alts) = do
  CoreExpr
e' <- CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e
  [Alt Id]
alts' <- (Alt Id -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM (Alt Id))
-> [Alt Id] -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM [Alt Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((\ (AltCon
c,[Id]
bs,CoreExpr
f) -> (CoreExpr -> Alt Id)
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM (Alt Id)
forall a b.
(a -> b)
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ CoreExpr
x -> AltCon -> [Id] -> CoreExpr -> Alt Id
forall {b}. AltCon -> [b] -> Expr b -> Alt b
mkAlt AltCon
c [Id]
bs CoreExpr
x) (CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
f)) ((AltCon, [Id], CoreExpr)
 -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM (Alt Id))
-> (Alt Id -> (AltCon, [Id], CoreExpr))
-> Alt Id
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM (Alt Id)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alt Id -> (AltCon, [Id], CoreExpr)
forall {b}. Alt b -> (AltCon, [b], Expr b)
getAlt) [Alt Id]
alts
  CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case CoreExpr
e' Id
b Type
ty [Alt Id]
alts')
extractAdv' (Cast CoreExpr
e CoercionR
c) = do
  CoreExpr
e' <- CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e
  CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr -> CoercionR -> CoreExpr
forall b. Expr b -> CoercionR -> Expr b
Cast CoreExpr
e' CoercionR
c)
extractAdv' (Tick CoreTickish
t CoreExpr
e) = do
  CoreExpr
e' <- CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e
  CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreTickish -> CoreExpr -> CoreExpr
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t CoreExpr
e')
extractAdv' e :: CoreExpr
e@(Let Rec{} CoreExpr
_) = CoreM CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall (m :: * -> *) a.
Monad m =>
m a -> WriterT [(Id, CoreExpr, CoreExpr)] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (CoreM CoreExpr
 -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr)
-> CoreM CoreExpr
-> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a b. (a -> b) -> a -> b
$ CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
extractAdv' (Let (NonRec Id
b CoreExpr
e1) CoreExpr
e2) = do
  CoreExpr
e1' <- CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e1
  CoreExpr
e2' <- CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv' CoreExpr
e2
  CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bind Id -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let (Id -> CoreExpr -> Bind Id
forall b. b -> Expr b -> Bind b
NonRec Id
b CoreExpr
e1') CoreExpr
e2')
extractAdv' e :: CoreExpr
e@Type{} = CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv' e :: CoreExpr
e@Var{} = CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv' e :: CoreExpr
e@Lit{} = CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv' e :: CoreExpr
e@Coercion{} = CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
forall a. a -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e



isDelayApp :: CoreExpr -> Bool
isDelayApp :: CoreExpr -> Bool
isDelayApp = (String -> Bool) -> CoreExpr -> Bool
isPrimApp (String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"delay")

isBoxApp :: CoreExpr -> Bool
isBoxApp :: CoreExpr -> Bool
isBoxApp = (String -> Bool) -> CoreExpr -> Bool
isPrimApp (\String
occ -> String
occ String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"Box" Bool -> Bool -> Bool
|| String
occ String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"box")

isAdvApp :: CoreExpr -> Bool
isAdvApp :: CoreExpr -> Bool
isAdvApp = (String -> Bool) -> CoreExpr -> Bool
isPrimApp (String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"adv")

isSelectApp :: CoreExpr -> Bool
isSelectApp :: CoreExpr -> Bool
isSelectApp = (String -> Bool) -> CoreExpr -> Bool
isPrimApp (String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"select")

isPrimApp :: (String -> Bool) -> CoreExpr -> Bool
isPrimApp :: (String -> Bool) -> CoreExpr -> Bool
isPrimApp String -> Bool
p (App CoreExpr
e CoreExpr
e')
  | CoreExpr -> Bool
forall {b}. Expr b -> Bool
isType CoreExpr
e' Bool -> Bool -> Bool
|| Bool -> Bool
not  (Type -> Bool
tcIsLiftedTypeKind((() :: Constraint) => Type -> Type
Type -> Type
typeKind ((() :: Constraint) => CoreExpr -> Type
CoreExpr -> Type
exprType CoreExpr
e'))) = (String -> Bool) -> CoreExpr -> Bool
isPrimApp String -> Bool
p CoreExpr
e
  | Bool
otherwise = Bool
False
isPrimApp String -> Bool
p (Cast CoreExpr
e CoercionR
_) = (String -> Bool) -> CoreExpr -> Bool
isPrimApp String -> Bool
p CoreExpr
e
isPrimApp String -> Bool
p (Tick CoreTickish
_ CoreExpr
e) = (String -> Bool) -> CoreExpr -> Bool
isPrimApp String -> Bool
p CoreExpr
e
isPrimApp String -> Bool
p (Var Id
v) = (String -> Bool) -> Id -> Bool
isPrimVar String -> Bool
p Id
v
isPrimApp String -> Bool
_ CoreExpr
_ = Bool
False

isPrimVar :: (String -> Bool) -> Var -> Bool
isPrimVar :: (String -> Bool) -> Id -> Bool
isPrimVar String -> Bool
p Id
v = Bool -> (Bool -> Bool) -> Maybe Bool -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Bool -> Bool
forall a. a -> a
id (Maybe Bool -> Bool) -> Maybe Bool -> Bool
forall a b. (a -> b) -> a -> b
$ do
  let name :: Name
name = Id -> Name
varName Id
v
  Module
mod <- Name -> Maybe Module
nameModule_maybe Name
name
  let occ :: String
occ = Name -> String
forall a. NamedThing a => a -> String
getOccString Name
name
  Bool -> Maybe Bool
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Bool
p String
occ
          Bool -> Bool -> Bool
&& ModuleName -> String
moduleNameString (Module -> ModuleName
forall unit. GenModule unit -> ModuleName
moduleName Module
mod) String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"AsyncRattus.InternalPrimitives")