implement case simplification

TODO.md

@@ -18,8 +18,8 @@
 simplifaction transformations required by the codegen
 
 - [ ] vectorisation
-- [ ] case simplification
+- [x] case simplification
 - [x] split fetch operation
 - [ ] right hoist fetch operation
-- [ ] register introduction
+- [x] register introduction
 

src/Grin.hs

@@ -27,7 +27,6 @@ data Exp
   | SReturn     Val
   | SStore      Val
   | SFetchI     Name (Maybe Int) -- fetch a full node or a single node item
---  | SFetchItem  Name Int
   | SUpdate     Name Val
   | SBlock      Exp
   -- Alt

src/Transformations.hs

@@ -319,24 +319,27 @@ type SubstMap = Map SimpleVal SimpleVal
 
 mapVals :: (Val -> Val) -> Exp -> Exp
 mapVals f = \case
-  SReturn val -> SReturn $ f val
+  ECase val alts -> ECase (f val) alts
   SApp name vals -> SApp name (map f vals)
-  exp -> exp -- TODO
+  SReturn val -> SReturn $ f val
+  SStore val -> SStore $ f val
+  SUpdate name val -> SUpdate name $ f val
+  exp -> exp
 
-substVals :: SubstMap -> Exp -> Exp
-substVals env = mapVals subst
-  where subst x = Map.findWithDefault x x env
+substExpVals :: SubstMap -> Exp -> Exp
+substExpVals env = mapVals (subst env)
+
+subst env x = Map.findWithDefault x x env
 
--- TODO: subst map composition
 caseSimplification :: Exp -> Exp
 caseSimplification e = ana builder (mempty, e) where
   builder :: (SubstMap, Exp) -> ExpF (SubstMap, Exp)
   builder (env, exp) =
     case exp of
-      ECase (VarTagNode tagVar vals) alts -> ECaseF (Var tagVar) (map (substAlt env vals) alts)
-      e -> (env,) <$> project (substVals env e)
+      ECase (VarTagNode tagVar vals) alts -> ECaseF (subst env $ Var tagVar) (map (substAlt env vals) alts)
+      e -> (env,) <$> project (substExpVals env e)
 
   substAlt env vals = \case
       Alt (NodePat tag vars) e -> (altEnv, Alt (TagPat tag) e)
-                             where altEnv = foldl' (\m (name,val) -> Map.insert (Var name) val m) env (zip vars vals)
+                             where altEnv = foldl' (\m (name,val) -> Map.insert (Var name) (subst env val) m) env (zip vars vals)
       alt -> (env, alt)