module DDC.Core.Transform.Unshare
(unshareModule)
where
import DDC.Core.Exp.Annot.AnTEC
import DDC.Core.Exp.Annot
import DDC.Core.Module
import DDC.Type.Transform.SubstituteT
import Data.Map (Map)
import qualified Data.Map.Strict as Map
unshareModule
:: (Ord n, Show n)
=> Module (AnTEC a n) n -> Module (AnTEC a n) n
unshareModule !mm
= let
importValuesNts
= [ let (iv', m) = addParamsImportValue iv
in ((n, iv'), m)
| (n, iv) <- moduleImportValues mm]
(importValues', ntssImport')
= unzip importValuesNts
(ntsBody, xx) = addParamsX $ moduleBody mm
nts' = Map.union (Map.unions ntssImport') ntsBody
xx' = addArgsX nts' xx
exportValues'
= [ (n, updateExportSource nts' ex)
| (n, ex) <- moduleExportValues mm ]
in mm { moduleBody = xx'
, moduleExportValues = exportValues'
, moduleImportValues = importValues' }
addParamsImportValue
:: ImportValue n -> (ImportValue n, Map n (Type n))
addParamsImportValue iv
= case iv of
ImportValueModule m n t (Just (nType, nValue, nBoxes))
-> case addParamsT t of
Just t'
-> ( ImportValueModule m n t'
(Just (nType, nValue + 1, nBoxes))
, Map.singleton n t')
Nothing
-> ( iv, Map.empty)
ImportValueModule{} -> (iv, Map.empty)
ImportValueSea{} -> (iv, Map.empty)
addParamsT :: Type n -> Maybe (Type n)
addParamsT tt
= case tt of
TVar{} -> Just $ tUnit `tFun` tt
TCon{} -> Just $ tUnit `tFun` tt
TForall b tBody
-> do tBody' <- addParamsT tBody
return $ TForall b tBody'
TApp{}
-> case takeTFun tt of
Nothing -> Just $ tUnit `tFun` tt
Just _ -> Nothing
TSum{}
-> Nothing
addParamsX
:: Ord n
=> Exp (AnTEC a n) n
-> ( Map n (Type n)
, Exp (AnTEC a n) n)
addParamsX xx
= case xx of
XLet a (LRec bxs) xBody
-> let (ns, bxs') = addParamsBXS a bxs
in ( ns
, XLet a (LRec bxs') xBody)
_ -> ( Map.empty
, xx)
addParamsBXS _a []
= (Map.empty, [])
addParamsBXS a ((b, x) : bxs)
= let (ns1, b', x') = addParamsBX a b x
(ns2, bxs') = addParamsBXS a bxs
in ( Map.union ns1 ns2
, (b', x') : bxs')
addParamsBX _ b@(BName n _) x
= case addParamsBodyX x of
Nothing
-> (Map.empty, b, x)
Just (x', t')
-> ( Map.singleton n t'
, replaceTypeOfBind t' b
, x')
addParamsBX _ b x
= (Map.empty, b, x)
addParamsBodyX xx
= case xx of
XLam{}
-> Nothing
XLAM a bParam xBody
-> case addParamsBodyX xBody of
Nothing
-> Nothing
Just (xBody', tBody')
-> let t' = TForall bParam tBody'
a' = a { annotType = t' }
in Just ( XLAM a' bParam xBody', t')
_
-> let a = annotOfExp xx
t' = tFun tUnit (annotType a)
a' = a { annotType = t' }
in Just (XLam a' (BNone tUnit) xx, t')
addArgsX :: (Show n, Ord n)
=> Map n (Type n)
-> Exp (AnTEC a n) n
-> Exp (AnTEC a n) n
addArgsX nts xx
= let downX = addArgsX nts
downLts = addArgsLts nts
downA = addArgsAlt nts
in case xx of
XVar _a (UName n)
-> case Map.lookup n nts of
Just tF -> fst $ wrapAppX xx tF
Nothing -> xx
XVar{} -> xx
XCon{} -> xx
XApp{} -> addArgsAppX nts xx []
XLAM a b xBody -> XLAM a b (downX xBody)
XLam a b xBody -> XLam a b (downX xBody)
XLet a lts xBody -> XLet a (downLts lts) (downX xBody)
XCase a xScrut as -> XCase a (downX xScrut) (map downA as)
XCast a c x -> XCast a c (downX x)
XType{} -> xx
XWitness{} -> xx
addArgsAppX !nts !xx !ats
= let downX = addArgsX nts
tA = annotType $ annotOfExp xx
in case xx of
XVar _a (UName n)
-> case Map.lookup n nts of
Just tF
-> let (x1, t1) = wrapAtsX xx tF ats
(x2, _) = wrapAppX x1 t1
in x2
Nothing
-> fst $ wrapAtsX xx tA ats
XVar{}
-> fst $ wrapAtsX xx tA ats
XCon{}
-> fst $ wrapAtsX xx tA ats
XApp _a1 x1 (XType a2 t)
-> addArgsAppX nts x1 ((a2, t) : ats)
XApp a x1 x2
-> XApp a (addArgsAppX nts x1 ats) (downX x2)
_ -> fst $ wrapAtsX xx tA ats
addArgsLts nts lts
= let downX = addArgsX nts
in case lts of
LLet b x -> LLet b (downX x)
LRec bxs -> LRec [(b, downX x) | (b, x) <- bxs]
LPrivate{} -> lts
addArgsAlt nts aa
= let downX = addArgsX nts
in case aa of
AAlt p x -> AAlt p (downX x)
wrapAppX :: Exp (AnTEC a n) n
-> Type n
-> (Exp (AnTEC a n) n, Type n)
wrapAppX xF tF
= case takeTFun tF of
Just (_, tResult)
-> let a = annotOfExp xF
aR = a { annotType = tResult }
aV = a { annotType = tF }
aU = a { annotType = tUnit }
xF' = mapAnnotOfExp (const aV) xF
in ( XApp aR xF' (xUnit aU)
, tResult)
Nothing
-> (xF, tF)
wrapAtsX !xF !tF []
= (xF, tF)
wrapAtsX !xF !tF ((aArg, tArg): ats)
= case tF of
TForall bParam tBody
-> let a = annotOfExp xF
tR = substituteT bParam tArg tBody
aR = a { annotType = tR }
aV = a { annotType = tF }
xF' = mapAnnotOfExp (const aV) xF
in wrapAtsX
(XApp aR xF' (XType aArg tArg))
tR ats
_ -> (xF, tF)
updateExportSource
:: Ord n
=> Map n (Type n) -> ExportSource n -> ExportSource n
updateExportSource mm ex
= case ex of
ExportSourceLocal n _t
-> case Map.lookup n mm of
Nothing -> ex
Just t' -> ExportSourceLocal n t'
ExportSourceLocalNoType _
-> ex