From 8463dc28875114971f8608b096c5251d13b5f7a3 Mon Sep 17 00:00:00 2001
From: Martin Fredin <fredin.martin@gmail.com>
Date: Sat, 29 Apr 2023 21:58:39 +0200
Subject: [PATCH] Small fix to lambda lifter

---
 src/LambdaLifter.hs       | 17 ++---------
 tests/TestLambdaLifter.hs | 60 +++++++++++++++++----------------------
 2 files changed, 29 insertions(+), 48 deletions(-)

diff --git a/src/LambdaLifter.hs b/src/LambdaLifter.hs
index 83d3466..5581814 100644
--- a/src/LambdaLifter.hs
+++ b/src/LambdaLifter.hs
@@ -28,17 +28,6 @@ lambdaLift (Program ds) = Program (datatypes ++ binds)
                                        _       -> False
     binds = map DBind $ (collectScs . abstract . freeVars) [b | DBind b <- ds]
 
--- lambdaLift (Program defs) = trace (printTree abst) $ Program $ datatypes ++ ll binds
---   where
---     abst = abstract frees
---     frees = freeVars [b | DBind b@(Bind (Ident s, _) _ _) <- binds, s == "f"]
---
---     ll = map DBind . collectScs . abstract . freeVars . map (\(DBind b) -> b)
---     (binds, datatypes) = partition isBind defs
---     isBind = \case
---       DBind _ -> True
---       _       -> False
-
 -- | Annotate free variables
 freeVars :: [Bind] -> [ABind]
 freeVars binds = [ let ae  = freeVarsExp [] e
@@ -172,9 +161,9 @@ abstractAnnExp Ann {frees, term = (annae, typ) } = case annae of
     AAbs x annae' -> do
         i <- nextNumber
         rhs <- abstractAnnExp annae''
-        let sc_name = Ident ("sc_" ++ show i)
-            sc      = (ELet (Bind (sc_name, typ) vars rhs) (EVar sc_name, typ), typ)
-        pure $ foldl applyFree sc frees
+        let sc_name  = Ident ("sc_" ++ show i)
+            e@(_, t) = foldl applyFree (EVar sc_name, typ) frees
+        pure (ELet (Bind (sc_name, typ) vars rhs) e ,t)
 
       where
         vars = frees <| (x, t_x) <|| ys
diff --git a/tests/TestLambdaLifter.hs b/tests/TestLambdaLifter.hs
index 79c78b2..d209819 100644
--- a/tests/TestLambdaLifter.hs
+++ b/tests/TestLambdaLifter.hs
@@ -58,49 +58,41 @@ abs_1 = undefined
 
 
 
-runPrintFree = print $ freeVarsExp [] (EAbs "x" (EVar "x", TVar' "a"), TVar' "a")
-
+runFreeVars = either putStrLn print (runFree s2)
 runAbstract = either putStrLn (putStrLn . printTree) (runAbs s2)
-  where
-    s = unlines [ "add : Int -> Int -> Int"
-                , "f : Int -> Int -> Int"
-                , "f x y = add x y"
-                , "f = \\x. (\\y. add x y)"
-                ]
-
-    s2 = unlines [ "data List (a) where"
-                 , "    Nil  : List (a)"
-                 , "    Cons : a -> List (a) -> List (a)"
-                 , "map : (a -> b) -> List (a) -> List (b)"
-                 , "add : Int -> Int -> Int"
-
-                 , "f : List (Int)"
-                 , "f = (\\x.\\ys. map (\\y. add y x) ys) 4 (Cons 1 (Cons 2 Nil))"
-                 ]
+runCollect = either putStrLn (putStrLn . printTree) (run s2)
 
 
-runCollect = either putStrLn (putStrLn . printTree) (run s)
-  where
-    s = unlines [ "data List (a) where"
-                , "    Nil  : List (a)"
-                , "    Cons : a -> List (a) -> List (a)"
-                , "add : Int -> Int -> Int"
-                , "map : (a -> b) -> List (a) -> List (b)"
-                , "map f xs = case xs of"
-                , "    Nil       => Nil"
-                , "    Cons x xs => Cons (f x) (map f xs)"
+s1 = unlines [ "add : Int -> Int -> Int"
+             , "f : Int -> Int -> Int"
+             , "f x y = add x y"
+             , "f = \\x. (\\y. add x y)"
+             ]
+
+s2 = unlines [ "data List (a) where"
+             , "    Nil  : List (a)"
+             , "    Cons : a -> List (a) -> List (a)"
+             , "add : Int -> Int -> Int"
+             , "map : (a -> b) -> List (a) -> List (b)"
+          -- , "map f xs = case xs of"
+          -- , "    Nil       => Nil"
+          -- , "    Cons x xs => Cons (f x) (map f xs)"
+
+             , "f : List (Int)"
+             , "f = (\\x.\\ys. map (\\y. add y x) ys) 4 (Cons 1 (Cons 2 Nil))"
+             ]
+
+s3 = "main = (\\plussq. (\\f. f (f 0)) (plussq 3)) (\\x. \\y. y + x + x)"
 
-                , "f : List (Int)"
-                , "f = (\\x.\\ys. map (\\y. add y x) ys) 4 (Cons 1 (Cons 2 Nil))"
-                ]
 
 
 run = fmap collectScs . runAbs
 
-runAbs s = do
-  Program ds <- run' s
-  pure $ (abstract . freeVars) [b | DBind b <- ds]
+runAbs = fmap abstract . runFree
 
+runFree s = do
+  Program ds <- run' s
+  pure $ freeVars [b | DBind b <- ds]
 
 run' =   fmap removeForall
       . reportTEVar