Compute temporary stack height in JuvixTree

src/Juvix/Compiler/Asm/Pipeline.hs

@@ -20,7 +20,7 @@ toReg' = validate >=> filterUnreachable >=> computeStackUsage >=> computePreallo
 -- | Perform transformations on JuvixAsm necessary before the translation to
 -- Nockma
 toNockma' :: (Members '[Error AsmError, Reader Options] r) => InfoTable -> Sem r InfoTable
-toNockma' = validate >=> filterUnreachable >=> computeTempHeight
+toNockma' = validate >=> filterUnreachable
 
 toReg :: (Members '[Error JuvixError, Reader EntryPoint] r) => InfoTable -> Sem r InfoTable
 toReg = mapReader fromEntryPoint . mapError (JuvixError @AsmError) . toReg'

src/Juvix/Compiler/Asm/Transformation.hs

@@ -3,12 +3,10 @@ module Juvix.Compiler.Asm.Transformation
     module Juvix.Compiler.Asm.Transformation.Prealloc,
     module Juvix.Compiler.Asm.Transformation.Validate,
     module Juvix.Compiler.Asm.Transformation.FilterUnreachable,
-    module Juvix.Compiler.Asm.Transformation.TempHeight,
   )
 where
 
 import Juvix.Compiler.Asm.Transformation.FilterUnreachable
 import Juvix.Compiler.Asm.Transformation.Prealloc
 import Juvix.Compiler.Asm.Transformation.StackUsage
-import Juvix.Compiler.Asm.Transformation.TempHeight
 import Juvix.Compiler.Asm.Transformation.Validate

src/Juvix/Compiler/Tree/Data/TransformationId.hs

@@ -9,6 +9,7 @@ data TransformationId
   | IdentityU
   | IdentityD
   | Apply
+  | TempHeight
   deriving stock (Data, Bounded, Enum, Show)
 
 data PipelineId
@@ -19,7 +20,7 @@ data PipelineId
 type TransformationLikeId = TransformationLikeId' TransformationId PipelineId
 
 toNockmaTransformations :: [TransformationId]
-toNockmaTransformations = [Apply]
+toNockmaTransformations = [Apply, TempHeight]
 
 toAsmTransformations :: [TransformationId]
 toAsmTransformations = []
@@ -31,6 +32,7 @@ instance TransformationId' TransformationId where
     IdentityU -> strIdentityU
     IdentityD -> strIdentityD
     Apply -> strApply
+    TempHeight -> strTempHeight
 
 instance PipelineId' TransformationId PipelineId where
   pipelineText :: PipelineId -> Text

src/Juvix/Compiler/Tree/Data/TransformationId/Strings.hs

@@ -19,3 +19,6 @@ strIdentityD = "identity-dmap"
 
 strApply :: Text
 strApply = "apply"
+
+strTempHeight :: Text
+strTempHeight = "temp-height"

src/Juvix/Compiler/Tree/Transformation.hs

@@ -9,6 +9,7 @@ import Juvix.Compiler.Tree.Data.TransformationId
 import Juvix.Compiler.Tree.Transformation.Apply
 import Juvix.Compiler.Tree.Transformation.Base
 import Juvix.Compiler.Tree.Transformation.Identity
+import Juvix.Compiler.Tree.Transformation.TempHeight
 
 applyTransformations :: forall r. [TransformationId] -> InfoTable -> Sem r InfoTable
 applyTransformations ts tbl = foldM (flip appTrans) tbl ts
@@ -19,3 +20,4 @@ applyTransformations ts tbl = foldM (flip appTrans) tbl ts
       IdentityU -> return . identityU
       IdentityD -> return . identityD
       Apply -> return . computeApply
+      TempHeight -> return . computeTempHeight

src/Juvix/Compiler/Tree/Transformation/TempHeight.hs

@@ -0,0 +1,25 @@
+module Juvix.Compiler.Tree.Transformation.TempHeight where
+
+import Juvix.Compiler.Tree.Extra.Recursors
+import Juvix.Compiler.Tree.Transformation.Base
+
+computeFunctionTempHeight :: Node -> Node
+computeFunctionTempHeight = umapN go
+  where
+    go :: Int -> Node -> Node
+    go k = \case
+      MemRef (DRef (TempRef r)) ->
+        let r' = set refTempTempHeight (Just k) r
+         in MemRef $ DRef (TempRef r')
+      MemRef (ConstrRef field@Field {_fieldRef = TempRef r}) ->
+        let r' = set refTempTempHeight (Just k) r
+         in MemRef
+              ( ConstrRef
+                  field
+                    { _fieldRef = TempRef r'
+                    }
+              )
+      node -> node
+
+computeTempHeight :: InfoTable -> InfoTable
+computeTempHeight = mapT (const computeFunctionTempHeight)