{-# LANGUAGE TupleSections #-}
module AsyncRattus.Plugin.Transform (
transform
) where
import GHC.Core.Opt.Monad
import GHC.Plugins
import AsyncRattus.Plugin.PrimExpr
import AsyncRattus.Plugin.Utils
import Data.Maybe (fromJust)
import Prelude hiding ((<>))
import Data.Functor ((<&>))
import Control.Applicative ((<|>))
import Data.Tuple (swap)
data Ctx = Ctx {
Ctx -> Maybe Id
fresh :: Maybe Var
}
emptyCtx :: Ctx
emptyCtx :: Ctx
emptyCtx = Ctx {
fresh :: Maybe Id
fresh = Maybe Id
forall a. Maybe a
Nothing
}
replaceVar :: Var -> Var -> Expr Var -> Expr Var
replaceVar :: Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep (Var Id
v) = if Id
v Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
match then Id -> Expr Id
forall b. Id -> Expr b
Var Id
rep else Id -> Expr Id
forall b. Id -> Expr b
Var Id
v
replaceVar Id
match Id
rep (App Expr Id
e Expr Id
e') = Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
e) (Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
e')
replaceVar Id
match Id
rep (Tick CoreTickish
_ Expr Id
e) = Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
e
replaceVar Id
match Id
rep (Lam Id
v Expr Id
e) = Id -> Expr Id -> Expr Id
forall b. b -> Expr b -> Expr b
Lam (if Id
v Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
match then Id
rep else Id
v) (Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
e)
replaceVar Id
match Id
rep (Let (NonRec Id
b Expr Id
e') Expr Id
e) =
Bind Id -> Expr Id -> Expr Id
forall b. Bind b -> Expr b -> Expr b
Let (Id -> Expr Id -> Bind Id
forall b. b -> Expr b -> Bind b
NonRec Id
newB (Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
e')) (Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
e)
where newB :: Id
newB = if Id
b Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
match then Id
rep else Id
b
replaceVar Id
match Id
rep (Cast Expr Id
e CoercionR
_) = Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
e
replaceVar Id
match Id
rep (Case Expr Id
e Id
b Type
t [Alt Id]
alts) =
Expr Id -> Id -> Type -> [Alt Id] -> Expr Id
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case Expr Id
newExpr Id
newB Type
t ((Alt Id -> Alt Id) -> [Alt Id] -> [Alt Id]
forall a b. (a -> b) -> [a] -> [b]
map (\(Alt AltCon
con [Id]
binds Expr Id
expr) -> AltCon -> [Id] -> Expr Id -> Alt Id
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
con ((Id -> Id) -> [Id] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (\Id
v -> if Id
v Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
match then Id
rep else Id
v) [Id]
binds) (Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
expr)) [Alt Id]
alts)
where newExpr :: Expr Id
newExpr = Id -> Id -> Expr Id -> Expr Id
replaceVar Id
match Id
rep Expr Id
e
newB :: Id
newB = if Id
b Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
match then Id
rep else Id
b
replaceVar Id
_ Id
_ Expr Id
e = Expr Id
e
transformPrim :: Ctx -> Expr Var -> CoreM (Expr Var, PrimInfo)
transformPrim :: Ctx -> Expr Id -> CoreM (Expr Id, PrimInfo)
transformPrim Ctx
ctx expr :: Expr Id
expr@(App Expr Id
e Expr Id
e') = case Expr Id -> Maybe PrimInfo
isPrimExpr Expr Id
expr of
Just primInfo :: PrimInfo
primInfo@(AdvApp Id
f TypedArg
_) -> do
Id
varAdv' <- CoreM Id
adv'Var
let newE :: Expr Id
newE = Id -> Id -> Expr Id -> Expr Id
replaceVar Id
f Id
varAdv' Expr Id
e
(Expr Id, PrimInfo) -> CoreM (Expr Id, PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App Expr Id
newE Expr Id
e') (Id -> Expr Id
forall b. Id -> Expr b
Var (Maybe Id -> Id
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Id -> Id) -> Maybe Id -> Id
forall a b. (a -> b) -> a -> b
$ Ctx -> Maybe Id
fresh Ctx
ctx)), PrimInfo
primInfo)
Just primInfo :: PrimInfo
primInfo@(SelectApp Id
f TypedArg
_ TypedArg
_) -> do
Id
varSelect' <- CoreM Id
select'Var
let newE :: Expr Id
newE = Id -> Id -> Expr Id -> Expr Id
replaceVar Id
f Id
varSelect' Expr Id
e
(Expr Id, PrimInfo) -> CoreM (Expr Id, PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App Expr Id
newE Expr Id
e') (Id -> Expr Id
forall b. Id -> Expr b
Var (Maybe Id -> Id
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Id -> Id) -> Maybe Id -> Id
forall a b. (a -> b) -> a -> b
$ Ctx -> Maybe Id
fresh Ctx
ctx)), PrimInfo
primInfo)
Just (DelayApp Id
_ Type
t) -> do
Id
bigDelayVar <- CoreM Id
bigDelay
TyCon
inputValueV <- CoreM TyCon
inputValueVar
let inputValueType :: Type
inputValueType = TyCon -> Type
mkTyConTy TyCon
inputValueV
Id
inpVar <- FastString -> Type -> Type -> CoreM Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> Type -> Type -> m Id
mkSysLocalM (String -> FastString
fsLit String
"inpV") Type
inputValueType Type
inputValueType
let ctx' :: Ctx
ctx' = Ctx
ctx {fresh = Just inpVar}
(Expr Id
newExpr, Maybe PrimInfo
maybePrimInfo) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx' Expr Id
e'
let primInfo :: PrimInfo
primInfo = Maybe PrimInfo -> PrimInfo
forall a. HasCallStack => Maybe a -> a
fromJust Maybe PrimInfo
maybePrimInfo
let lambdaExpr :: Expr Id
lambdaExpr = Id -> Expr Id -> Expr Id
forall b. b -> Expr b -> Expr b
Lam Id
inpVar Expr Id
newExpr
Expr Id
clockCode <- PrimInfo -> CoreM (Expr Id)
constructClockExtractionCode PrimInfo
primInfo
(Expr Id, PrimInfo) -> CoreM (Expr Id, PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Id -> Expr Id
forall b. Id -> Expr b
Var Id
bigDelayVar) (Type -> Expr Id
forall b. Type -> Expr b
Type Type
t)) Expr Id
clockCode) Expr Id
lambdaExpr, PrimInfo
primInfo)
Just PrimInfo
primInfo -> do
String -> CoreM (Expr Id, PrimInfo)
forall a. HasCallStack => String -> a
error (String -> CoreM (Expr Id, PrimInfo))
-> String -> CoreM (Expr Id, PrimInfo)
forall a b. (a -> b) -> a -> b
$ SDoc -> String
showSDocUnsafe (SDoc -> String) -> SDoc -> String
forall a b. (a -> b) -> a -> b
$ String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"transformPrim: Cannot transform " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Prim -> SDoc
forall a. Outputable a => a -> SDoc
ppr (PrimInfo -> Prim
prim PrimInfo
primInfo)
Maybe PrimInfo
Nothing -> String -> CoreM (Expr Id, PrimInfo)
forall a. HasCallStack => String -> a
error String
"Cannot transform non-prim applications"
transformPrim Ctx
_ Expr Id
_ = do
String -> CoreM (Expr Id, PrimInfo)
forall a. HasCallStack => String -> a
error String
"Cannot transform anything else than prim applications"
transform :: CoreExpr -> CoreM CoreExpr
transform :: Expr Id -> CoreM (Expr Id)
transform Expr Id
expr = (Expr Id, Maybe PrimInfo) -> Expr Id
forall a b. (a, b) -> a
fst ((Expr Id, Maybe PrimInfo) -> Expr Id)
-> CoreM (Expr Id, Maybe PrimInfo) -> CoreM (Expr Id)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
emptyCtx Expr Id
expr
transform' :: Ctx -> CoreExpr -> CoreM (CoreExpr, Maybe PrimInfo)
transform' :: Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx expr :: Expr Id
expr@(App Expr Id
e Expr Id
e') = case Expr Id -> Maybe PrimInfo
isPrimExpr Expr Id
expr of
Just (BoxApp Id
_) -> do
(Expr Id
newExpr, Maybe PrimInfo
primInfo) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e'
(Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App Expr Id
e Expr Id
newExpr, Maybe PrimInfo
primInfo)
Just PrimInfo
_ -> do
(Expr Id
newExpr, PrimInfo
primInfo) <- Ctx -> Expr Id -> CoreM (Expr Id, PrimInfo)
transformPrim Ctx
ctx Expr Id
expr
(Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id
newExpr, PrimInfo -> Maybe PrimInfo
forall a. a -> Maybe a
Just PrimInfo
primInfo)
Maybe PrimInfo
Nothing -> do
(Expr Id
newExpr, Maybe PrimInfo
primInfo) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e
(Expr Id
newExpr', Maybe PrimInfo
primInfo') <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e'
(Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App Expr Id
newExpr Expr Id
newExpr', Maybe PrimInfo
primInfo Maybe PrimInfo -> Maybe PrimInfo -> Maybe PrimInfo
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe PrimInfo
primInfo')
transform' Ctx
ctx (Lam Id
b Expr Id
rhs) = do
(Expr Id
newExpr, Maybe PrimInfo
primInfo) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
rhs
(Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Id -> Expr Id -> Expr Id
forall b. b -> Expr b -> Expr b
Lam Id
b Expr Id
newExpr, Maybe PrimInfo
primInfo)
transform' Ctx
ctx (Let (NonRec Id
b Expr Id
rhs) Expr Id
e) = do
(Expr Id
newRhs, Maybe PrimInfo
primInfo) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
rhs
(Expr Id
newExpr, Maybe PrimInfo
primInfo') <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e
(Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bind Id -> Expr Id -> Expr Id
forall b. Bind b -> Expr b -> Expr b
Let (Id -> Expr Id -> Bind Id
forall b. b -> Expr b -> Bind b
NonRec Id
b Expr Id
newRhs) Expr Id
newExpr, Maybe PrimInfo
primInfo Maybe PrimInfo -> Maybe PrimInfo -> Maybe PrimInfo
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe PrimInfo
primInfo')
transform' Ctx
ctx (Let (Rec [(Id, Expr Id)]
binds) Expr Id
e) = do
[(Id, (Expr Id, Maybe PrimInfo))]
transformedBinds <- ((Id, Expr Id) -> CoreM (Id, (Expr Id, Maybe PrimInfo)))
-> [(Id, Expr Id)] -> CoreM [(Id, (Expr Id, Maybe PrimInfo))]
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
b, Expr Id
bindE) -> ((Expr Id, Maybe PrimInfo) -> (Id, (Expr Id, Maybe PrimInfo)))
-> CoreM (Expr Id, Maybe PrimInfo)
-> CoreM (Id, (Expr Id, Maybe PrimInfo))
forall a b. (a -> b) -> CoreM a -> CoreM b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Id
b,) (Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
bindE)) [(Id, Expr Id)]
binds
(Expr Id
e', Maybe PrimInfo
mPi) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e
let primInfos :: [Maybe PrimInfo]
primInfos = ((Id, (Expr Id, Maybe PrimInfo)) -> Maybe PrimInfo)
-> [(Id, (Expr Id, Maybe PrimInfo))] -> [Maybe PrimInfo]
forall a b. (a -> b) -> [a] -> [b]
map (\(Id
_, (Expr Id
_, Maybe PrimInfo
p)) -> Maybe PrimInfo
p) [(Id, (Expr Id, Maybe PrimInfo))]
transformedBinds
let firstPrimInfo :: Maybe PrimInfo
firstPrimInfo = (Maybe PrimInfo -> Maybe PrimInfo -> Maybe PrimInfo)
-> Maybe PrimInfo -> [Maybe PrimInfo] -> Maybe PrimInfo
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Maybe PrimInfo -> Maybe PrimInfo -> Maybe PrimInfo
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(<|>) Maybe PrimInfo
mPi [Maybe PrimInfo]
primInfos
[(Id, Expr Id)]
newBinds <- ((Id, (Expr Id, Maybe PrimInfo)) -> CoreM (Id, Expr Id))
-> [(Id, (Expr Id, Maybe PrimInfo))] -> CoreM [(Id, Expr 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 (\(Id
b, (Expr Id
e, Maybe PrimInfo
_)) -> (Id, Expr Id) -> CoreM (Id, Expr Id)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Id
b, Expr Id
e)) [(Id, (Expr Id, Maybe PrimInfo))]
transformedBinds
(Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bind Id -> Expr Id -> Expr Id
forall b. Bind b -> Expr b -> Expr b
Let ([(Id, Expr Id)] -> Bind Id
forall b. [(b, Expr b)] -> Bind b
Rec [(Id, Expr Id)]
newBinds) Expr Id
e', Maybe PrimInfo
firstPrimInfo)
transform' Ctx
ctx (Case Expr Id
e Id
b Type
t [Alt Id]
alts) = do
(Expr Id
expr, Maybe PrimInfo
primInfo) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e
[(Maybe PrimInfo, Alt Id)]
transformed <- (Alt Id -> CoreM (Maybe PrimInfo, Alt Id))
-> [Alt Id] -> CoreM [(Maybe PrimInfo, 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 (\(Alt AltCon
con [Id]
binds Expr Id
expr) -> Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
expr CoreM (Expr Id, Maybe PrimInfo)
-> ((Expr Id, Maybe PrimInfo) -> (Maybe PrimInfo, Alt Id))
-> CoreM (Maybe PrimInfo, Alt Id)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> (Expr Id -> Alt Id)
-> (Maybe PrimInfo, Expr Id) -> (Maybe PrimInfo, Alt Id)
forall a b. (a -> b) -> (Maybe PrimInfo, a) -> (Maybe PrimInfo, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (AltCon -> [Id] -> Expr Id -> Alt Id
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
con [Id]
binds) ((Maybe PrimInfo, Expr Id) -> (Maybe PrimInfo, Alt Id))
-> ((Expr Id, Maybe PrimInfo) -> (Maybe PrimInfo, Expr Id))
-> (Expr Id, Maybe PrimInfo)
-> (Maybe PrimInfo, Alt Id)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr Id, Maybe PrimInfo) -> (Maybe PrimInfo, Expr Id)
forall a b. (a, b) -> (b, a)
swap) [Alt Id]
alts
let firstPrimInfo :: Maybe PrimInfo
firstPrimInfo = (Maybe PrimInfo -> (Maybe PrimInfo, Alt Id) -> Maybe PrimInfo)
-> Maybe PrimInfo -> [(Maybe PrimInfo, Alt Id)] -> Maybe PrimInfo
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Maybe PrimInfo
acc (Maybe PrimInfo
p, Alt Id
_) -> Maybe PrimInfo
acc Maybe PrimInfo -> Maybe PrimInfo -> Maybe PrimInfo
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe PrimInfo
p) Maybe PrimInfo
primInfo [(Maybe PrimInfo, Alt Id)]
transformed
let alts' :: [Alt Id]
alts' = ((Maybe PrimInfo, Alt Id) -> Alt Id)
-> [(Maybe PrimInfo, Alt Id)] -> [Alt Id]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe PrimInfo, Alt Id) -> Alt Id
forall a b. (a, b) -> b
snd [(Maybe PrimInfo, Alt Id)]
transformed
(Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> Id -> Type -> [Alt Id] -> Expr Id
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case Expr Id
expr Id
b Type
t [Alt Id]
alts', Maybe PrimInfo
firstPrimInfo)
transform' Ctx
ctx (Cast Expr Id
e CoercionR
c) = do (Expr Id
e' , Maybe PrimInfo
p) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e; (Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> CoercionR -> Expr Id
forall b. Expr b -> CoercionR -> Expr b
Cast Expr Id
e' CoercionR
c, Maybe PrimInfo
p)
transform' Ctx
ctx (Tick CoreTickish
t Expr Id
e) = do (Expr Id
e' , Maybe PrimInfo
p) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e; (Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreTickish -> Expr Id -> Expr Id
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t Expr Id
e', Maybe PrimInfo
p)
transform' Ctx
_ Expr Id
e = (Expr Id, Maybe PrimInfo) -> CoreM (Expr Id, Maybe PrimInfo)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id
e, Maybe PrimInfo
forall a. Maybe a
Nothing)
constructClockExtractionCode :: PrimInfo -> CoreM CoreExpr
(AdvApp Id
_ TypedArg
arg) = TypedArg -> CoreM (Expr Id)
createClockCode TypedArg
arg
constructClockExtractionCode (SelectApp Id
_ TypedArg
arg TypedArg
arg2) =
TypedArg -> TypedArg -> CoreM (Expr Id)
clockUnion TypedArg
arg TypedArg
arg2
constructClockExtractionCode PrimInfo
primInfo = String -> CoreM (Expr Id)
forall a. HasCallStack => String -> a
error (String -> CoreM (Expr Id)) -> String -> CoreM (Expr Id)
forall a b. (a -> b) -> a -> b
$ String
"Cannot construct clock for prim " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SDoc -> String
showSDocUnsafe (Prim -> SDoc
forall a. Outputable a => a -> SDoc
ppr (PrimInfo -> Prim
prim PrimInfo
primInfo))
createClockCode :: (Var, Type) -> CoreM CoreExpr
createClockCode :: TypedArg -> CoreM (Expr Id)
createClockCode (Id
argV, Type
argT) = do
Id
extractClock <- CoreM Id
extractClockVar
Expr Id -> CoreM (Expr Id)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> CoreM (Expr Id)) -> Expr Id -> CoreM (Expr Id)
forall a b. (a -> b) -> a -> b
$ Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App (Id -> Expr Id
forall b. Id -> Expr b
Var Id
extractClock) (Type -> Expr Id
forall b. Type -> Expr b
Type Type
argT)) (Id -> Expr Id
forall b. Id -> Expr b
Var Id
argV)
clockUnion :: (Var,Type) -> (Var, Type) -> CoreM CoreExpr
clockUnion :: TypedArg -> TypedArg -> CoreM (Expr Id)
clockUnion TypedArg
arg TypedArg
arg2 = do
Expr Id
clock1Code <- TypedArg -> CoreM (Expr Id)
createClockCode TypedArg
arg
Expr Id
clock2Code <- TypedArg -> CoreM (Expr Id)
createClockCode TypedArg
arg2
Id
unionVar' <- CoreM Id
unionVar
Expr Id -> CoreM (Expr Id)
forall a. a -> CoreM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr Id -> CoreM (Expr Id)) -> Expr Id -> CoreM (Expr Id)
forall a b. (a -> b) -> a -> b
$
Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App
(
Expr Id -> Expr Id -> Expr Id
forall b. Expr b -> Expr b -> Expr b
App
(
(Id -> Expr Id
forall b. Id -> Expr b
Var Id
unionVar')
)
Expr Id
clock1Code
)
Expr Id
clock2Code