Tweak and add tests

danieljharvey · Dec 13, 2023 · 38dd883 · 38dd883
1 parent d8d6b2e
commit 38dd883
Show file tree

Hide file tree

Showing 6 changed files with 268 additions and 219 deletions.
diff --git a/smol-core/src/Smol/Core/Transform.hs b/smol-core/src/Smol/Core/Transform.hs
@@ -14,4 +14,17 @@ transform ::
   (Ord ann, Ord (dep Identifier)) =>
   Expr dep ann ->
   Expr dep ann
-transform = removeUnused . etaReduce . betaReduce . floatDown . floatUp . flattenLets . inline
+transform = transformPass . transformPass . transformPass
+
+transformPass ::
+  (Ord ann, Ord (dep Identifier)) =>
+  Expr dep ann ->
+  Expr dep ann
+transformPass =
+  removeUnused
+    . etaReduce
+    . floatDown
+    . flattenLets
+    . floatUp
+    . betaReduce
+    . inline
diff --git a/smol-core/src/Smol/Core/Transform/BetaReduce.hs b/smol-core/src/Smol/Core/Transform/BetaReduce.hs
@@ -15,6 +15,8 @@ betaReduceInternal (EApp ann (EAnn _ _ (ELambda _ann ident body)) val) =
   betaReduceInternal $ ELet ann ident val (betaReduceInternal body)
 betaReduceInternal (EApp annA (EApp annB (ELambda _ identA (ELambda _ identB body)) valA) valB) =
   betaReduceInternal $ ELet annA identA valA (ELet annB identB valB (betaReduceInternal body))
+betaReduceInternal (EApp annA (EApp annB (EAnn _ _ (ELambda _ identA (ELambda _ identB body))) valA) valB) =
+  betaReduceInternal $ ELet annA identA valA (ELet annB identB valB (betaReduceInternal body))
 betaReduceInternal (EIf _ (EPrim _ (PBool True)) thenExpr _) =
   betaReduceInternal thenExpr
 betaReduceInternal (EIf _ (EPrim _ (PBool False)) _ elseExpr) =

diff --git a/smol-core/src/Smol/Core/Transform/FlattenLets.hs b/smol-core/src/Smol/Core/Transform/FlattenLets.hs
@@ -5,7 +5,7 @@ import Smol.Core.Types
 
 -- | We don't want `let a = (let b = 1 in b + 1) in a + 1`
 -- instead we want `let b = 1; let a = b + 1; a + 1
-flattenLets :: Expr var ann -> Expr var ann
+flattenLets :: (Eq (dep Identifier)) => Expr dep ann -> Expr dep ann
 flattenLets (ELet ann ident (ELet ann' ident' expr' body') body) =
   flattenLets $
     ELet
@@ -18,4 +18,6 @@ flattenLets (ELet ann ident (ELet ann' ident' expr' body') body) =
           (flattenLets body')
           (flattenLets body)
       )
+-- also, turn `let a = b in a` into `b`
+flattenLets (ELet _ ident body (EVar _ ident')) | ident == ident' = body
 flattenLets other = mapExpr flattenLets other