volatility info

anoma · Nov 28, 2024 · a08fd7f · a08fd7f
1 parent c45bd9a
commit a08fd7f
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 10 deletions.
diff --git a/src/Juvix/Compiler/Core/Info/VolatilityInfo.hs b/src/Juvix/Compiler/Core/Info/VolatilityInfo.hs
@@ -0,0 +1,21 @@
+module Juvix.Compiler.Core.Info.VolatilityInfo where
+
+import Juvix.Compiler.Core.Extra
+import Juvix.Compiler.Core.Info qualified as Info
+
+newtype VolatilityInfo = VolatilityInfo
+  { _infoIsVolatile :: Bool
+  }
+
+instance IsInfo VolatilityInfo
+
+kVolatilityInfo :: Key VolatilityInfo
+kVolatilityInfo = Proxy
+
+makeLenses ''VolatilityInfo
+
+isVolatile :: Info -> Bool
+isVolatile i =
+  case Info.lookup kVolatilityInfo i of
+    Just VolatilityInfo {..} -> _infoIsVolatile
+    Nothing -> False
diff --git a/src/Juvix/Compiler/Core/Transformation/Optimize/LetFolding.hs b/src/Juvix/Compiler/Core/Transformation/Optimize/LetFolding.hs
@@ -17,6 +17,7 @@ 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.Info.VolatilityInfo qualified as Info
 import Juvix.Compiler.Core.Transformation.Base
 
 convertNode :: (Module -> BinderList Binder -> Node -> Bool) -> Module -> Node -> Node
@@ -25,10 +26,11 @@ convertNode isFoldable md = rmapL go
     go :: ([BinderChange] -> Node -> Node) -> BinderList Binder -> Node -> Node
     go recur bl = \case
       NLet Let {..}
-        | ( isImmediate md (_letItem ^. letItemValue)
-              || Info.freeVarOccurrences 0 _letBody <= 1
-              || isFoldable md bl (_letItem ^. letItemValue)
-          )
+        | not (Info.isVolatile _letInfo)
+            && ( isImmediate md (_letItem ^. letItemValue)
+                   || Info.freeVarOccurrences 0 _letBody <= 1
+                   || isFoldable md bl (_letItem ^. letItemValue)
+               )
             && not (Info.hasDebugOps _letBody) ->
             go (recur . (mkBCRemove b val' :)) (BL.cons b bl) _letBody
         where

diff --git a/src/Juvix/Compiler/Core/Transformation/Optimize/LoopHoisting.hs b/src/Juvix/Compiler/Core/Transformation/Optimize/LoopHoisting.hs
@@ -1,7 +1,9 @@
 module Juvix.Compiler.Core.Transformation.Optimize.LoopHoisting (loopHoisting) where
 
 import Juvix.Compiler.Core.Extra
+import Juvix.Compiler.Core.Info qualified as Info
 import Juvix.Compiler.Core.Info.FreeVarsInfo qualified as Info
+import Juvix.Compiler.Core.Info.VolatilityInfo qualified as Info
 import Juvix.Compiler.Core.Transformation.Base
 import Juvix.Compiler.Core.Transformation.ComputeTypeInfo (computeNodeType')
 
@@ -31,12 +33,13 @@ loopHoisting md = mapT (const (umapL go)) md
     goLamApp bl sym info h arg argNum args'
       | null subterms = goApp bl sym (mkApp info h arg) argNum args'
       | otherwise =
-          mkLets'
-            (map (\node -> (computeNodeType' md bl node, node)) subterms')
-            ( adjustLetBoundVars
-                . shift n
-                $ (mkApps (mkApp info h (reLambdasRev lams body')) args')
-            )
+          setLetsVolatile n $
+            mkLets'
+              (map (\node -> (computeNodeType' md bl node, node)) subterms')
+              ( adjustLetBoundVars
+                  . shift n
+                  $ (mkApps (mkApp info h (reLambdasRev lams body')) args')
+              )
       where
         (lams, body) = unfoldLambdasRev arg
         (subterms, body') = extractMaximalInvariantSubterms (length lams) body
@@ -73,3 +76,16 @@ loopHoisting md = mapT (const (umapL go)) md
           NVar Var {..}
             | _varIndex < 0 -> mkVar' (n - _varIndex - 1)
           _ -> node
+
+    setLetsVolatile :: Int -> Node -> Node
+    setLetsVolatile n
+      | n == 0 = id
+      | otherwise = \case
+          NLet Let {..} ->
+            NLet
+              Let
+                { _letInfo = Info.insert (Info.VolatilityInfo True) _letInfo,
+                  _letBody = setLetsVolatile (n - 1) _letBody,
+                  _letItem
+                }
+          node -> node