From 01fb15786ce0abece717c66e0a2626c7c28892a0 Mon Sep 17 00:00:00 2001
From: Lukasz Czajka <lukasz@heliax.dev>
Date: Fri, 13 Sep 2024 14:55:16 +0200
Subject: [PATCH] style changes

---
 src/Juvix/Compiler/Core/Info/FreeVarsInfo.hs | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/Juvix/Compiler/Core/Info/FreeVarsInfo.hs b/src/Juvix/Compiler/Core/Info/FreeVarsInfo.hs
index b9be178967..312c6cfe61 100644
--- a/src/Juvix/Compiler/Core/Info/FreeVarsInfo.hs
+++ b/src/Juvix/Compiler/Core/Info/FreeVarsInfo.hs
@@ -21,6 +21,8 @@ makeLenses ''FreeVarsInfo
 computeFreeVarsInfo :: Node -> Node
 computeFreeVarsInfo = computeFreeVarsInfo' 1
 
+-- | `lambdaMultiplier` specifies how much to multiply the free variable count
+-- for variables under lambdas
 computeFreeVarsInfo' :: Int -> Node -> Node
 computeFreeVarsInfo' lambdaMultiplier = umap go
   where
@@ -36,9 +38,7 @@ computeFreeVarsInfo' lambdaMultiplier = umap go
           fvi =
             FreeVarsInfo
               . fmap (* lambdaMultiplier)
-              . Map.mapKeysMonotonic (\idx -> idx - 1)
-              . Map.filterWithKey (\idx _ -> idx >= 1)
-              $ getFreeVarsInfo _lambdaBody ^. infoFreeVars
+              $ getFreeVars 1 _lambdaBody
       _ ->
         modifyInfo (Info.insert fvi) node
         where
@@ -46,15 +46,18 @@ computeFreeVarsInfo' lambdaMultiplier = umap go
             FreeVarsInfo $
               foldr
                 ( \NodeChild {..} acc ->
-                    Map.unionWith (+) acc
-                      . Map.mapKeysMonotonic (\idx -> idx - _childBindersNum)
-                      . Map.filterWithKey
-                        (\idx _ -> idx >= _childBindersNum)
-                      $ getFreeVarsInfo _childNode ^. infoFreeVars
+                    Map.unionWith (+) acc $
+                      getFreeVars _childBindersNum _childNode
                 )
                 mempty
                 (children node)
 
+    getFreeVars :: Int -> Node -> Map Index Int
+    getFreeVars bindersNum node =
+      Map.mapKeysMonotonic (\idx -> idx - bindersNum)
+        . Map.filterWithKey (\idx _ -> idx >= bindersNum)
+        $ getFreeVarsInfo node ^. infoFreeVars
+
 getFreeVarsInfo :: Node -> FreeVarsInfo
 getFreeVarsInfo = fromJust . Info.lookup kFreeVarsInfo . getInfo