|
| 1 | +module HelVM.HelPS.HS2Lazy.PatComp.PatternMatchCompiler ( |
| 2 | + compilePatternMatch |
| 3 | + , patBindings |
| 4 | +) where |
| 5 | + |
| 6 | +import HelVM.HelPS.HS2Lazy.Unsafe |
| 7 | + |
| 8 | +import HS2Lazy.PPrint () |
| 9 | +import HS2Lazy.Syntax |
| 10 | + |
| 11 | +import Data.List (partition) |
| 12 | + |
| 13 | +import Prelude hiding (Alt, Ap, Const) |
| 14 | +import qualified Relude.Unsafe as Unsafe |
| 15 | + |
| 16 | + |
| 17 | +type PatComp = State Int |
| 18 | + |
| 19 | +compilePatternMatch :: Program -> Program |
| 20 | +compilePatternMatch pgm = evalState (pcProgram pgm) 0 |
| 21 | + |
| 22 | +pcProgram :: Program -> PatComp Program |
| 23 | +pcProgram = traverse pcBindGroup |
| 24 | + |
| 25 | +pcBindGroup :: BindGroup -> PatComp BindGroup |
| 26 | +pcBindGroup (es, iss) = (,) <$> traverse pcExpl es <*> traverse (traverse pcImpl) iss |
| 27 | + |
| 28 | +pcExpl :: Expl -> PatComp Expl |
| 29 | +pcExpl (i, sc, alts) = go <$> pcAlts alts where go alt = (i, sc, [alt]) |
| 30 | + |
| 31 | +pcImpl :: Impl -> PatComp Impl |
| 32 | +pcImpl (i, alts) = go <$> pcAlts alts where go alt = (i, [alt]) |
| 33 | + |
| 34 | +pcAlts :: [Alt] -> PatComp Alt |
| 35 | +pcAlts [(ps, Rhs e)] | all isPVar ps = go <$> pcExpr e where go e' = (ps , Rhs e') |
| 36 | +pcAlts qs = go =<< newVars (length $ fst $ unsafeHead qs) where |
| 37 | + go us = go' <$> match us qs matchError where |
| 38 | + go' rhs = (map PVar us , Rhs rhs) |
| 39 | + |
| 40 | +isPVar :: Pat -> Bool |
| 41 | +isPVar (PVar _) = True |
| 42 | +isPVar _ = False |
| 43 | + |
| 44 | +pcRhs :: Rhs -> Expr -> PatComp Expr |
| 45 | +pcRhs (Rhs e) _ = pcExpr e |
| 46 | +pcRhs (Where bg rhs) def = fmap (Let bg) (pcRhs rhs def) |
| 47 | +pcRhs (Guarded gds) def = foldr makeIf def <$> traverse pcGuard gds |
| 48 | + |
| 49 | +pcGuard :: (Expr, Expr) -> PatComp (Expr, Expr) |
| 50 | +pcGuard (e1, e2) = (,) <$> pcExpr e1 <*> pcExpr e2 |
| 51 | + |
| 52 | +pcExpr :: Expr -> PatComp Expr |
| 53 | +pcExpr (Ap e1 e2) = Ap <$> pcExpr e1 <*> pcExpr e2 |
| 54 | +pcExpr (Let bg e) = Let <$> pcBindGroup bg <*> pcExpr e |
| 55 | +pcExpr (Lambda a) = fmap Lambda (pcAlts [a]) |
| 56 | +pcExpr (Case e pes) = go =<< pcExpr e where |
| 57 | + go (Var v) = match' v |
| 58 | + go e'' = goLet e'' =<< newVar |
| 59 | + goLet e'' v = Let (bind1 v e'') <$> match' v |
| 60 | + match' v' = match [v'] qs matchError |
| 61 | + qs = [([p] , rhs) | (p , rhs) <- pes] |
| 62 | +pcExpr (ESign e sc) = ESign <$> pcExpr e <*> pure sc |
| 63 | +pcExpr c = pure c |
| 64 | + |
| 65 | +matchError :: Expr |
| 66 | +matchError = Ap (Var "error") (Lit $ LitStr "Non-exhaustive patterns") |
| 67 | + |
| 68 | +patBindings :: Expr -> Pat -> [Impl] |
| 69 | +patBindings v (PVar i) = [(i, [([], Rhs v)])] |
| 70 | +patBindings _ PWildcard = [] |
| 71 | +patBindings v (PAs i p) = (i, [([], Rhs v)]) : patBindings v p |
| 72 | +patBindings _ (PLit _) = [] |
| 73 | +patBindings v (PCon con pats) = concat [patBindings (makeSel con n v) p | (p, n) <- zip pats [1..]] |
| 74 | + |
| 75 | +makeSel :: Const -> Int -> Expr -> Expr |
| 76 | +makeSel con i e = expr where |
| 77 | + vs = ["@@" ++ show v | v <- [1 .. (conArity con)]] |
| 78 | + body = Rhs $ Var $ vs Unsafe.!! (i - 1) |
| 79 | + receiver' = receiver vs body |
| 80 | + expr = ap e [ifReceiver i' con receiver' eError | i' <- [1 .. (tyconNumCon $ conTycon con)]] |
| 81 | + |
| 82 | +type Equation = Alt |
| 83 | + |
| 84 | +isVar :: Equation -> Bool |
| 85 | +isVar (p : _ , _) = test p where |
| 86 | + test (PVar _) = True |
| 87 | + test PWildcard = True |
| 88 | + test (PAs _ p') = test p' |
| 89 | + test (PLit _) = False |
| 90 | + test (PCon _ _) = False |
| 91 | +isVar _ = error "isVar" |
| 92 | + |
| 93 | +match :: [Id] -> [Equation] -> Expr -> PatComp Expr |
| 94 | +match [] qs def = foldrM pcRhs def (map snd qs) |
| 95 | +match us qs def = foldrM (matchVarCon us) def (partitionEqns isVar qs) |
| 96 | + |
| 97 | +matchVarCon :: [Id] -> [Equation] -> Expr -> PatComp Expr |
| 98 | +matchVarCon us@(u : _) qs def = go $ unsafeHead $ fst $ unsafeHead qs' where |
| 99 | + go (PLit _) = bindDefault (matchLit us qs') def |
| 100 | + go (PCon _ _) = bindDefault (matchCon us qs') def |
| 101 | + go _ = matchVar us qs' def |
| 102 | + qs' = map sub qs |
| 103 | + sub (PAs v p : ps, rhs) = sub (p : ps, Where (bind1 v (Var u)) rhs) |
| 104 | + sub (ps, rhs) = (ps, rhs) |
| 105 | +matchVarCon _ _ _ = error "matchVarCon" |
| 106 | + |
| 107 | +matchVar :: [Id] -> [Equation] -> Expr -> PatComp Expr |
| 108 | +matchVar (u : us) qs def = match us (map sub qs) def where |
| 109 | + sub (PVar v : ps , rhs) = (ps , Where (bind1 v (Var u)) rhs) |
| 110 | + sub (PWildcard : ps , rhs) = (ps , rhs) |
| 111 | + sub _ = error "sub" |
| 112 | +matchVar _ _ _ = error "matchVar" |
| 113 | + |
| 114 | +bindDefault :: (Expr -> PatComp Expr) -> Expr -> PatComp Expr |
| 115 | +bindDefault f def |
| 116 | + | simple def = f def |
| 117 | + | otherwise = go =<< newVar |
| 118 | + where |
| 119 | + go v = Let (bind1 v def) <$> f (Var v) |
| 120 | + simple _ = True |
| 121 | + |
| 122 | +matchLit :: [Id] -> [Equation] -> Expr -> PatComp Expr |
| 123 | +matchLit us qs def = foldr makeIf def <$> traverse (matchLitClause us def) (groupLit qs) |
| 124 | + |
| 125 | +matchLitClause :: [Id] -> Expr -> (Literal, [Equation]) -> PatComp (Expr, Expr) |
| 126 | +matchLitClause (u : us) def (lit , qs) = go <$> match us [(ps , rhs) | (_ : ps , rhs) <- qs] def |
| 127 | + where go e = (ap (Var "&eq") [Var u , Lit lit] , e) |
| 128 | +matchLitClause _ _ _ = error "matchLitClause" |
| 129 | + |
| 130 | +groupLit :: [Equation] -> [(Literal, [Equation])] |
| 131 | +groupLit [] = [] |
| 132 | +groupLit qs@((PLit l:_,_):_) = (l, qs') : groupLit qs'' where |
| 133 | + (qs', qs'') = partition go qs |
| 134 | + go (PLit l' : _ , _) = l == l' |
| 135 | + go _ = False |
| 136 | +groupLit _ = error "groupLit" |
| 137 | + |
| 138 | +matchCon :: [Id] -> [Equation] -> Expr -> PatComp Expr |
| 139 | +matchCon us qs def |
| 140 | + | isCovered grps = flip (foldr makeIf) <$> clauses <*> lastClause |
| 141 | + | otherwise = foldr makeIf def <$> traverse (matchConClause us def) grps |
| 142 | + where |
| 143 | + clauses = traverse (matchConClause us def) (Unsafe.init grps) |
| 144 | + lastClause = matchConLastClause us def (Unsafe.last grps) |
| 145 | + grps = groupCon qs |
| 146 | + |
| 147 | +matchConClause :: [Id] -> Expr -> (Const, [Equation]) -> PatComp (Expr, Expr) |
| 148 | +matchConClause (u : us) def (con, qs) = go =<< newVars (conArity con) where |
| 149 | + go us' = go' <$> match (us' ++ us) [(ps' ++ ps, rhs) | (PCon _ ps' : ps , rhs) <- qs] def where |
| 150 | + go' body = (cond, expr) where |
| 151 | + cond = makeTagEq con (Var u) |
| 152 | + expr = ap (Var u) [ifReceiver i con receiver' eError | i <- [1 .. (tyconNumCon $ conTycon con)]] |
| 153 | + receiver' = receiver us' $ Rhs body |
| 154 | +matchConClause _ _ _ = error "matchConClause" |
| 155 | + |
| 156 | +matchConLastClause :: [Id] -> Expr -> (Const, [Equation]) -> PatComp Expr |
| 157 | +matchConLastClause us def grp = snd <$> matchConClause us def grp |
| 158 | + |
| 159 | +groupCon :: [Equation] -> [(Const, [Equation])] |
| 160 | +groupCon [] = [] |
| 161 | +groupCon qs@((PCon c _ : _ , _) : _) = (c , qs') : groupCon qs'' where |
| 162 | + (qs' , qs'') = partition go qs |
| 163 | + go (PCon c' _ : _ , _) = c == c' |
| 164 | + go _ = False |
| 165 | +groupCon _ = error "groupCon" |
| 166 | + |
| 167 | +isCovered :: [(Const, [Equation])] -> Bool |
| 168 | +isCovered grps = go $ tyconNumCon $ conTycon $ fst $ unsafeHead grps where |
| 169 | + go n = n == 0 || length grps == n |
| 170 | + |
| 171 | +newVars :: Int -> PatComp [Id] |
| 172 | +newVars k = go =<< get where go n = put (n + k) $> ['@' : show (n + i) | i <- [1 .. k]] |
| 173 | + |
| 174 | +newVar :: PatComp Id |
| 175 | +newVar = go =<< get where go n = put (n + 1) $> ('@' : show (n + 1)) |
| 176 | + |
| 177 | +makeIf :: (Expr , Expr) -> Expr -> Expr |
| 178 | +makeIf (c , e) e' = ap (Var "IF") [c , e , e'] |
| 179 | + |
| 180 | +makeTagEq :: Const -> Expr -> Expr |
| 181 | +makeTagEq con e = ap e es where |
| 182 | + arities = tyconArities (conTycon con) |
| 183 | + es = [test a (b == conTag con) | (a, b) <- zip arities [1..]] |
| 184 | + test arity b = Lambda ([PVar ('_':show n) | n <- [1 .. arity]] , Rhs $ ifThenElse b eTrue eFalse) |
| 185 | + |
| 186 | +bind1 :: Id -> Expr -> BindGroup |
| 187 | +bind1 v e = ([], [[(v, [([], Rhs e)])]]) |
| 188 | + |
| 189 | +ifReceiver :: Int -> Const -> p -> p -> p |
| 190 | +ifReceiver i con = ifThenElse $ i == conTag con |
| 191 | + |
| 192 | +receiver :: [Id] -> Rhs -> Expr |
| 193 | +receiver vs body = Lambda ([PVar v | v <- vs] , body) |
| 194 | + |
| 195 | +eError :: Expr |
| 196 | +eError = Ap (Var "error") (Lit $ LitStr "!?") |
| 197 | + |
| 198 | +---- |
| 199 | + |
| 200 | +partitionEqns :: Eq b => (a -> b) -> [a] -> [[a]] |
| 201 | +partitionEqns _ [] = [] |
| 202 | +partitionEqns _ [x] = [[x]] |
| 203 | +partitionEqns f (x : xs@(x' : _)) |
| 204 | + | f x == f x' = tack x (partitionEqns f xs) |
| 205 | + | otherwise = [x] : partitionEqns f xs |
| 206 | + |
| 207 | +foldrM :: Monad m => (b -> a -> m a) -> a -> [b] -> m a |
| 208 | +foldrM f a xs = foldlM (flip f) a (reverse xs) |
| 209 | + |
| 210 | +ifThenElse :: Bool -> a -> a -> a |
| 211 | +ifThenElse True x _ = x |
| 212 | +ifThenElse False _ y = y |
0 commit comments