module DDC.Core.Flow.Transform.Wind
( RefInfo(..)
, windModule)
where
import DDC.Core.Module
import DDC.Core.Exp
import DDC.Core.Flow
import DDC.Core.Flow.Prim
import DDC.Core.Compounds
import DDC.Core.Flow.Compounds
(tNat, dcNat, dcTupleN, dcBool, tTupleN)
import qualified Data.Map as Map
import Data.Map (Map)
data RefInfo
= RefInfo
{ refInfoName :: Name
, refInfoType :: Type Name
, refInfoCurrent :: Name
, refInfoVersionNumber :: Int }
data RefMap
= RefMap (Map Name RefInfo)
refMapZero :: RefMap
refMapZero = RefMap Map.empty
refMapElems :: RefMap -> [RefInfo]
refMapElems (RefMap mm)
= Map.elems mm
insertRefInfo :: RefInfo -> RefMap -> RefMap
insertRefInfo info (RefMap mm)
= RefMap (Map.insert (refInfoName info) info mm)
lookupRefInfo :: RefMap -> Name -> Maybe RefInfo
lookupRefInfo (RefMap mm) n
= Map.lookup n mm
nameOfRefInfo :: RefInfo -> Maybe Name
nameOfRefInfo info
= Just $ NameVarMod (refInfoName info) (show $ refInfoVersionNumber info)
bumpVersionOfRefInfo :: RefInfo -> RefInfo
bumpVersionOfRefInfo info
= info { refInfoVersionNumber = refInfoVersionNumber info + 1 }
bumpVersionInRefMap :: Name -> RefMap -> RefMap
bumpVersionInRefMap n (RefMap mm)
= RefMap $ Map.update (Just . bumpVersionOfRefInfo) n mm
bumpAllVersionsInRefMap :: RefMap -> RefMap
bumpAllVersionsInRefMap mm
= foldr bumpVersionInRefMap mm $ map refInfoName $ refMapElems mm
data Context
= ContextLoop
{ contextLoopName :: Name
, contextLoopCounter :: Name
, contextLoopAccs :: [Name] }
| ContextGuard
{
contextGuardCounter :: Name
, contextGuardFlag :: Bool }
deriving Show
isContextLoop :: Context -> Bool
isContextLoop cc
= case cc of
ContextLoop{} -> True
_ -> False
makeTailCallFromContexts :: a -> RefMap -> [Context] -> Exp a Name
makeTailCallFromContexts a refMap context@(ContextLoop nLoop _ _ : _)
= let
xLoop = XVar a (UName nLoop)
xArgs = slurpArgUpdates a refMap [] context
in xApps a xLoop xArgs
makeTailCallFromContexts _ _ contexts
= error $ unlines
[ "ddc-core-flow.makeTailCallFromContexts"
, " Can't make a tailcall for this context."
, " context = " ++ show contexts ]
slurpArgUpdates
:: a
-> RefMap
-> [(Name, Exp a Name)]
-> [Context]
-> [Exp a Name]
slurpArgUpdates a refMap [] (ContextLoop _ nCounter nAccs : more)
= let
nxCounter'
= ( nCounter
, xIncrement a (XVar a (UName nCounter)) )
nxAccs'
= [ (nAcc, XVar a (UName nAcc'))
| nAcc <- nAccs
, let Just info = lookupRefInfo refMap nAcc
, let Just nAcc' = nameOfRefInfo info ]
in slurpArgUpdates a refMap (nxCounter' : nxAccs') more
slurpArgUpdates a refMap args (ContextGuard nCounter flag : more)
| flag == True
= let
update [] = []
update ((n, x) : args')
| n == nCounter = (n, xIncrement a x) : update args'
| otherwise = (n, x) : update args'
in slurpArgUpdates a refMap (update args) more
| otherwise
= slurpArgUpdates a refMap args more
slurpArgUpdates _ _ _ (ContextLoop{} : _)
= error $ unlines
[ "ddc-core-flow.slurpArgUpdates"
, " Nested loops are not supported." ]
slurpArgUpdates _ _ args []
= map snd args
xIncrement :: a -> Exp a Name -> Exp a Name
xIncrement a xx
= xApps a (XVar a (UPrim (NamePrimArith PrimArithAdd)
(typePrimArith PrimArithAdd)))
[ XType a tNat, xx, XCon a (dcNat 1) ]
xSubInt :: a -> Exp a Name -> Exp a Name -> Exp a Name
xSubInt a x1 x2
= xApps a (XVar a (UPrim (NamePrimArith PrimArithSub)
(typePrimArith PrimArithSub)))
[ XType a tNat, x1, x2]
windModule :: Module () Name -> Module () Name
windModule m
= let body' = windModuleBodyX (moduleBody m)
in m { moduleBody = body' }
windModuleBodyX :: Exp () Name -> Exp () Name
windModuleBodyX xx
= case xx of
XLet a (LLet b x1) x2
-> let x1' = windBodyX refMapZero [] x1
x2' = windModuleBodyX x2
in XLet a (LLet b x1') x2'
XLet a (LRec bxs) x2
-> let bxs' = [(b, windBodyX refMapZero [] x) | (b, x) <- bxs]
x2' = windModuleBodyX x2
in XLet a (LRec bxs') x2'
XLet a lts x2
-> let x2' = windModuleBodyX x2
in XLet a lts x2'
_ -> xx
windBodyX
:: RefMap
-> [Context]
-> Exp () Name
-> Exp () Name
windBodyX refMap context xx
= let down = windBodyX refMap context
in case xx of
XLet a (LLet (BName nRef _) x) x2
| Just ( NameOpStore OpStoreNew
, [XType _ tElem, xVal] ) <- takeXPrimApps x
-> let
info = RefInfo
{ refInfoName = nRef
, refInfoType = tElem
, refInfoCurrent = nInit
, refInfoVersionNumber = 0 }
Just nInit = nameOfRefInfo info
refMap' = insertRefInfo info refMap
in XLet a (LLet (BName nInit tElem) xVal)
(windBodyX refMap' context x2)
XLet a (LLet bResult x) x2
| Just ( NameOpStore OpStoreRead
, [XType _ _tElem, XVar _ (UName nRef)] )
<- takeXPrimApps x
, Just info <- lookupRefInfo refMap nRef
, Just nVal <- nameOfRefInfo info
-> XLet a (LLet bResult (XVar a (UName nVal)))
(windBodyX refMap context x2)
XLet a (LLet (BNone _) x) x2
| Just ( NameOpStore OpStoreWrite
, [XType _ _tElem, XVar _ (UName nRef), xVal])
<- takeXPrimApps x
, refMap' <- bumpVersionInRefMap nRef refMap
, Just info <- lookupRefInfo refMap' nRef
, Just nVal <- nameOfRefInfo info
, tVal <- refInfoType info
-> XLet a (LLet (BName nVal tVal) xVal)
(windBodyX refMap' context x2)
XLet a (LLet (BNone _) x) x2
| Just ( NameOpControl OpControlLoopN
, [ XType _ tK, xLength
, XLam _ bIx@(BName nIx _) xBody]) <- takeXPrimApps x
-> let
nLoop = NameVar "loop"
bLoop = BName nLoop tLoop
uLoop = UName nLoop
nLength = NameVarMod nLoop "length"
bLength = BName nLength tNat
uLength = UName nLength
refMap_init = refMap
refMap_body = bumpAllVersionsInRefMap refMap
refMap_final = bumpAllVersionsInRefMap refMap_body
bsAccs = [ BName nVar (refInfoType info)
| info <- refMapElems refMap_body
, let Just nVar = nameOfRefInfo info ]
usAccs = takeSubstBoundsOfBinds bsAccs
tsAccs = map typeOfBind bsAccs
tIndex = typeOfBind bIx
tResult = loopResultT tsAccs
tLoop = foldr tFun tResult (tIndex : tsAccs)
context' = context
++ [ ContextLoop
{ contextLoopName = nLoop
, contextLoopCounter = nIx
, contextLoopAccs = map refInfoName
$ refMapElems refMap_body } ]
xBody' = windBodyX refMap_body context' xBody
xDriver = xLams a (bIx : bsAccs)
$ XCase a (xSubInt a (XVar a uLength) (XVar a (UName nIx)))
[ AAlt (PData (dcNat 0) []) xResult
, AAlt PDefault xBody' ]
xResult = loopResultX a
tsAccs
[XVar a u | u <- usAccs]
xsInit = XCon a (dcNat 0)
: [ XVar a (UName nVar)
| info <- refMapElems refMap_init
, let Just nVar = nameOfRefInfo info ]
bsFinal = [ BName nVar (refInfoType info)
| info <- refMapElems refMap_final
, let Just nVar = nameOfRefInfo info ]
x2' = windBodyX refMap_final context x2
in XLet a (LLet bLength (xNatOfRateNat tK xLength))
$ XLet a (LRec [(bLoop, xDriver)])
$ runUnpackLoop
a
tsAccs
(xApps a (XVar a uLoop) xsInit)
bsFinal
x2'
XLet a (LLet (BNone _) x) x2
| Just ( NameOpControl OpControlGuard
, [ XVar _ (UName nCountRef)
, xFlag
, XLam _ bCount xBody ]) <- takeXPrimApps x
-> let
Just infoCount = lookupRefInfo refMap nCountRef
Just nCount = nameOfRefInfo infoCount
context' = context
++ [ ContextGuard
{ contextGuardCounter = nCountRef
, contextGuardFlag = True } ]
xBody' = XLet a (LLet bCount (XVar a (UName nCount)))
$ windBodyX refMap context' xBody
in XCase a xFlag
[ AAlt (PData (dcBool True) []) xBody'
, AAlt PDefault (down x2) ]
XCon a dc
| any isContextLoop context
, dc == dcUnit
-> makeTailCallFromContexts a refMap context
XApp{}
| Just ( NameOpControl (OpControlSplit n)
, [ XType _ tK, xN, xBranch1, xBranch2 ]) <- takeXPrimApps xx
-> let xBranch1' = down xBranch1
xBranch2' = down xBranch2
in xSplit n tK xN xBranch1' xBranch2'
XVar{} -> xx
XCon{} -> xx
XLAM a b x -> XLAM a b (down x)
XLam a b x -> XLam a b (down x)
XApp{} -> xx
XLet a (LLet b x) x2
-> XLet a (LLet b (windBodyX refMap [] x))
(down x2)
XLet a (LRec bxs) x2
-> XLet a (LRec [(b, windBodyX refMap [] x) | (b, x) <- bxs])
(down x2)
XLet a lts x2
-> XLet a lts (down x2)
XCase{}
-> error $ unlines
[ "ddc-core-flow.windBodyX"
, " case-expressions not supported yet" ]
XCast a c x
-> let x' = windBodyX refMap context x
in XCast a c x'
XType{} -> xx
XWitness{} -> xx
type TypeF = Type Name
type ExpF = Exp () Name
xNatOfRateNat :: Type Name -> Exp () Name -> Exp () Name
xNatOfRateNat tK xR
= xApps ()
(xVarOpConcrete OpConcreteNatOfRateNat)
[XType () tK, xR]
xVarOpConcrete :: OpConcrete -> Exp () Name
xVarOpConcrete op
= XVar () (UPrim (NameOpConcrete op) (typeOpConcrete op))
xSplit :: Int
-> TypeF
-> ExpF
-> ExpF -> ExpF -> ExpF
xSplit n tK xRN xDownFn xTailFn
= xApps ()
(xVarOpControl $ OpControlSplit n)
[ XType () tK, xRN, xDownFn, xTailFn ]
xVarOpControl :: OpControl -> Exp () Name
xVarOpControl op
= XVar () (UPrim (NameOpControl op) (typeOpControl op))
loopResultT :: [Type Name] -> Type Name
loopResultT tsAccs
= case tsAccs of
[] -> tUnit
[tAcc] -> tAcc
_ -> tTupleN tsAccs
loopResultX :: a -> [Type Name] -> [Exp a Name] -> Exp a Name
loopResultX a tsAccs xsAccs
= case xsAccs of
[] -> xUnit a
[x] -> x
_ -> xApps a (XCon a (dcTupleN $ length tsAccs))
([XType a t | t <- tsAccs] ++ xsAccs)
runUnpackLoop
:: a
-> [Type Name]
-> Exp a Name
-> [Bind Name]
-> Exp a Name
-> Exp a Name
runUnpackLoop a tsAccs xRunLoop bsAcc xCont
| [] <- tsAccs
= XLet a (LLet (BNone tUnit) xRunLoop) xCont
| [_t] <- tsAccs
, [b] <- bsAcc
= XLet a (LLet b xRunLoop) xCont
| otherwise
= XCase a xRunLoop
[ AAlt (PData (dcTupleN $ length tsAccs) bsAcc) xCont ]