Precompute debug operations info in linear time (#3038)

After * #2973 let-folding stopped being linear in the size of the input program. All transformations should be linear whenever possible. This PR makes let-folding linear again.
anoma · Sep 14, 2024 · 66427cd · 66427cd
1 parent a63314d
commit 66427cd
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 4 deletions.
diff --git a/src/Juvix/Compiler/Core/Extra/Utils.hs b/src/Juvix/Compiler/Core/Extra/Utils.hs
@@ -184,7 +184,7 @@ isDebugOp = \case
       OpTrace -> True
       OpFail -> True
       OpSeq -> True
-      OpAssert -> False
+      OpAssert -> True
       OpAnomaByteArrayFromAnomaContents -> False
       OpAnomaByteArrayToAnomaContents -> False
       OpAnomaDecode -> False

diff --git a/src/Juvix/Compiler/Core/Info/DebugOpsInfo.hs b/src/Juvix/Compiler/Core/Info/DebugOpsInfo.hs
@@ -0,0 +1,38 @@
+module Juvix.Compiler.Core.Info.DebugOpsInfo where
+
+import Juvix.Compiler.Core.Extra
+import Juvix.Compiler.Core.Info qualified as Info
+
+newtype DebugOpsInfo = DebugOpsInfo
+  { _infoHasDebugOps :: Bool
+  }
+
+instance IsInfo DebugOpsInfo
+
+kDebugOpsInfo :: Key DebugOpsInfo
+kDebugOpsInfo = Proxy
+
+makeLenses ''DebugOpsInfo
+
+-- | Computes debug operations info for each subnode.
+computeDebugOpsInfo :: Node -> Node
+computeDebugOpsInfo = umap go
+  where
+    go :: Node -> Node
+    go node
+      | isDebugOp node =
+          modifyInfo (Info.insert (DebugOpsInfo True)) node
+      | otherwise =
+          modifyInfo (Info.insert dbi) node
+      where
+        dbi =
+          DebugOpsInfo
+            . or
+            . map (hasDebugOps . (^. childNode))
+            $ children node
+
+getDebugOpsInfo :: Node -> DebugOpsInfo
+getDebugOpsInfo = fromJust . Info.lookup kDebugOpsInfo . getInfo
+
+hasDebugOps :: Node -> Bool
+hasDebugOps = (^. infoHasDebugOps) . getDebugOpsInfo
diff --git a/src/Juvix/Compiler/Core/Transformation/Optimize/LetFolding.hs b/src/Juvix/Compiler/Core/Transformation/Optimize/LetFolding.hs
@@ -15,6 +15,7 @@ module Juvix.Compiler.Core.Transformation.Optimize.LetFolding (letFolding, letFo
 
 import Juvix.Compiler.Core.Data.BinderList qualified as BL
 import Juvix.Compiler.Core.Extra
+import Juvix.Compiler.Core.Info.DebugOpsInfo as Info
 import Juvix.Compiler.Core.Info.FreeVarsInfo as Info
 import Juvix.Compiler.Core.Transformation.Base
 
@@ -28,7 +29,7 @@ convertNode isFoldable md = rmapL go
               || Info.freeVarOccurrences 0 _letBody <= 1
               || isFoldable md bl (_letItem ^. letItemValue)
           )
-            && not (containsDebugOps _letBody) ->
+            && not (Info.hasDebugOps _letBody) ->
             go (recur . (mkBCRemove b val' :)) (BL.cons b bl) _letBody
         where
           val' = go recur bl (_letItem ^. letItemValue)
@@ -39,13 +40,15 @@ convertNode isFoldable md = rmapL go
 letFolding' :: (Module -> BinderList Binder -> Node -> Bool) -> Module -> Module
 letFolding' isFoldable tab =
   mapAllNodes
-    ( removeInfo kFreeVarsInfo
+    ( removeInfo kDebugOpsInfo
+        . removeInfo kFreeVarsInfo
         . convertNode isFoldable tab
-        . computeFreeVarsInfo' 2
         -- 2 is the lambda multiplier factor which guarantees that every free
         -- variable under a lambda is counted at least twice, preventing let
         -- folding for let-bound variables (with non-immediate values) that
         -- occur under lambdas
+        . computeFreeVarsInfo' 2
+        . computeDebugOpsInfo
     )
     tab