lambdaMultiplier in FreeVarsInfo computation

src/Juvix/Compiler/Core/Info/FreeVarsInfo.hs

@@ -19,26 +19,38 @@ makeLenses ''FreeVarsInfo
 -- | Computes free variable info for each subnode. Assumption: no subnode is a
 -- closure.
 computeFreeVarsInfo :: Node -> Node
-computeFreeVarsInfo = umap go
+computeFreeVarsInfo = computeFreeVarsInfo' 1
+
+computeFreeVarsInfo' :: Int -> Node -> Node
+computeFreeVarsInfo' lambdaMultiplier = umap go
   where
     go :: Node -> Node
     go node = case node of
       NVar Var {..} ->
         mkVar (Info.insert fvi _varInfo) _varIndex
         where
           fvi = FreeVarsInfo (Map.singleton _varIndex 1)
+      NLam Lambda {..} ->
+        modifyInfo (Info.insert fvi) node
+        where
+          fvi =
+            FreeVarsInfo
+              . fmap (* lambdaMultiplier)
+              . Map.mapKeysMonotonic (\idx -> idx - 1)
+              . Map.filterWithKey (\idx _ -> idx >= 1)
+              $ getFreeVarsInfo _lambdaBody ^. infoFreeVars
       _ ->
         modifyInfo (Info.insert fvi) node
         where
           fvi =
             FreeVarsInfo $
               foldr
                 ( \NodeChild {..} acc ->
-                    Map.unionWith (+) acc $
-                      Map.mapKeysMonotonic (\idx -> idx - _childBindersNum) $
-                        Map.filterWithKey
-                          (\idx _ -> idx >= _childBindersNum)
-                          (getFreeVarsInfo _childNode ^. infoFreeVars)
+                    Map.unionWith (+) acc
+                      . Map.mapKeysMonotonic (\idx -> idx - _childBindersNum)
+                      . Map.filterWithKey
+                        (\idx _ -> idx >= _childBindersNum)
+                      $ getFreeVarsInfo _childNode ^. infoFreeVars
                 )
                 mempty
                 (children node)

src/Juvix/Compiler/Core/Transformation/Optimize/LetFolding.hs

@@ -40,7 +40,7 @@ letFolding' isFoldable tab =
   mapAllNodes
     ( removeInfo kFreeVarsInfo
         . convertNode isFoldable tab
-        . computeFreeVarsInfo
+        . computeFreeVarsInfo' 2
     )
     tab
 

tests/Compilation/positive/test059.juvix

@@ -1,15 +1,15 @@
 -- builtin list
 module test059;
 
-import Stdlib.Prelude open hiding {head};
+import Stdlib.Prelude open;
 
 mylist : List Nat := [1; 2; 3 + 1];
 
 mylist2 : List (List Nat) := [[10]; [2]; 3 + 1 :: nil];
 
-head : {a : Type} -> a -> List a -> a
+head' : {a : Type} -> a -> List a -> a
   | a [] := a
   | a [x; _] := x
   | _ (h :: _) := h;
 
-main : Nat := head 50 mylist + head 50 (head [] mylist2);
+main : Nat := head' 50 mylist + head' 50 (head' [] mylist2);