Precompute maximum heap allocation

.gitignore

@@ -56,7 +56,7 @@ TAGS
 # other
 .DS_Store
 
-/runtime/include
+/runtime/include/
 _build/
 _build.*/
 *.agdai

app/Commands/Dev/Asm/Run.hs

@@ -7,6 +7,7 @@ import Juvix.Compiler.Asm.Error qualified as Asm
 import Juvix.Compiler.Asm.Extra qualified as Asm
 import Juvix.Compiler.Asm.Interpreter qualified as Asm
 import Juvix.Compiler.Asm.Pretty qualified as Asm
+import Juvix.Compiler.Asm.Transformation.Validate qualified as Asm
 import Juvix.Compiler.Asm.Translation.FromSource qualified as Asm
 
 runCommand :: forall r. Members '[Embed IO, App] r => AsmRunOptions -> Sem r ()
@@ -15,7 +16,7 @@ runCommand opts = do
   case Asm.runParser file s of
     Left err -> exitJuvixError (JuvixError err)
     Right tab ->
-      let v = if opts ^. asmRunNoValidate then Nothing else Asm.validate tab
+      let v = if opts ^. asmRunNoValidate then Nothing else Asm.validate' tab
        in case v of
             Just err ->
               exitJuvixError (JuvixError err)

app/Commands/Dev/Asm/Validate.hs

@@ -2,7 +2,7 @@ module Commands.Dev.Asm.Validate where
 
 import Commands.Base
 import Commands.Dev.Asm.Validate.Options
-import Juvix.Compiler.Asm.Extra qualified as Asm
+import Juvix.Compiler.Asm.Transformation.Validate qualified as Asm
 import Juvix.Compiler.Asm.Translation.FromSource qualified as Asm
 
 runCommand :: forall r. Members '[Embed IO, App] r => AsmValidateOptions -> Sem r ()
@@ -11,7 +11,7 @@ runCommand opts = do
   case Asm.runParser file s of
     Left err -> exitJuvixError (JuvixError err)
     Right tab -> do
-      case Asm.validate tab of
+      case Asm.validate' tab of
         Just err ->
           exitJuvixError (JuvixError err)
         Nothing ->

runtime/src/juvix/limits.h

@@ -30,17 +30,13 @@
 #define MAX_FIELDS 255U
 
 #define MAX_CONSTR_ARGS MAX_FIELDS
-// Max number of fields minus the extra field in a closure.
-#define MAX_FUNCTION_ARGS (MAX_FIELDS - 1)
+// Max number of fields minus the extra fields in a closure.
+#define MAX_FUNCTION_ARGS (MAX_FIELDS - 2)
 
 #define MAX_CSTRING_LENGTH (MAX_FIELDS * sizeof(word_t) - 1)
 
 #define MAX_LOCAL_VARS (PAGE_SIZE / 2 / sizeof(word_t))
 
-// The maximum number of words that can be allocated on the heap by an
-// invocation of the dispatch loop (CallClosures)
-#define MAX_DISPATCH_ALLOC (MAX_FUNCTION_ARGS + CLOSURE_SKIP + 1)
-
 /*****************************************/
 /* Static asserts */
 

src/Juvix/Compiler/Asm/Extra.hs

@@ -1,78 +1,12 @@
 module Juvix.Compiler.Asm.Extra
-  ( module Juvix.Compiler.Asm.Extra,
-    module Juvix.Compiler.Asm.Extra.Base,
+  ( module Juvix.Compiler.Asm.Extra.Base,
     module Juvix.Compiler.Asm.Extra.Type,
     module Juvix.Compiler.Asm.Extra.Recursors,
     module Juvix.Compiler.Asm.Error,
   )
 where
 
-import Data.HashMap.Strict qualified as HashMap
-import Juvix.Compiler.Asm.Data.InfoTable
 import Juvix.Compiler.Asm.Error
 import Juvix.Compiler.Asm.Extra.Base
 import Juvix.Compiler.Asm.Extra.Recursors
 import Juvix.Compiler.Asm.Extra.Type
-import Juvix.Compiler.Asm.Language
-
-validateCode :: forall r. Member (Error AsmError) r => InfoTable -> Arguments -> Code -> Sem r ()
-validateCode tab args = void . recurse sig args
-  where
-    sig :: RecursorSig Memory r ()
-    sig =
-      RecursorSig
-        { _recursorInfoTable = tab,
-          _recurseInstr = \_ _ -> return (),
-          _recurseBranch = \_ _ _ _ -> return (),
-          _recurseCase = \_ _ _ _ -> return ()
-        }
-
-validateFunction :: Member (Error AsmError) r => InfoTable -> FunctionInfo -> Sem r ()
-validateFunction tab fi = validateCode tab (argumentsFromFunctionInfo fi) (fi ^. functionCode)
-
-validateInfoTable :: Member (Error AsmError) r => InfoTable -> Sem r InfoTable
-validateInfoTable tab = do
-  mapM_ (validateFunction tab) (HashMap.elems (tab ^. infoFunctions))
-  return tab
-
-validate :: InfoTable -> Maybe AsmError
-validate tab =
-  case run $ runError $ validateInfoTable tab of
-    Left err -> Just err
-    _ -> Nothing
-
-computeFunctionStackUsage :: Member (Error AsmError) r => InfoTable -> FunctionInfo -> Sem r FunctionInfo
-computeFunctionStackUsage tab fi = do
-  ps <- snd <$> recurseS sig initialStackInfo (fi ^. functionCode)
-  let maxValueStack = maximum (map fst ps)
-      maxTempStack = maximum (map snd ps)
-  return
-    fi
-      { _functionMaxValueStackHeight = maxValueStack,
-        _functionMaxTempStackHeight = maxTempStack
-      }
-  where
-    sig :: RecursorSig StackInfo r (Int, Int)
-    sig =
-      RecursorSig
-        { _recursorInfoTable = tab,
-          _recurseInstr = \si _ -> return (si ^. stackInfoValueStackHeight, si ^. stackInfoTempStackHeight),
-          _recurseBranch = \si _ l r ->
-            return
-              ( max (si ^. stackInfoValueStackHeight) (max (maximum (map fst l)) (maximum (map fst r))),
-                max (si ^. stackInfoTempStackHeight) (max (maximum (map snd l)) (maximum (map snd r)))
-              ),
-          _recurseCase = \si _ cs md ->
-            return
-              ( max (si ^. stackInfoValueStackHeight) (max (maximum (map (maximum . map fst) cs)) (maybe 0 (maximum . map fst) md)),
-                max (si ^. stackInfoTempStackHeight) (max (maximum (map (maximum . map snd) cs)) (maybe 0 (maximum . map snd) md))
-              )
-        }
-
-computeStackUsage :: Member (Error AsmError) r => InfoTable -> Sem r InfoTable
-computeStackUsage tab = do
-  fns <- mapM (computeFunctionStackUsage tab) (tab ^. infoFunctions)
-  return tab {_infoFunctions = fns}
-
-computeStackUsage' :: InfoTable -> Either AsmError InfoTable
-computeStackUsage' tab = run $ runError $ computeStackUsage tab

src/Juvix/Compiler/Asm/Extra/Base.hs

@@ -32,3 +32,11 @@ isFinalInstr = \case
   TailCallClosures {} -> True
   Failure -> True
   _ -> False
+
+getConstrSize :: MemRep -> Int -> Int
+getConstrSize rep argsNum = case rep of
+  MemRepConstr -> 1 + argsNum
+  MemRepTag -> 0
+  MemRepTuple -> argsNum
+  MemRepUnit -> 0
+  MemRepUnpacked {} -> 0

src/Juvix/Compiler/Asm/Extra/Recursors.hs

@@ -113,6 +113,8 @@ recurse' sig = go True
               return mem
             Failure ->
               return mem
+            Prealloc {} ->
+              return mem
             AllocConstr tag -> do
               let ci = getConstrInfo (sig ^. recursorInfoTable) tag
                   n = ci ^. constructorArgsNum
@@ -294,8 +296,11 @@ initialStackInfo = StackInfo {_stackInfoValueStackHeight = 0, _stackInfoTempStac
 -- program code which need only stack height information. Also, the code using
 -- the simplified recursor can itself be simpler if it doesn't need the extra
 -- info provided by the full recursor.
-recurseS :: forall r a. Member (Error AsmError) r => RecursorSig StackInfo r a -> StackInfo -> Code -> Sem r (StackInfo, [a])
-recurseS sig = go
+recurseS :: forall r a. Member (Error AsmError) r => RecursorSig StackInfo r a -> Code -> Sem r [a]
+recurseS sig code = snd <$> recurseS' sig initialStackInfo code
+
+recurseS' :: forall r a. Member (Error AsmError) r => RecursorSig StackInfo r a -> StackInfo -> Code -> Sem r (StackInfo, [a])
+recurseS' sig = go
   where
     go :: StackInfo -> Code -> Sem r (StackInfo, [a])
     go si = \case
@@ -353,6 +358,8 @@ recurseS sig = go
               return si
             Failure ->
               return si
+            Prealloc {} ->
+              return si
             AllocConstr tag -> do
               let ci = getConstrInfo (sig ^. recursorInfoTable) tag
                   n = ci ^. constructorArgsNum
@@ -363,17 +370,17 @@ recurseS sig = go
                 stackInfoPopValueStack (_allocClosureArgsNum - 1) si
             ExtendClosure InstrExtendClosure {..} ->
               return $
-                stackInfoPopValueStack (_extendClosureArgsNum - 1) si
+                stackInfoPopValueStack _extendClosureArgsNum si
             Call x ->
               fixStackCall si x
             TailCall x ->
-              fixStackCall si x
+              fixStackCall (dropTempStack si) x
             CallClosures x ->
               fixStackCallClosures si x
             TailCallClosures x ->
-              fixStackCallClosures si x
+              fixStackCallClosures (dropTempStack si) x
             Return ->
-              return si
+              return (dropTempStack si)
 
         fixStackBinOp :: StackInfo -> Sem r StackInfo
         fixStackBinOp si = return $ stackInfoPopValueStack 1 si
@@ -432,3 +439,54 @@ recurseS sig = go
 
     stackInfoPopTempStack :: Int -> StackInfo -> StackInfo
     stackInfoPopTempStack n si = si {_stackInfoTempStackHeight = si ^. stackInfoTempStackHeight - n}
+
+    dropTempStack :: StackInfo -> StackInfo
+    dropTempStack si = si {_stackInfoTempStackHeight = 0}
+
+-- | Fold signature. Contains read-only fold parameters.
+data FoldSig m r a = FoldSig
+  { _foldInfoTable :: InfoTable,
+    _foldAdjust :: a -> a,
+    _foldInstr :: m -> CmdInstr -> a -> Sem r a,
+    _foldBranch :: m -> CmdBranch -> a -> a -> a -> Sem r a,
+    _foldCase :: m -> CmdCase -> [a] -> Maybe a -> a -> Sem r a
+  }
+
+makeLenses ''FoldSig
+
+foldS :: forall r a. Member (Error AsmError) r => FoldSig StackInfo r a -> Code -> a -> Sem r a
+foldS sig code a = snd <$> foldS' sig initialStackInfo code a
+
+foldS' :: forall r a. Member (Error AsmError) r => FoldSig StackInfo r a -> StackInfo -> Code -> a -> Sem r (StackInfo, a)
+foldS' sig si code acc = do
+  (si', fs) <- recurseS' sig' si code
+  a' <- compose fs acc
+  return (si', a')
+  where
+    sig' :: RecursorSig StackInfo r (a -> Sem r a)
+    sig' =
+      RecursorSig
+        { _recursorInfoTable = sig ^. foldInfoTable,
+          _recurseInstr = \s cmd -> return ((sig ^. foldInstr) s cmd),
+          _recurseBranch = \s cmd br1 br2 ->
+            return
+              ( \a -> do
+                  let a' = (sig ^. foldAdjust) a
+                  a1 <- compose br1 a'
+                  a2 <- compose br2 a'
+                  (sig ^. foldBranch) s cmd a1 a2 a
+              ),
+          _recurseCase = \s cmd brs md ->
+            return
+              ( \a -> do
+                  let a' = (sig ^. foldAdjust) a
+                  as <- mapM (`compose` a') brs
+                  ad <- case md of
+                    Just d -> Just <$> compose d a'
+                    Nothing -> return Nothing
+                  (sig ^. foldCase) s cmd as ad a
+              )
+        }
+
+    compose :: [a -> Sem r a] -> a -> Sem r a
+    compose lst x = foldr (=<<) (return x) lst

src/Juvix/Compiler/Asm/Interpreter.hs

@@ -113,6 +113,8 @@ runCodeR infoTable funInfo = goCode (funInfo ^. functionCode) >> popLastValueSta
       Failure -> do
         v <- topValueStack
         runtimeError $ mappend "failure: " (printVal v)
+      Prealloc {} ->
+        goCode cont
       AllocConstr tag -> do
         let ci = getConstrInfo infoTable tag
         args <- replicateM (ci ^. constructorArgsNum) popValueStack

src/Juvix/Compiler/Asm/Language.hs

@@ -35,7 +35,7 @@ data Value
 data MemValue
   = -- | A direct memory reference.
     DRef DirectRef
-  | -- | ConstrRef references is an indirect reference to a field (argument) of
+  | -- | ConstrRef is an indirect reference to a field (argument) of
     --  a constructor: field k holds the (k+1)th argument.
     ConstrRef Field
 
@@ -87,6 +87,9 @@ data Instruction
   | -- | Interrupt execution with a runtime error printing the value on top of
     -- the stack. JVA opcode: 'fail'.
     Failure
+  | -- | Preallocate memory. This instruction is inserted automatically before
+    -- translation to JuvixReg. It does not occur in JVA files.
+    Prealloc InstrPrealloc
   | -- | Allocate constructor data with a given tag. The n arguments (the number n
     -- determined by the constant tag) are popped from the stack and stored at
     -- increasing offsets (stack[0]: field 0, stack[1]: field 1, ...,
@@ -168,6 +171,11 @@ data Opcode
     -- by two, and push the result. JVA opcode: 'eq'.
     ValEq
 
+newtype InstrPrealloc = InstrPrealloc
+  { -- | How many words to preallocate?
+    _preallocWordsNum :: Int
+  }
+
 data InstrAllocClosure = InstrAllocClosure
   { _allocClosureFunSymbol :: Symbol,
     -- | The number of supplied arguments to be stored in the closure.

src/Juvix/Compiler/Asm/Options.hs

@@ -0,0 +1,18 @@
+module Juvix.Compiler.Asm.Options
+  ( module Juvix.Compiler.Asm.Options,
+    module Juvix.Compiler.Backend,
+  )
+where
+
+import Juvix.Compiler.Backend
+import Juvix.Prelude
+
+data Options = Options
+  { _optDebug :: Bool,
+    _optLimits :: Limits
+  }
+
+makeLenses ''Options
+
+getClosureSize :: Options -> Int -> Int
+getClosureSize opts argsNum = opts ^. optLimits . limitsClosureHeadSize + argsNum

src/Juvix/Compiler/Asm/Pipeline.hs

@@ -3,8 +3,10 @@ module Juvix.Compiler.Asm.Pipeline where
 import Juvix.Compiler.Asm.Data.InfoTable
 import Juvix.Compiler.Asm.Extra
 import Juvix.Compiler.Asm.Language
+import Juvix.Compiler.Asm.Options
+import Juvix.Compiler.Asm.Transformation
 
 -- | Perform transformations on JuvixAsm necessary before the translation to
 -- JuvixReg
-toReg :: Member (Error AsmError) r => InfoTable -> Sem r InfoTable
-toReg = validateInfoTable >=> computeStackUsage
+toReg :: Members '[Error AsmError, Reader Options] r => InfoTable -> Sem r InfoTable
+toReg = validate >=> computeStackUsage >=> computePrealloc

src/Juvix/Compiler/Asm/Pretty/Base.hs

@@ -194,6 +194,7 @@ instance PrettyCode Instruction where
     Trace -> return $ pretty ("trace" :: String)
     Dump -> return $ pretty ("dump" :: String)
     Failure -> return $ pretty ("fail" :: String)
+    Prealloc {} -> return $ pretty ("prealloc" :: String)
     AllocConstr {} -> return $ pretty ("alloc" :: String)
     AllocClosure {} -> return $ pretty ("calloc" :: String)
     ExtendClosure {} -> return $ pretty ("cextend" :: String)

src/Juvix/Compiler/Asm/Transformation.hs

@@ -0,0 +1,10 @@
+module Juvix.Compiler.Asm.Transformation
+  ( module Juvix.Compiler.Asm.Transformation.StackUsage,
+    module Juvix.Compiler.Asm.Transformation.Prealloc,
+    module Juvix.Compiler.Asm.Transformation.Validate,
+  )
+where
+
+import Juvix.Compiler.Asm.Transformation.Prealloc
+import Juvix.Compiler.Asm.Transformation.StackUsage
+import Juvix.Compiler.Asm.Transformation.Validate

src/Juvix/Compiler/Asm/Transformation/Base.hs

@@ -0,0 +1,27 @@
+module Juvix.Compiler.Asm.Transformation.Base
+  ( module Juvix.Compiler.Asm.Transformation.Base,
+    module Juvix.Compiler.Asm.Data.InfoTable,
+    module Juvix.Compiler.Asm.Extra,
+    module Juvix.Compiler.Asm.Language,
+  )
+where
+
+import Juvix.Compiler.Asm.Data.InfoTable
+import Juvix.Compiler.Asm.Extra
+import Juvix.Compiler.Asm.Language
+
+liftCodeTransformation :: Monad m => (Code -> m Code) -> FunctionInfo -> m FunctionInfo
+liftCodeTransformation f fi = do
+  code <- f (fi ^. functionCode)
+  return fi {_functionCode = code}
+
+liftFunctionTransformation :: Monad m => (FunctionInfo -> m FunctionInfo) -> InfoTable -> m InfoTable
+liftFunctionTransformation f tab = do
+  fns <- mapM f (tab ^. infoFunctions)
+  return tab {_infoFunctions = fns}
+
+runTransformation :: (InfoTable -> Sem '[Error AsmError] InfoTable) -> InfoTable -> Either AsmError InfoTable
+runTransformation trans tab =
+  case run $ runError $ trans tab of
+    Left err -> Left err
+    Right tab' -> Right tab'