{-# 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
    -- The checking pass has ensured that there are not advances on different
    -- clocks. Thus we can just pick the first PrimInfo we find.
    (Expr Id
expr, Maybe PrimInfo
primInfo) <- Ctx -> Expr Id -> CoreM (Expr Id, Maybe PrimInfo)
transform' Ctx
ctx Expr Id
e

    -- For each alternative, transform it and save the maybePrimInfo-value
    [(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

    -- Of all the primInfos we have, pick the first one. This is safe because
    -- the checking pass has ensured that the clocks of all primitives.
    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
constructClockExtractionCode :: PrimInfo -> CoreM (Expr Id)
constructClockExtractionCode (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)

-- Generate code for union of two clocks.
-- clockUnion (aVar, aType) (bVar, bType) returns the AST for:
--  Set.union (extractClock aVar) (extractClock bVar)

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