style changes

anoma · Sep 13, 2024 · ca15171 · ca15171
1 parent ff3b663
commit ca15171
Showing 1 changed file with 11 additions and 8 deletions.
diff --git 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,25 +38,26 @@ 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
           fvi =
             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