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

{-# LANGUAGE CPP #-}

module Rattus.Plugin.SingleTick
  (toSingleTick) where

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

  
import Rattus.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' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CoreExpr -> CoreM CoreExpr
toSingleTick) [(Id, CoreExpr)]
bs
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. Bind b -> Expr b -> Expr b
Let (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
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. Bind b -> Expr b -> Expr b
Let (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' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((\ (AltCon
c,[Id]
bs,CoreExpr
f) -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ CoreExpr
x -> forall {b}. AltCon -> [b] -> Expr b -> Alt b
mkAlt AltCon
c [Id]
bs CoreExpr
x) (CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
f)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {b}. Alt b -> (AltCon, [b], Expr b)
getAlt) [Alt Id]
alts
  forall (m :: * -> *) a. Monad m => a -> m a
return (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
  forall (m :: * -> *) a. Monad m => a -> m a
return (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
  forall (m :: * -> *) a. Monad m => a -> m a
return (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) <- 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' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\ (Id
x,CoreExpr
a,CoreExpr
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
  forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets' [(Id, CoreExpr, CoreExpr)]
advs' (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) <- forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e2)
      [(Id, CoreExpr)]
advs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CoreExpr -> CoreM CoreExpr
toSingleTick) [(Id, CoreExpr)]
advs
      forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
advs' (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
      forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1' CoreExpr
e2')

toSingleTick e :: CoreExpr
e@Type{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
toSingleTick e :: CoreExpr
e@Var{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
toSingleTick e :: CoreExpr
e@Lit{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
toSingleTick e :: CoreExpr
e@Coercion{} = 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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\CoreExpr
e' (Id
x,CoreExpr
b) -> forall b. Bind b -> Expr b -> Expr b
Let (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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\CoreExpr
e' (Id
x,CoreExpr
a,CoreExpr
b) -> forall b. Bind b -> Expr b -> Expr b
Let (forall b. b -> Expr b -> Bind b
NonRec Id
x (forall b. Expr b -> Expr b -> Expr b
App CoreExpr
a CoreExpr
b)) CoreExpr
e') CoreExpr
e [(Id, CoreExpr, CoreExpr)]
ls

isVar :: CoreExpr -> Bool
isVar :: CoreExpr -> Bool
isVar (App CoreExpr
e CoreExpr
e')
  | forall {b}. Expr b -> Bool
isType CoreExpr
e' Bool -> Bool -> Bool
|| Bool -> Bool
not  (Type -> Bool
tcIsLiftedTypeKind(HasDebugCallStack => Type -> Type
typeKind (CoreExpr -> Type
exprType CoreExpr
e'))) = CoreExpr -> Bool
isVar CoreExpr
e
  | Bool
otherwise = Bool
False
isVar (Cast CoreExpr
e CoercionR
_) = CoreExpr -> Bool
isVar CoreExpr
e
isVar (Tick CoreTickish
_ CoreExpr
e) = CoreExpr -> Bool
isVar CoreExpr
e
isVar (Var Id
_) = Bool
True
isVar CoreExpr
_ = Bool
False


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

-- 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 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
isDelayApp CoreExpr
e1 = do
      (CoreExpr
e2', [(Id, CoreExpr)]
advs) <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e2)
      [(Id, CoreExpr)]
advs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv) [(Id, CoreExpr)]
advs
      forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
advs' (forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 CoreExpr
e2'))
  | CoreExpr -> Bool
isBoxApp CoreExpr
e1 = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift 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
      forall (m :: * -> *) a. Monad m => a -> m a
return (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) <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ 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' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\ (Id
x,CoreExpr
a,CoreExpr
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
  forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
advs' (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' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((\ (AltCon
c,[Id]
bs,CoreExpr
f) -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ CoreExpr
x -> forall {b}. AltCon -> [b] -> Expr b -> Alt b
mkAlt AltCon
c [Id]
bs CoreExpr
x) (CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
f)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {b}. Alt b -> (AltCon, [b], Expr b)
getAlt) [Alt Id]
alts
  forall (m :: * -> *) a. Monad m => a -> m a
return (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
  forall (m :: * -> *) a. Monad m => a -> m a
return (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
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t CoreExpr
e')
extractAdv e :: CoreExpr
e@(Let Rec{} CoreExpr
_) = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift 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
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. Bind b -> Expr b -> Expr b
Let (forall b. b -> Expr b -> Bind b
NonRec Id
b CoreExpr
e1') CoreExpr
e2')
extractAdv e :: CoreExpr
e@Type{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv e :: CoreExpr
e@Var{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv e :: CoreExpr
e@Lit{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv e :: CoreExpr
e@Coercion{} = 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.
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 <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (forall (m :: * -> *).
MonadUnique m =>
FastString -> CoreExpr -> m Id
mkSysLocalFromExpr (String -> FastString
fsLit String
"adv") CoreExpr
e)
       forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell [(Id
x,CoreExpr
e1,CoreExpr
e2)]
       forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. Id -> Expr b
Var Id
x)
  | CoreExpr -> Bool
isDelayApp CoreExpr
e1 = do
      (CoreExpr
e2', [(Id, CoreExpr)]
advs) <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (CoreExpr -> WriterT [(Id, CoreExpr)] CoreM CoreExpr
extractAdv CoreExpr
e2)
      [(Id, CoreExpr)]
advs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM CoreExpr -> WriterT [(Id, CoreExpr, CoreExpr)] CoreM CoreExpr
extractAdv') [(Id, CoreExpr)]
advs
      forall (m :: * -> *) a. Monad m => a -> m a
return ([(Id, CoreExpr)] -> CoreExpr -> CoreExpr
foldLets [(Id, CoreExpr)]
advs' (forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1 CoreExpr
e2'))
  | CoreExpr -> Bool
isBoxApp CoreExpr
e1 = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift 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
      forall (m :: * -> *) a. Monad m => a -> m a
return (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
  forall (m :: * -> *) a. Monad m => a -> m a
return (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' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((\ (AltCon
c,[Id]
bs,CoreExpr
f) -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ CoreExpr
x -> 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)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {b}. Alt b -> (AltCon, [b], Expr b)
getAlt) [Alt Id]
alts
  forall (m :: * -> *) a. Monad m => a -> m a
return (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
  forall (m :: * -> *) a. Monad m => a -> m a
return (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
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t CoreExpr
e')
extractAdv' e :: CoreExpr
e@(Let Rec{} CoreExpr
_) = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift 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
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. Bind b -> Expr b -> Expr b
Let (forall b. b -> Expr b -> Bind b
NonRec Id
b CoreExpr
e1') CoreExpr
e2')
extractAdv' e :: CoreExpr
e@Type{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv' e :: CoreExpr
e@Var{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv' e :: CoreExpr
e@Lit{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
extractAdv' e :: CoreExpr
e@Coercion{} = forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e



isDelayApp :: CoreExpr -> Bool
isDelayApp :: CoreExpr -> Bool
isDelayApp = (String -> Bool) -> CoreExpr -> Bool
isPrimApp (\String
occ -> String
occ 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 forall a. Eq a => a -> a -> Bool
== String
"Box" Bool -> Bool -> Bool
|| String
occ forall a. Eq a => a -> a -> Bool
== String
"box")

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


isPrimApp :: (String -> Bool) -> CoreExpr -> Bool
isPrimApp :: (String -> Bool) -> CoreExpr -> Bool
isPrimApp String -> Bool
p (App CoreExpr
e CoreExpr
e')
  | forall {b}. Expr b -> Bool
isType CoreExpr
e' Bool -> Bool -> Bool
|| Bool -> Bool
not  (Type -> Bool
tcIsLiftedTypeKind(HasDebugCallStack => Type -> Type
typeKind (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 = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False forall a. a -> a
id 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 = forall a. NamedThing a => a -> String
getOccString Name
name
  forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Bool
p String
occ
          Bool -> Bool -> Bool
&& ((ModuleName -> String
moduleNameString (forall unit. GenModule unit -> ModuleName
moduleName Module
mod) forall a. Eq a => a -> a -> Bool
== String
"Rattus.Internal") Bool -> Bool -> Bool
||
          ModuleName -> String
moduleNameString (forall unit. GenModule unit -> ModuleName
moduleName Module
mod) forall a. Eq a => a -> a -> Bool
== String
"Rattus.Primitives"))