ddc-source-tetra-0.4.3.1: DDC/Source/Tetra/Transform/Prep.hs
{-# LANGUAGE TypeFamilies, OverloadedStrings #-}
-- | A light simplification pass before conversion of desugared code to Core.
module DDC.Source.Tetra.Transform.Prep
( type S, evalState, newVar
, desugarModule)
where
import DDC.Source.Tetra.Module
import DDC.Source.Tetra.Exp
import Data.Monoid
import Data.Text (Text)
import Data.Map (Map)
import qualified Control.Monad.State.Strict as S
import qualified Data.Text as Text
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
---------------------------------------------------------------------------------------------------
-- | State holding a variable name prefix and counter to
-- create fresh variable names.
type S = S.State (Bool, Text, Int)
-- | Evaluate a desguaring computation,
-- using the given prefix for freshly introduced variables.
evalState :: Text -> S a -> a
evalState n c
= S.evalState c (False, n, 0)
-- | Allocate a new named variable, yielding its associated bind and bound.
newVar :: Text -> S (Bind, Bound)
newVar pre
= do (p, n, i) <- S.get
let name = pre <> "$" <> n <> Text.pack (show i)
S.put (p, n, i + 1)
return (BName name, UName name)
-- | Set the progress flag in the state.
progress :: S ()
progress
= do (_, n, i) <- S.get
S.put (True, n, i)
---------------------------------------------------------------------------------------------------
-- | Desguar a module.
--
-- We keep applying the prep transforms we have until they
-- stop making progress.
--
desugarModule :: Module Source -> S (Module Source)
desugarModule mm
= do (_, n, i) <- S.get
S.put (False, n, i)
mm' <- desugarModule1 mm
(p', _, _) <- S.get
if p' then desugarModule mm'
else return mm'
-- | Prepare a source module for conversion to core.
desugarModule1 :: Module Source -> S (Module Source)
desugarModule1 mm
= do ts' <- mapM desugarTop $ moduleTops mm
return $ mm { moduleTops = ts' }
---------------------------------------------------------------------------------------------------
-- | Desugar a top-level definition.
desugarTop :: Top Source -> S (Top Source)
desugarTop tt
= case tt of
TopType{} -> return tt
TopData{} -> return tt
TopClause sp cl -> TopClause sp <$> desugarCl Map.empty cl
---------------------------------------------------------------------------------------------------
-- | Desugar a clause.
desugarCl
:: Map Name Name
-> Clause -> S Clause
desugarCl rns cl
= case cl of
SSig{}
-> return cl
SLet a b ps gxs
-> do ps' <- mapM desugarP ps
gxs' <- mapM (desugarGX rns) gxs
return $ SLet a b ps' gxs'
desugarP :: Param -> S Param
desugarP pp
= case pp of
MType{} -> return pp
MWitness{} -> return pp
MValue w mt -> MValue <$> desugarW w <*> return mt
---------------------------------------------------------------------------------------------------
-- | Desugar a guarded expression.
desugarGX
:: Map Name Name
-> GuardedExp -> S GuardedExp
desugarGX rns gx
= case gx of
GGuard g gx' -> GGuard <$> desugarG rns g <*> desugarGX rns gx'
GExp x -> GExp <$> desugarX rns x
---------------------------------------------------------------------------------------------------
-- | Desugar a guard.
desugarG :: Map Name Name
-> Guard -> S Guard
desugarG rns gg
= case gg of
GPat p x -> GPat <$> desugarW p <*> desugarX rns x
GPred x -> GPred <$> desugarX rns x
GDefault -> return GDefault
---------------------------------------------------------------------------------------------------
-- | Desugar an expression.
desugarX :: Map Name Name -- ^ Renamed bound variables.
-> Exp -> S Exp
desugarX rns xx
= case xx of
-- Lift out nested box casts.
-- This speculatively allocates the inner box,
-- but means it's easier to find (run (box x)) pairs
--
-- let b1 = box (let b2 = box x3
-- in x2)
-- in x1
--
-- => let b2 = box x3 in
-- let b1 = box x2 in
-- x1
--
-- This transform makes b2 scope over x1 where it didn't before,
-- so we rename it along the way to avoid variable clashes.
--
XLet (LLet b1
(XCast CastBox
(XLet (LLet (XBindVarMT (BName n2) mt2)
(XCast CastBox x3))
x2)))
x1
-> do
progress
-- Make a new name for b2 and desugar x2 to force the rename.
(b2', (UName n2')) <- newVar "x"
x2' <- desugarX (Map.insert n2 n2' rns) x2
desugarX rns
$ XLet (LLet (XBindVarMT b2' mt2) (XCast CastBox x3))
$ XLet (LLet b1 (XCast CastBox x2'))
$ x1
-- Eliminate trivial v1 = v2 bindings.
XLet (LLet (XBindVarMT (BName n1) _) (XVar (UName n2))) x1
-> do let rns' = Map.insert n1 n2 rns
progress
desugarX rns' x1
-- The match desugarer introduces case alternatives where the pattern
-- is just a variable, which we can convert to a let-expression.
XCase x0 ( AAltCase (PVar b) [GExp x1] : _)
-> do progress
desugarX rns
$ XLet (LLet (XBindVarMT b Nothing) x0)
$ x1
-- If the first pattern is a default and none of the other alternatives
-- constrain the type of the scrutinee then the core type inferencer
-- won't be able to determine the match type.
XCase _x0 alts@(AAltCase PDefault [GExp x1] : _)
| null [ p | AAltCase p@(PData _ _) _ <- alts]
-> do progress
desugarX rns x1
-- Translate out varible patterns.
-- The core language does not include them, so we bind the
-- scrutinee with a new name and substitute that for the
-- name bound by the variable patterns.
XCase x0 alts
-- Only do the rewrite if at least one expression has
-- such a variable pattern.
| ns <- [n | AAltCase (PVar n) _ <- alts]
, not $ null ns
-> do
progress
-- Desugar the scrutinee.
x0' <- desugarX rns x0
-- New variable to bind the scrutinee.
(b, u@(UName nScrut)) <- newVar "xScrut"
-- For each alternative, if it has a variable pattern
-- then substitute the new name for it in the alternative.
let desugarAlt (AAltCase (PVar (BName n1)) gxs)
= do let rns' = Map.insert n1 nScrut rns
gxs' <- mapM (desugarGX rns') gxs
return $ AAltCase PDefault gxs'
desugarAlt (AAltCase p gxs)
= do gxs' <- mapM (desugarGX rns) gxs
return $ AAltCase p gxs'
alts' <- mapM desugarAlt alts
-- The final expression.
return
$ XLet (LLet (XBindVarMT b Nothing) x0')
$ XCase (XVar u) alts'
-- Eliminate (run (box x)) pairs.
XCast CastBox (XCast CastRun x)
-> do progress
desugarX rns x
-- Lookup renames from the variable rename map.
XVar (UName n0)
-> let sink entered n
= case Map.lookup n rns of
Just n'
| Set.member n' entered
-> n'
| otherwise
-> sink (Set.insert n' entered) n'
Nothing -> n
in do
let n0' = sink Set.empty n0
if n0 /= n0'
then do
progress
return $ XVar (UName n0')
else return xx
-- Convert XWhere to let expressions.
XWhere _sp x cls
-> do x' <- desugarX rns x
cls' <- mapM (desugarCl rns) cls
return $ XLet (LGroup cls') x'
-- Boilerplate.
XAnnot a x -> XAnnot a <$> desugarX rns x
XVar{} -> return xx
XPrim{} -> return xx
XCon{} -> return xx
XLAM mb x -> XLAM mb <$> desugarX rns x
XLam mb x -> XLam mb <$> desugarX rns x
XApp x1 x2 -> XApp <$> desugarX rns x1 <*> desugarX rns x2
XLet lts x -> XLet <$> desugarLts rns lts <*> desugarX rns x
XCase x as -> XCase <$> desugarX rns x <*> mapM (desugarAC rns) as
XCast c x -> XCast c <$> desugarX rns x
XType{} -> return xx
XWitness{} -> return xx
XDefix sp xs -> XDefix sp <$> mapM (desugarX rns) xs
XInfixOp{} -> return xx
XInfixVar{} -> return xx
XMatch sp as x -> XMatch sp <$> mapM (desugarAM rns) as <*> desugarX rns x
XLamPat sp p mt x -> XLamPat sp p mt <$> desugarX rns x
XLamCase sp alts -> XLamCase sp <$> mapM (desugarAC rns) alts
---------------------------------------------------------------------------------------------------
-- | Desugar a case alternative.
desugarAC
:: Map Name Name
-> AltCase -> S AltCase
desugarAC rns aa
= case aa of
AAltCase p gxs
-> AAltCase <$> desugarW p <*> mapM (desugarGX rns) gxs
-- | Desugar a match alternative.
desugarAM
:: Map Name Name
-> AltMatch -> S AltMatch
desugarAM rns (AAltMatch gx)
= AAltMatch <$> desugarGX rns gx
-- | Desugar a pattern.
desugarW :: Pat -> S Pat
desugarW pp
= case pp of
-- Convert var binders where the variable is a wild card to
-- the default pattern. We can't convert plain variable patterns
-- to core.
PVar BNone
-> do progress
return PDefault
PDefault -> return PDefault
PAt b p -> PAt b <$> desugarW p
PVar b -> return $ PVar b
PData dc ps -> PData dc <$> mapM desugarW ps
---------------------------------------------------------------------------------------------------
-- | Desugar some let-bindings.
desugarLts
:: Map Name Name
-> Lets -> S Lets
desugarLts rns lts
= case lts of
LLet mb x -> LLet mb <$> desugarX rns x
LPrivate{} -> return lts
LGroup cls -> LGroup <$> mapM (desugarCl rns) cls
LRec bxs
-> do let (bs, xs) = unzip bxs
xs' <- mapM (desugarX rns) xs
return $ LRec $ zip bs xs'