ddc-source-tetra-0.4.3.1: DDC/Source/Tetra/Transform/Matches.hs
{-# LANGUAGE TypeFamilies, OverloadedStrings #-}
-- | Desugar match expressions to case expressions.
--
-- In a match expression if matching fails in one block of guards then
-- we skip to the next block. This introduces join point at the start
-- of every block of guards execpt the first one, which we need to flatten
-- out when converting to plain case expressions.
--
-- We also merge multiple clauses for the same function into a single one
-- while we're here.
--
module DDC.Source.Tetra.Transform.Matches
( type S, evalState, newVar
, desugarModule)
where
import DDC.Source.Tetra.Module
import DDC.Source.Tetra.Prim
import DDC.Source.Tetra.Exp
import Data.Monoid
import Data.Text (Text)
import qualified DDC.Data.SourcePos as SP
import qualified Control.Monad.State as S
import qualified Data.Text as Text
-------------------------------------------------------------------------------
-- | Desugar match expressions to case expressions in a module.
desugarModule :: Module Source -> S (Module Source)
desugarModule mm
= do ts' <- desugarTops $ moduleTops mm
return $ mm { moduleTops = ts' }
-------------------------------------------------------------------------------
-- | Desugar top-level definitions.
desugarTops :: [Top Source] -> S [Top Source]
desugarTops ts
= do let tsType = [t | t@TopType{} <- ts]
let tsData = [t | t@TopData{} <- ts]
let spCls = [(sp, cl) | TopClause sp cl <- ts]
-- We may have multiple clauses for the same function in a single
-- group, so we need to pass them all to the clause group
-- desugarer at once.
spCls' <- desugarClGroup spCls
return $ tsType
++ tsData
++ [TopClause sp cl | (sp, cl) <- spCls']
-------------------------------------------------------------------------------
-- | Desugar a clause group.
desugarClGroup :: [(SP, Clause)] -> S [(SP, Clause)]
desugarClGroup spcls0
= loop spcls0
where
-- We've reached the end of the list of clauses.
loop []
= return []
-- Signatures do not need desugaring.
loop ((sp, cl@SSig{}) : cls)
= do cls' <- loop cls
return $ (sp, cl) : cls'
-- We have a let-clause.
loop ( (sp, SLet sp1 (XBindVarMT b1 mt1) ps1 gxs1) : cls)
= loop cls >>= \cls'
-> case cls' of
-- Consecutive clauses are for the same function.
(_, SLet _sp2 (XBindVarMT b2 _mt2) ps2 [GExp xNext]) : clsRest
| b1 == b2
-> do
-- Flatten out guards, wrapping the next expression
-- with case expressions to implement them.
xBody_inner <- flattenGXs gxs1 xNext
-- Recursively desugar the flattened expression
xBody_rec <- desugarX sp xBody_inner
-- Intoduce new let-bindings to handle the case
-- where different clauses name their parameters
-- differently.
(ps1', _ps2', xBody_join)
<- joinParams ps1 ps2 xBody_rec
return $ (sp, SLet sp1 (XBindVarMT b1 mt1) ps1'
[GExp xBody_join])
: clsRest
-- Consecutive clauses are not for the same function.
_ -> do let xError = makeXErrorDefault
(Text.pack $ SP.sourcePosSource sp1)
(fromIntegral $ SP.sourcePosLine sp1)
-- Flatten out guards, wrapping the error expression
-- with case expressions to implement them.
xBody_inner <- flattenGXs gxs1 xError
-- Recursively desugar the flattened expression.
xBody' <- desugarX sp xBody_inner
return $ (sp, SLet sp1 (XBindVarMT b1 mt1) ps1
[GExp xBody'])
: cls'
-- | Given corresponding parameters for earlier and later clauses,
-- introduce let bindings to handle differences in parameter naming.
joinParams :: [Param] -> [Param] -> Exp
-> S ([Param], [Param], Exp)
joinParams [] ps2 xx
= return ([], ps2, xx)
joinParams ps1 [] xx
= return (ps1, [], xx)
joinParams (p1:ps1) (p2:ps2) xx
= do
(p1', p2', mLets) <- joinParam p1 p2
(ps1', ps2', xx') <- joinParams ps1 ps2 xx
case mLets of
Nothing
-> return (p1' : ps1', p2' : ps2', xx')
Just lts
-> return (p1' : ps1', p2' : ps2', XLet lts xx')
-- | Given corresponding parameters for earlier and later clauses,
-- introduce let bindings to handle differences in parameter naming.
joinParam :: Param -> Param
-> S (Param, Param, Maybe Lets)
joinParam p1 p2
= case (p1, p2) of
-- When an earlier pattern does not bind the argument to a variable
-- then we need to introduce a new variable so we can pass the
-- same argument to successive clauses.
( MValue pat1 mt1
, MValue (PVar (BName n2)) mt2)
| isAnonPat pat1
-> do (b, u) <- newVar "m"
let lts = LLet (XBindVarMT (BName n2) mt2) (XVar u)
return (MValue (PVar b) mt1, p2, Just lts)
-- When earlier clauses bind the argument using a different variable
-- than later ones then we need to add a synonym.
( MValue (PVar (BName n1)) _mt1
, MValue (PVar (BName n2)) mt2)
| n1 /= n2
-> do let lts = LLet (XBindVarMT (BName n2) mt2) (XVar (UName n1))
return (p1, p2, Just lts)
_ -> return (p1, p2, Nothing)
-- | Check if this pattern does not bind a variable.
isAnonPat :: Pat -> Bool
isAnonPat pp
= case pp of
PDefault -> True
PVar BAnon -> True
_ -> False
-------------------------------------------------------------------------------
-- | Desugar an expression.
desugarX :: SP -> Exp -> S Exp
desugarX sp xx
= case xx of
-- Boilerplate.
XAnnot sp' x -> XAnnot sp' <$> desugarX sp' x
XVar{} -> pure xx
XPrim{} -> pure xx
XCon{} -> pure xx
XLam b x -> XLam b <$> desugarX sp x
XLAM b x -> XLAM b <$> desugarX sp x
XApp x1 x2 -> XApp <$> desugarX sp x1 <*> desugarX sp x2
XLet lts x -> XLet <$> desugarLts sp lts <*> desugarX sp x
XCast c x -> XCast c <$> desugarX sp x
XType{} -> pure xx
XWitness{} -> pure xx
XDefix a xs -> XDefix a <$> mapM (desugarX sp) xs
XInfixOp{} -> pure xx
XInfixVar{} -> pure xx
-- Desugar case expressions.
XCase x alts
-> XCase <$> desugarX sp x
<*> mapM (desugarAC sp) alts
-- Desugar match expressions into case expressions.
XMatch _ alts xFail
-> do let gxs = [gx | AAltMatch gx <- alts]
xFlat <- flattenGXs gxs xFail
xFlat' <- desugarX sp xFlat
return xFlat'
XWhere sp' x cls
-> do x' <- desugarX sp' x
let spcls = [(sp', cl) | cl <- cls]
spcls' <- desugarClGroup spcls
return $ XWhere sp' x' (map snd spcls')
XLamPat sp' w mt x
-> XLamPat sp' w mt <$> desugarX sp x
XLamCase sp' alts
-> XLamCase sp' <$> mapM (desugarAC sp) alts
-------------------------------------------------------------------------------
-- | Desugar some let bindings.
desugarLts :: SP -> Lets -> S Lets
desugarLts sp lts
= case lts of
LLet mb x -> LLet mb <$> desugarX sp x
LRec bxs
-> do let (bs, xs) = unzip bxs
xs' <- mapM (desugarX sp) xs
return $ LRec $ zip bs xs'
LPrivate{} -> return lts
LGroup cls
-> do let spcls = zip (repeat sp) cls
spcls' <- desugarClGroup spcls
return $ LGroup $ map snd spcls'
-------------------------------------------------------------------------------
-- | Desugar a guarded expression.
desugarGX :: SP -> GuardedExp -> S GuardedExp
desugarGX sp gx
= case gx of
GGuard g gx' -> GGuard <$> desugarG sp g <*> desugarGX sp gx'
GExp x -> GExp <$> desugarX sp x
-------------------------------------------------------------------------------
-- | Desugar a guard.
desugarG :: SP -> Guard -> S Guard
desugarG sp g
= case g of
GPat p x -> GPat p <$> desugarX sp x
GPred x -> GPred <$> desugarX sp x
GDefault -> pure GDefault
-------------------------------------------------------------------------------
-- | Desugar a case alternative.
desugarAC :: SP -> AltCase -> S AltCase
desugarAC sp (AAltCase p gxs)
= do gxs' <- mapM (desugarGX sp) gxs
return $ AAltCase p gxs'
-------------------------------------------------------------------------------
-- | Desugar some guards to a case-expression.
-- At runtime, if none of the guards match then run the provided
-- fall-though computation.
flattenGXs :: [GuardedExp] -> Exp -> S Exp
flattenGXs gs0 fail0
= go gs0 fail0
where
-- Desugar list of guarded expressions.
go [] cont
= return cont
go [g] cont
= go1 g cont
go (g : gs) cont
= do gs' <- go gs cont
go1 g gs'
-- Desugar single guarded expression.
go1 (GExp x1) _
= return x1
go1 (GGuard GDefault gs) cont
= go1 gs cont
-- Simple cases where we can avoid introducing the continuation.
go1 (GGuard (GPred g1) (GExp x1)) cont
= return
$ XCase g1 [ AAltCase PTrue [GExp x1]
, AAltCase PDefault [GExp cont] ]
go1 (GGuard (GPat p1 g1) (GExp x1)) cont
= return
$ XCase g1 [ AAltCase p1 [GExp x1]
, AAltCase PDefault [GExp cont]]
-- Cases that use a continuation function as a join point.
-- We need this when desugaring general pattern alternatives,
-- as each group of guards can be reached from multiple places.
go1 (GGuard (GPred x1) gs) cont
= do (b, u) <- newVar "m"
x' <- go1 gs (XRun (XVar u))
return
$ XLet (LLet (XBindVarMT b Nothing) (XBox cont))
$ XCase x1 [ AAltCase PTrue [GExp x']
, AAltCase PDefault [GExp (XRun (XVar u)) ]]
go1 (GGuard (GPat p1 x1) gs) cont
= do (b, u) <- newVar "m"
x' <- go1 gs (XRun (XVar u))
return
$ XLet (LLet (XBindVarMT b Nothing) (XBox cont))
$ XCase x1 [ AAltCase p1 [GExp x']
, AAltCase PDefault [GExp (XRun (XVar u)) ]]
-------------------------------------------------------------------------------
-- | Source position.
type SP = SP.SourcePos
-- | State holding a variable name prefix and counter to
-- create fresh variable names.
type S = S.State (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 (n, 0)
-- | Allocate a new named variable, yielding its associated bind and bound.
newVar :: Text -> S (Bind, Bound)
newVar pre
= do (n, i) <- S.get
let name = pre <> "$" <> n <> Text.pack (show i)
S.put (n, i + 1)
return (BName name, UName name)