ddc-source-tetra-0.4.3.1: DDC/Source/Tetra/Transform/Expand.hs
{-# LANGUAGE TypeFamilies, UndecidableInstances #-}
-- | Look at type signatures and add quantifiers to bind any free type
-- variables.
--
-- Given
--
-- @
-- mapS (f : a -> S e b) (xx : List a) : S e (List b)
-- = box case xx of
-- Nil -> Nil
-- Cons x xs -> Cons (run f x) (run mapS f xs)
-- @
--
-- We get:
--
-- @
-- mapS [a e b : ?] (f : a -> S e b) (xx : List a) : S e (List b)
-- = /\(a e b : ?). box case xx of
-- Nil -> Nil
-- Cons x xs -> Cons (run f x) (run mapS f xs)
-- @
--
module DDC.Source.Tetra.Transform.Expand
( expandModule
, Expand (..))
where
import DDC.Source.Tetra.Collect.FreeVars
import DDC.Source.Tetra.Exp
import DDC.Source.Tetra.DataDef
import DDC.Source.Tetra.Module
import DDC.Data.SourcePos
import Data.Function
import DDC.Source.Tetra.Env (Env)
import Data.Maybe
import qualified DDC.Source.Tetra.Env as Env
import qualified Data.Set as Set
import qualified Data.List as List
-- | Run the expander on the given module.
expandModule :: SourcePos -> Module Source -> Module Source
expandModule sp mm
= expand sp Env.empty mm
---------------------------------------------------------------------------------------------------
class Expand c where
-- | Add quantifiers to the types of binders. Also add holes for missing
-- type arguments.
expand :: SourcePos -> Env -> c -> c
---------------------------------------------------------------------------------------------------
instance Expand (Module Source) where
expand = expandM
expandM a env mm
= let
-- Add quantifiers to the types of bindings, and also slurp
-- out the contribution to the top-level environment from each binding.
-- We need to do this in an initial binding because each top-level
-- thing is in-scope of all the others.
preTop p
= case p of
TopClause aT (SLet aL bm ps gxs)
-> let (bm', ps') = expandQuantParams env bm ps
in ( TopClause aT (SLet aL bm' ps' gxs)
, Env.extendDaVarMT bm' Env.empty)
-- Clauses should have already desugared.
TopClause _ SSig{} -> (p, Env.empty)
TopData _ def -> (p, envOfDataDef def)
TopType{} -> (p, Env.empty)
(tops_quant, envs)
= unzip $ map preTop $ moduleTops mm
-- Build the compound top-level environment.
env' = Env.unions $ env : envs
-- Expand all the top-level definitions.
tops' = map (expand a env')
$ tops_quant
in mm { moduleTops = tops' }
---------------------------------------------------------------------------------------------------
instance Expand (Top Source) where
expand = expandT
expandT _a env top
= case top of
TopClause a1 (SLet a2 bm ps gxs)
-> let env' = Env.extendDaVarMT bm env
env'' = List.foldl' (flip extendParam) env' ps
gxs' = map (expand a2 env'') gxs
in TopClause a1 (SLet a2 bm ps gxs')
TopClause _ (SSig{}) -> top
TopData{} -> top
TopType{} -> top
---------------------------------------------------------------------------------------------------
instance Expand Exp where
expand = downX
downX a env xx
= case xx of
XAnnot a' x
-> downX a' env x
-- Invoke the expander --------
XVar{} -> xx
XCon{} -> xx
XPrim{} -> xx
XApp{}
| (x1, xas) <- takeXAppsWithAnnots xx
-> let x1' = expand a env x1
xas' = [ (expand (fromMaybe a a') env x, a')
| (x, a') <- xas ]
in makeXAppsWithAnnots x1' xas'
XLet (LLet b x1) x2
-> let x1' = expand a env x1
env' = Env.extendDaVarMT b env
x2' = expand a env' x2
in XLet (LLet b x1') x2'
XLet (LRec bxs) x2
-> let (bs, xs) = unzip bxs
env' = Env.extendsDaVarMT bs env
xs' = map (expand a env') xs
bxs' = zip bs xs'
x2' = expand a env' x2
in XLet (LRec bxs') x2'
XLet (LGroup cs) x2
-> let cs' = map (downCX a env) cs
bs = [b | SLet _ b _ _ <- cs']
env' = Env.extendsDaVarMT bs env
x2' = downX a env' x2
in XLet (LGroup cs') x2'
-- Boilerplate ----------------
XLAM bm@(XBindVarMT b _) x
-> let env' = env & Env.extendTyVar' b
x' = expand a env' x
in XLAM bm x'
XLam bm@(XBindVarMT b _) x
-> let env' = env & Env.extendDaVar' b
x' = expand a env' x
in XLam bm x'
XLet (LPrivate bts mR bxs) x2
-> let env' = env & Env.extendsTyVar' bts
& Env.extendsDaVar bxs
x2' = expand a env' x2
in XLet (LPrivate bts mR bxs) x2'
XCase x alts -> XCase (downX a env x)
(map (downA a env) alts)
XCast c x -> XCast c (downX a env x)
XType{} -> xx
XWitness{} -> xx
XDefix a' xs -> XDefix a' (map (downX a' env) xs)
XInfixOp{} -> xx
XInfixVar{} -> xx
XMatch a' as x -> XMatch a' (map (downMA a' env) as) (downX a' env x)
XWhere a' x cls -> XWhere a' (downX a' env x) (map (downCX a env) cls)
XLamPat a' p mt x
-> let env' = extendPat p env
in XLamPat a' p mt (downX a' env' x)
XLamCase a' alts -> XLamCase a (map (downA a' env) alts)
---------------------------------------------------------------------------------------------------
instance Expand Clause where
expand a env cl
= downCX a env cl
downCX _a env cl
= case expandQuantClause env cl of
(_, SSig{})
-> cl
(env', SLet a mt ps gxs)
-> let gxs' = map (downGX a env') gxs
in SLet a mt ps gxs'
---------------------------------------------------------------------------------------------------
instance Expand GuardedExp where
expand = downGX
downGX a env (GGuard g x)
= let g' = expand a env g
env' = extendGuard g' env
in GGuard g' (expand a env' x)
downGX a env (GExp x)
= let x' = expand a env x
in GExp x'
---------------------------------------------------------------------------------------------------
instance Expand Guard where
expand = downG
downG a env gg
= case gg of
GPat p x
-> let env' = extendPat p env
x' = expand a env' x
in GPat p x'
GPred x
-> let x' = expand a env x
in GPred x'
GDefault
-> GDefault
---------------------------------------------------------------------------------------------------
instance Expand AltCase where
expand = downA
downA a env alt
= case alt of
AAltCase p gsx
-> let env' = extendPat p env
gsx' = map (expand a env') gsx
in AAltCase p gsx'
---------------------------------------------------------------------------------------------------
instance Expand AltMatch where
expand = downMA
downMA a env alt
= case alt of
AAltMatch gx -> AAltMatch (downGX a env gx)
---------------------------------------------------------------------------------------------------
-- | Extend a type environment with the variables bound by the given pattern.
extendPat :: Pat -> Env -> Env
extendPat ww env
= case ww of
PDefault -> env
PAt b p -> extendPat p $ Env.union env (Env.singletonDaVar' b)
PVar b -> Env.union env (Env.singletonDaVar' b)
PData{} -> env
extendParam :: Param -> Env -> Env
extendParam pp env
= case pp of
MType b _ -> Env.union env (Env.singletonTyVar' b)
MWitness b _ -> Env.union env (Env.singletonDaVar' b)
MValue p _ -> extendPat p env
-- | Extend a type environment with the variables bound by the given guard.
extendGuard :: Guard -> Env -> Env
extendGuard gg tenv
= case gg of
GPat w _ -> extendPat w tenv
_ -> tenv
---------------------------------------------------------------------------------------------------
expandQuantClause :: Env -> Clause -> (Env, Clause)
expandQuantClause env cc
= case cc of
SSig{}
-> (env, cc)
SLet a mt ps gxs
-> let (mt', ps') = expandQuantParams env mt ps
in (env, SLet a mt' ps' gxs)
-- | Expand missing quantifiers in types of bindings.
--
-- If a binding mentions type variables that are not in scope then add new
-- quantifiers to its type, as well as matching type lambdas.
--
expandQuantParams
:: Env -- ^ Current environment.
-> BindVarMT -- ^ Type of binding.
-> [Param] -- ^ Parameters of binding.
-> (BindVarMT, [Param]) -- ^ Expanded type and body of binding.
expandQuantParams env bmBind ps
| XBindVarMT bBind (Just tBind) <- bmBind
, fvs <- freeVarsT env tBind
, not $ Set.null fvs
= let
-- Make new binders for each of the free type variables.
-- We shouldn't have any holes or indices in the incoming type,
-- but don't have a way to specify this in the type of the AST.
makeBind u
= case u of
UName n -> Just $ BName n
UHole -> error "ddc-source-tetra.expandQuant: not expanding hole in type"
UIx{} -> error "ddc-source-tetra.expandQuant: not expanding deBruijn type var"
Just bsNew = sequence $ map makeBind $ Set.toList fvs
-- Attach quantifiers to the front of the old type,
-- using a hole bound to indicate we want the type inferencer,
-- to infer a kind or this.
k = TVar UHole
tBind' = foldr (\b t -> TApp (TCon (TyConForall k)) (TAbs b k t)) tBind bsNew
-- Attach type lambdas to the front of the term
-- using a matching hole bound on the type abstraction.
-- We could instead just not include a kind, but use
-- a hole so the form of the term matches the form
-- of its type.
ps' = [MType b Nothing | b <- bsNew] ++ ps
in (XBindVarMT bBind (Just tBind'), ps')
| otherwise
= (bmBind, ps)