module DDC.Core.Transform.Thread
( Thread (..)
, Config (..)
, injectStateType)
where
import DDC.Core.Compounds
import DDC.Core.Module
import DDC.Core.Exp
import DDC.Base.Pretty
import DDC.Core.Transform.Reannotate
import DDC.Core.Check (AnTEC (..))
import DDC.Type.Env (KindEnv, TypeEnv)
import qualified DDC.Type.Env as Env
import qualified DDC.Core.Check as Check
data Config a n
= Config
{
configCheckConfig :: Check.Config n
, configThreadMe :: n -> Type n -> Maybe (Type n)
, configTokenType :: Type n
, configVoidType :: Type n
, configWrapResultType :: Type n -> Type n
, configWrapResultExp :: Exp (AnTEC a n) n -> Exp (AnTEC a n) n
-> Exp a n
, configThreadPat :: n -> Maybe (Bind n -> [Bind n] -> Pat n)
}
class Thread (c :: * -> * -> *) where
thread :: (Ord n, Show n, Pretty n)
=> Config a n
-> KindEnv n -> TypeEnv n
-> c (AnTEC a n) n
-> c a n
instance Thread Module where
thread config kenv tenv mm
= let body' = threadModuleBody config kenv tenv (moduleBody mm)
in mm { moduleBody = body' }
data Context n
= ContextRec n
| ContextFun n
deriving Eq
threadModuleBody
:: (Ord n, Show n, Pretty n)
=> Config a n
-> KindEnv n -> TypeEnv n
-> Exp (AnTEC a n) n
-> Exp a n
threadModuleBody config kenv tenv xx
= case xx of
XLet a lts x
-> let lts' = threadTopLets config kenv tenv lts
(bks, bts) = bindsOfLets lts
kenv' = Env.extends bks kenv
tenv' = Env.extends bts tenv
x' = threadModuleBody config kenv' tenv' x
in XLet (annotTail a) lts' x'
_ -> reannotate annotTail xx
threadTopLets
:: (Ord n, Show n, Pretty n)
=> Config a n
-> KindEnv n -> TypeEnv n
-> Lets (AnTEC a n) n
-> Lets a n
threadTopLets config kenv tenv lts
= case lts of
LLet b x
-> let (b', x') = threadTopBind config [] kenv tenv b x
in LLet b' x'
LRec bxs
-> let tenv' = Env.extends (map fst bxs) tenv
bxs' = [ threadTopBind config [ContextRec n] kenv tenv' b x
| (b, x) <- bxs
, let BName n _ = b ]
in LRec bxs'
_ -> reannotate annotTail lts
threadTopBind
:: (Ord n, Show n, Pretty n)
=> Config a n
-> [Context n]
-> KindEnv n -> TypeEnv n
-> Bind n -> Exp (AnTEC a n) n
-> (Bind n, Exp a n)
threadTopBind config context kenv tenv b xBody
= let tBind = typeOfBind b
tBind' = injectStateType config tBind
b' = replaceTypeOfBind tBind' b
tenv' = Env.extend b' tenv
tsArgs = fst $ takeTFunAllArgResult tBind'
in ( b'
, threadProc config context kenv tenv' xBody tsArgs)
threadArg
:: (Ord n, Show n, Pretty n)
=> Config a n
-> [Context n]
-> KindEnv n -> TypeEnv n
-> Type n -> Exp (AnTEC a n) n
-> Exp a n
threadArg config context kenv tenv t xx
= case xx of
XLam{} -> threadProcArg config context kenv tenv t xx
XLAM{} -> threadProcArg config context kenv tenv t xx
_ -> reannotate annotTail xx
threadProcArg config context kenv tenv t xx
= let tsArgs = fst $ takeTFunAllArgResult t
in threadProc config context kenv tenv xx tsArgs
threadProc
:: (Ord n, Show n, Pretty n)
=> Config a n
-> [Context n]
-> KindEnv n -> TypeEnv n
-> Exp (AnTEC a n) n
-> [Type n]
-> Exp a n
threadProc config context kenv tenv xx []
= threadProcBody config context kenv tenv xx
threadProc config context kenv tenv xx (t : tsArgs)
= case xx of
XLAM a b x
-> let kenv' = Env.extend b kenv
x' = threadProc config context kenv' tenv x tsArgs
in XLAM (annotTail a) b x'
XLam a b x
-> let tenv' = Env.extend b tenv
x' = threadProc config context kenv tenv' x tsArgs
in XLam (annotTail a) b x'
_ | Just a <- takeAnnotOfExp xx
, t == configTokenType config
-> let b' = BAnon (configTokenType config)
tenv' = Env.extend b' tenv
x' = threadProc config context kenv tenv' xx tsArgs
in XLam (annotTail a) b' x'
_ -> threadProcBody config context kenv tenv xx
threadProcBody
:: (Ord n, Show n, Pretty n)
=> Config a n
-> [Context n]
-> KindEnv n -> TypeEnv n
-> Exp (AnTEC a n) n
-> Exp a n
threadProcBody config context kenv tenv xx
= case xx of
XLet a (LRec bxs) x2
-> let bxs' = [threadTopBind config
(context ++ [ContextRec n])
kenv tenv b x
| (b, x) <- bxs
, let BName n _ = b ]
tenv' = Env.extends (map fst bxs) tenv
x2' = threadProcBody config
(context ++ [ContextFun n
| (b, _x) <- bxs
, let BName n _ = b ])
kenv tenv' x2
in XLet (annotTail a) (LRec bxs') x2'
XLet _ (LLet b x) x2
| Just (XVar a u, xsArgs) <- takeXApps x
, Just n <- takeNameOfBound u
, Just tOld <- Env.lookup u tenv
, Just tNew <- configThreadMe config n tOld
, Just mkPat <- configThreadPat config n
-> let
tWorld = configTokenType config
xsArgs' = xsArgs ++ [XVar a (UIx 0)]
tsArgs = fst $ takeTFunAllArgResult tNew
xsArgs'' = zipWith (threadArg config context kenv tenv) tsArgs xsArgs'
u' = replaceTypeOfBound tNew u
x' = xApps (annotTail a) (XVar (annotTail a) u') xsArgs''
tenv' = Env.extend b tenv
x2' = threadProcBody config context kenv tenv' x2
pat' = mkPat (BAnon tWorld) [b]
in XCase (annotTail a) x' [AAlt pat' x2']
XLet a (LLet b x1) x2
| Just (XVar _ (UName n), _xsArgs) <- takeXApps x1
, elem (ContextFun n) context
, Just mkPat <- configThreadPat config n
-> let
tWorld = configTokenType config
a' = annotTail a
x1' = XApp a' (reannotate annotTail x1) (XVar a' (UIx 0))
x2' = threadProcBody config context kenv tenv x2
pat' = mkPat (BAnon tWorld) [b]
in XCase (annotTail a) x1' [AAlt pat' x2']
XLet a lts x
-> let (bks, bts) = bindsOfLets lts
kenv' = Env.extends bks kenv
tenv' = Env.extends bts tenv
lts' = reannotate annotTail lts
x' = threadProcBody config context kenv' tenv' x
in XLet (annotTail a) lts' x'
XCase a xScrut [AAlt (PData _dc bs) xBody]
| Just ((XVar _ (UName n), _xsArgs)) <- takeXApps xScrut
, elem (ContextFun n) context
, Just mkPat <- configThreadPat config n
-> let
a' = annotTail a
tWorld = configTokenType config
xScrut' = XApp a' (reannotate annotTail xScrut) (XVar a' (UIx 0))
pat' = mkPat (BAnon tWorld) bs
alt' = threadAlt config context kenv tenv
(AAlt pat' xBody)
in XCase (annotTail a) xScrut' [alt']
XCase a x alts
-> let alts' = map (threadAlt config context kenv tenv) alts
x' = reannotate annotTail x
in XCase (annotTail a) x' alts'
XLAM{} -> error "ddc-core-simpl.Thread: unexpected XLAM"
XLam{} -> error "ddc-core-simpl.Thread: unexpected XLam"
XCast{} -> error "ddc-core-simpl.Thread: unexpected cast."
XType t -> XType t
XWitness w
-> XWitness (reannotate annotTail w)
XApp a _ _
| Just ((XVar _ (UName n), _xsArgs)) <- takeXApps xx
, elem (ContextRec n) context
-> let a' = annotTail a
in XApp a' (reannotate annotTail xx)
(XVar a' (UIx 0))
_
| otherwise
-> let Just a = takeAnnotOfExp xx
a' = AnTEC (configTokenType config)
(tBot kEffect)
(tBot kClosure)
(annotTail a)
xWorld = XVar a' (UIx 0)
wrap = configWrapResultExp config
in wrap xWorld xx
threadAlt
:: (Ord n, Show n, Pretty n)
=> Config a n
-> [Context n]
-> KindEnv n -> TypeEnv n
-> Alt (AnTEC a n) n
-> Alt a n
threadAlt config context kenv tenv (AAlt pat xx)
= case pat of
PDefault
-> AAlt pat (threadProcBody config context kenv tenv xx)
PData _ bs
-> let tenv' = Env.extends bs tenv
in AAlt pat (threadProcBody config context kenv tenv' xx)
injectStateType :: Eq n => Config a n -> Type n -> Type n
injectStateType config tt
= let down = injectStateType config
in case tt of
TForall b x
-> TForall b (down x)
TApp{}
| (tsArg@(_ : _), tResult) <- takeTFunArgResult tt
-> let tsArg' = tsArg ++ [configTokenType config]
tResult' = injectStateType config tResult
in foldr tFunPE tResult' tsArg'
_ | tt == configTokenType config -> tt
| tt == configVoidType config -> configTokenType config
| otherwise -> configWrapResultType config tt