module DDC.Source.Tetra.Transform.Expand
( Config (..)
, configDefault
, Expand (..))
where
import DDC.Source.Tetra.Compounds
import DDC.Source.Tetra.Predicates
import DDC.Source.Tetra.DataDef
import DDC.Source.Tetra.Module
import DDC.Source.Tetra.Prim
import DDC.Source.Tetra.Exp
import DDC.Type.Collect
import DDC.Type.Env (KindEnv, TypeEnv)
import qualified DDC.Type.Env as Env
import qualified Data.Set as Set
data Config a n
= Config
{
configMakeTypeHole :: Kind n -> Type n }
configDefault :: Config a Name
configDefault
= Config
{ configMakeTypeHole = \k -> TVar (UPrim NameHole k)}
class Expand (c :: * -> * -> *) where
expand
:: Ord n
=> Config a n
-> KindEnv n -> TypeEnv n
-> c a n -> c a n
instance Expand Module where
expand config kenv tenv mm
= let
preTop p
= case p of
TopBind a b x
-> let (b', x') = expandQuant a config kenv (b, x)
in (TopBind a b' x', Env.singleton b')
TopData _ def
-> (p, typeEnvOfDataDef def)
(tops_quant, tenvs)
= unzip $ map preTop $ moduleTops mm
tenv' = Env.unions $ tenv : tenvs
tops' = map (expand config kenv tenv')
$ tops_quant
in mm { moduleTops = tops' }
instance Expand Top where
expand config kenv tenv top
= case top of
TopBind a b x
-> let tenv' = Env.extend b tenv
x' = expand config kenv tenv' x
in TopBind a b x'
TopData{}
-> top
instance Expand Exp where
expand config kenv tenv xx
= let down = expand config kenv tenv
in case xx of
XVar{}
-> expandApp config kenv tenv xx []
XCon{}
-> expandApp config kenv tenv xx []
XApp{}
| (x1, xas) <- takeXAppsWithAnnots xx
-> if isXVar x1 || isXCon x1
then let xas' = [ (expand config kenv tenv x, a) | (x, a) <- xas ]
in expandApp config kenv tenv x1 xas'
else let x1' = expand config kenv tenv x1
xas' = [ (expand config kenv tenv x, a) | (x, a) <- xas ]
in makeXAppsWithAnnots x1' xas'
XLet a (LLet b x1) x2
-> let
(b_quant, x1_quant)
= expandQuant a config kenv (b, x1)
tenv' = Env.extend b_quant tenv
x1' = expand config kenv tenv' x1_quant
x2' = expand config kenv tenv' x2
in XLet a (LLet b x1') x2'
XLet a (LRec bxs) x2
-> let
(bs_quant, xs_quant)
= unzip
$ [expandQuant a config kenv (b, x) | (b, x) <- bxs]
tenv' = Env.extends bs_quant tenv
xs' = map (expand config kenv tenv') xs_quant
x2' = expand config kenv tenv' x2
in XLet a (LRec (zip bs_quant xs')) x2'
XLAM a b x
-> let kenv' = Env.extend b kenv
x' = expand config kenv' tenv x
in XLAM a b x'
XLam a b x
-> let tenv' = Env.extend b tenv
x' = expand config kenv tenv' x
in XLam a b x'
XLet a (LPrivate bts mR bxs) x2
-> let tenv' = Env.extends bts kenv
kenv' = Env.extends bxs tenv
x2' = expand config kenv' tenv' x2
in XLet a (LPrivate bts mR bxs) x2'
XCase a x alts -> XCase a (down x) (map down alts)
XCast a c x -> XCast a c (down x)
XType{} -> xx
XWitness{} -> xx
XDefix a xs -> XDefix a (map down xs)
XInfixOp{} -> xx
XInfixVar{} -> xx
instance Expand Alt where
expand config kenv tenv alt
= case alt of
AAlt PDefault x2
-> let x2' = expand config kenv tenv x2
in AAlt PDefault x2'
AAlt (PData dc bs) x2
-> let tenv' = Env.extends bs tenv
x2' = expand config kenv tenv' x2
in AAlt (PData dc bs) x2'
expandQuant
:: Ord n
=> a
-> Config a n
-> KindEnv n
-> (Bind n, Exp a n)
-> (Bind n, Exp a n)
expandQuant a config kenv (b, x)
| fvs <- freeVarsT kenv (typeOfBind b)
, not $ Set.null fvs
= let
kHole = configMakeTypeHole config sComp
makeBind u
= case u of
UName n -> Just $ BName n kHole
UIx{} -> Just $ BAnon kHole
_ -> Nothing
Just bsNew = sequence $ map makeBind $ Set.toList fvs
t' = foldr TForall (typeOfBind b) bsNew
b' = replaceTypeOfBind t' b
x' = foldr (XLAM a) x bsNew
in (b', x')
| otherwise
= (b, x)
expandApp
:: Ord n
=> Config a n
-> KindEnv n
-> TypeEnv n
-> Exp a n
-> [(Exp a n, a)]
-> Exp a n
expandApp config _kenv tenv x0 xas0
| Just (a, u) <- slurpVarConBound x0
, Just tt <- Env.lookup u tenv
, not $ isBot tt
= let
go t xas
= case (t, xas) of
(TForall _b t2, (x1@(XType _ _t1'), a1) : xas')
-> (x1, a1) : go t2 xas'
(TForall b t2, xas')
-> let k = typeOfBind b
Just a0 = takeAnnotOfExp x0
xh = XType a0 (configMakeTypeHole config k)
in (xh, a) : go t2 xas'
_ -> xas
xas_expanded
= go tt xas0
in makeXAppsWithAnnots x0 xas_expanded
| otherwise
= makeXAppsWithAnnots x0 xas0
slurpVarConBound :: Exp a n -> Maybe (a, Bound n)
slurpVarConBound xx
= case xx of
XVar a u -> Just (a, u)
XCon a dc
| DaConBound n <- dc -> Just (a, UName n)
| DaConPrim n t <- dc -> Just (a, UPrim n t)
_ -> Nothing