From f4bb6a52c31d8c2e349a62b963a23cd9bc539aad Mon Sep 17 00:00:00 2001
From: Lukasz Czajka <lukasz@heliax.dev>
Date: Fri, 13 Sep 2024 19:30:08 +0200
Subject: [PATCH] precompute debug ops info

---
 src/Juvix/Compiler/Core/Extra/Utils.hs        |  2 +-
 src/Juvix/Compiler/Core/Info/DebugOpsInfo.hs  | 38 +++++++++++++++++++
 .../Transformation/Optimize/LetFolding.hs     |  9 +++--
 3 files changed, 45 insertions(+), 4 deletions(-)
 create mode 100644 src/Juvix/Compiler/Core/Info/DebugOpsInfo.hs

diff --git a/src/Juvix/Compiler/Core/Extra/Utils.hs b/src/Juvix/Compiler/Core/Extra/Utils.hs
index de26939d97..183d4ba247 100644
--- 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
new file mode 100644
index 0000000000..14e0cc61e2
--- /dev/null
+++ 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
index c8752b1398..ab9fa789af 100644
--- 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