Implement the dynamic dispatch loop in JuvixAsm (#2556)

* Closes #2555 
* Depends on #2554

app/Commands/Dev/Asm/Compile.hs

@@ -30,7 +30,12 @@ runCommand opts = do
           ensureDir buildDir
           cFile <- inputCFile file
           embed $ TIO.writeFile (toFilePath cFile) _resultCCode
-          Compile.runCommand opts {_compileInputFile = Just (AppPath (preFileFromAbs cFile) False)}
+          outfile <- Compile.outputFile opts file
+          Compile.runCommand
+            opts
+              { _compileInputFile = Just (AppPath (preFileFromAbs cFile) False),
+                _compileOutputFile = Just (AppPath (preFileFromAbs outfile) False)
+              }
   where
     getFile :: Sem r (Path Abs File)
     getFile = getMainFile (opts ^. compileInputFile)

cntlines.sh

@@ -4,14 +4,15 @@ function count() {
     cloc $1 | grep 'SUM:' | awk '{print $5}'
 }
 
-function count_pir () {
-    find $1 -name '*.pir' -print | xargs sed '/^[[:space:]]*$/d' | wc -l | tr -d ' '
+function count_ext () {
+    find $2 -name $1 -print | xargs sed '/^[[:space:]]*$/d' | wc -l | tr -d ' '
 }
 
 RUNTIME_C=$(count runtime/src/juvix)
-RUNTIME_VAMPIR=$(count_pir runtime/src/vampir)
+RUNTIME_VAMPIR=$(count_ext '*.pir' runtime/src/vampir)
+RUNTIME_JVA=$(count_ext '*.jva' runtime/src/asm)
 
-RUNTIME=$((RUNTIME_C+RUNTIME_VAMPIR))
+RUNTIME=$((RUNTIME_C+RUNTIME_VAMPIR+RUNTIME_JVA))
 
 BACKENDC=$(count src/Juvix/Compiler/Backend/C/)
 GEB=$(count src/Juvix/Compiler/Backend/Geb/)
@@ -52,6 +53,7 @@ echo "   JuvixAsm: $ASM LOC"
 echo "   JuvixCore: $CORE LOC"
 echo "Runtime: $RUNTIME LOC"
 echo "   C runtime: $RUNTIME_C LOC"
+echo "   JuvixAsm runtime: $RUNTIME_JVA LOC"
 echo "   VampIR runtime: $RUNTIME_VAMPIR LOC"
 echo "Other: $OTHER LOC"
 echo "   Application: $APP LOC"
@@ -61,4 +63,4 @@ echo "   Data: $DATA LOC"
 echo "   Prelude: $PRELUDE LOC"
 echo "Tests: $TESTS LOC"
 echo ""
-echo "Total: $TOTAL Haskell LOC + $RUNTIME_C C LOC + $RUNTIME_VAMPIR VampIR LOC"
+echo "Total: $TOTAL Haskell LOC + $RUNTIME_C C LOC + $RUNTIME_JVA JuvixAsm LOC + $RUNTIME_VAMPIR VampIR LOC"

runtime/src/asm/apply.jva

@@ -0,0 +1,193 @@
+
+function juvix_apply_1(*, *) : * {
+    push arg[0];
+    argsnum;
+    push 1;
+    eq;
+    br {
+        true: { -- argsnum = 1
+            push arg[1];
+            push arg[0];
+            tcall $ 1;
+        }
+        false: { -- argsnum > 1
+            push arg[1];
+            push arg[0];
+            cextend 1;
+            ret;
+        }
+    };
+}
+
+function juvix_apply_2(*, *, *) : * {
+    push arg[0];
+    argsnum;
+    tsave n {
+        push n;
+        push 2;
+        eq;
+        br {
+            true: { -- argsnum = 2
+                push arg[2];
+                push arg[1];
+                push arg[0];
+                tcall $ 2;
+            }
+            false: {
+                push n;
+                push 1;
+                eq;
+                br {
+                    true: { -- argsnum = 1
+                        push arg[2];
+                        push arg[1];
+                        push arg[0];
+                        call $ 1;
+                        tcall juvix_apply_1;
+                    }
+                    false: { -- argsnum > 2
+                        push arg[2];
+                        push arg[1];
+                        push arg[0];
+                        cextend 2;
+                        ret;
+                    }
+                };
+            }
+        };
+    };
+}
+
+function juvix_apply_3(*, *, *, *) : * {
+    push arg[0];
+    argsnum;
+    tsave n {
+        push n;
+        push 3;
+        eq;
+        br {
+            true: { -- argsnum = 3
+                push arg[3];
+                push arg[2];
+                push arg[1];
+                push arg[0];
+                tcall $ 3;
+            }
+            false: {
+                push n;
+                push 3;
+                lt;
+                br {
+                    true: { -- argsnum > 3
+                        push arg[3];
+                        push arg[2];
+                        push arg[1];
+                        push arg[0];
+                        cextend 3;
+                        ret;
+                    }
+                    false: { -- argsnum <= 2
+                        push n;
+                        push 2;
+                        eq;
+                        br {
+                            true: { -- argsnum = 2
+                                push arg[3];
+                                push arg[2];
+                                push arg[1];
+                                push arg[0];
+                                call $ 2;
+                                tcall juvix_apply_1;
+                            }
+                            false: { -- argsnum = 1
+                                push arg[3];
+                                push arg[2];
+                                push arg[1];
+                                push arg[0];
+                                call $ 1;
+                                tcall juvix_apply_2;
+                            }
+                        };
+                    }
+                };
+            }
+        };
+    };
+}
+
+function juvix_apply_4(*, *, *, *, *) : * {
+    push arg[0];
+    argsnum;
+    tsave n {
+        push n;
+        push 4;
+        eq;
+        br {
+            true: { -- argsnum = 4
+                push arg[4];
+                push arg[3];
+                push arg[2];
+                push arg[1];
+                push arg[0];
+                tcall $ 4;
+            }
+            false: {
+                push n;
+                push 4;
+                lt;
+                br {
+                    true: { -- argsnum > 4
+                        push arg[4];
+                        push arg[3];
+                        push arg[2];
+                        push arg[1];
+                        push arg[0];
+                        cextend 4;
+                        ret;
+                    }
+                    false: { -- argsnum <= 3
+                        push n;
+                        push 3;
+                        eq;
+                        br {
+                            true: { -- argsnum = 3
+                                push arg[4];
+                                push arg[3];
+                                push arg[2];
+                                push arg[1];
+                                push arg[0];
+                                call $ 3;
+                                tcall juvix_apply_1;
+                            }
+                            false: {
+                                push n;
+                                push 2;
+                                eq;
+                                br {
+                                    true: { -- argsnum = 2
+                                        push arg[4];
+                                        push arg[3];
+                                        push arg[2];
+                                        push arg[1];
+                                        push arg[0];
+                                        call $ 2;
+                                        tcall juvix_apply_2;
+                                    }
+                                    false: { -- argsnum = 1
+                                        push arg[4];
+                                        push arg[3];
+                                        push arg[2];
+                                        push arg[1];
+                                        push arg[0];
+                                        call $ 1;
+                                        tcall juvix_apply_3;
+                                    }
+                                };
+                            }
+                        };
+                    }
+                };
+            }
+        };
+    };
+}

runtime/src/juvix/api.h

@@ -112,6 +112,8 @@ closure_label:
         error_exit();      \
     } while (0)
 
+#define JUVIX_ARGS_NUM(var, val) (var = make_smallint(get_closure_largs(val)))
+
 #define JUVIX_ALLOC_INT(var, val) (var = make_smallint(val))
 // ALLOC_CONSTR_BOXED(var, uid, nargs)
 // ALLOC_CONSTR_BOXED_TAG(var, uid)

src/Juvix/Compiler/Asm/Data/InfoTableBuilder.hs

@@ -42,9 +42,11 @@ emptyBuilderState =
     }
 
 runInfoTableBuilder :: Sem (InfoTableBuilder ': r) a -> Sem r (InfoTable, a)
-runInfoTableBuilder =
-  fmap (first (^. stateInfoTable))
-    . runState emptyBuilderState
+runInfoTableBuilder = fmap (first (^. stateInfoTable)) . runInfoTableBuilder' emptyBuilderState
+
+runInfoTableBuilder' :: BuilderState -> Sem (InfoTableBuilder ': r) a -> Sem r (BuilderState, a)
+runInfoTableBuilder' bs =
+  runState bs
     . reinterpret interp
   where
     interp :: InfoTableBuilder m a -> Sem (State BuilderState ': r) a

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

@@ -0,0 +1,60 @@
+module Juvix.Compiler.Asm.Extra.Apply where
+
+import Data.FileEmbed qualified as FE
+import Data.HashMap.Strict qualified as HashMap
+import Data.Text.Encoding
+import Juvix.Compiler.Asm.Data.InfoTable
+import Juvix.Compiler.Asm.Data.InfoTableBuilder
+import Juvix.Compiler.Asm.Language
+import Juvix.Compiler.Asm.Translation.FromSource
+
+data ApplyBuiltins = ApplyBuiltins
+  { -- | The number of `juvix_apply_n` functions.
+    _applyBuiltinsNum :: Int,
+    -- | Maps `n` to the function `juvix_apply_n`.
+    _applyBuiltinsMap :: HashMap Int Symbol
+  }
+
+makeLenses ''ApplyBuiltins
+
+addApplyBuiltins :: InfoTable -> (ApplyBuiltins, InfoTable)
+addApplyBuiltins tab = (blts, bs' ^. stateInfoTable)
+  where
+    nextSymbol = maximum (0 : HashMap.keys (tab ^. infoFunctions) ++ HashMap.keys (tab ^. infoInductives)) + 1
+    nextUserId = maximum (0 : mapMaybe getUserTag (HashMap.keys (tab ^. infoConstrs))) + 1
+
+    bs :: BuilderState
+    bs =
+      BuilderState
+        { _stateNextSymbol = nextSymbol,
+          _stateNextUserTag = nextUserId,
+          _stateInfoTable = tab,
+          _stateIdents = mempty
+        }
+
+    bs' :: BuilderState
+    bs' =
+      fromRight impossible $
+        parseText' bs $
+          decodeUtf8 $(FE.makeRelativeToProject "runtime/src/asm/apply.jva" >>= FE.embedFile)
+
+    blts :: ApplyBuiltins
+    blts =
+      ApplyBuiltins
+        { _applyBuiltinsNum = 4,
+          _applyBuiltinsMap =
+            HashMap.fromList $ map mkApply [1 .. 4]
+        }
+
+    mkApply :: Int -> (Int, Symbol)
+    mkApply x = (x, f)
+      where
+        idt = "juvix_apply_" <> show x
+        f = case fromJust $ HashMap.lookup idt (bs' ^. stateIdents) of
+          IdentFun s -> s
+          _ -> impossible
+
+    getUserTag :: Tag -> Maybe Word
+    getUserTag = \case
+      BuiltinTag {} -> Nothing
+      UserTag x -> Just x

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

@@ -113,6 +113,12 @@ recurse' sig = go True
               return mem
             Failure ->
               return $ pushValueStack TyDynamic (popValueStack 1 mem)
+            ArgsNum -> do
+              when (null (mem ^. memoryValueStack)) $
+                throw $
+                  AsmError loc "empty value stack"
+              checkFunType (topValueStack' 0 mem)
+              return $ pushValueStack mkTypeInteger (popValueStack 1 mem)
             Prealloc {} ->
               return mem
             AllocConstr tag -> do
@@ -384,6 +390,9 @@ recurseS' sig = go
               return si
             Failure ->
               return si
+            ArgsNum ->
+              -- push + pop = nop
+              return si
             Prealloc {} ->
               return si
             AllocConstr tag -> do

src/Juvix/Compiler/Asm/Interpreter.hs

@@ -128,6 +128,14 @@ runCodeR infoTable funInfo = goCode (funInfo ^. functionCode) >> popLastValueSta
       Failure -> do
         v <- topValueStack
         runtimeError $ mappend "failure: " (printVal v)
+      ArgsNum -> do
+        v <- popValueStack
+        case v of
+          ValClosure cl -> do
+            let n = lookupFunInfo infoTable (cl ^. closureSymbol) ^. functionArgsNum - length (cl ^. closureArgs)
+            pushValueStack (ValInteger (toInteger n))
+            goCode cont
+          _ -> runtimeError "invalid operation: expected closure on top of value stack"
       Prealloc {} ->
         goCode cont
       AllocConstr tag -> do

src/Juvix/Compiler/Asm/Language.hs

@@ -95,6 +95,10 @@ data Instruction
   | -- | Interrupt execution with a runtime error printing the value on top of
     -- the stack. JVA opcode: 'fail'.
     Failure
+  | -- | Computes the number of expected arguments for the closure on top of the
+    -- stack, pops the stack and pushes the result on top of the stack. JVA
+    -- opcode: 'argsnum'.
+    ArgsNum
   | -- | Preallocate memory. This instruction is inserted automatically before
     -- translation to JuvixReg. It does not occur in JVA files.
     Prealloc InstrPrealloc

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

@@ -280,6 +280,7 @@ instance PrettyCode Instruction where
     Trace -> return $ primitive Str.instrTrace
     Dump -> return $ primitive Str.instrDump
     Failure -> return $ primitive Str.instrFailure
+    ArgsNum -> return $ primitive Str.instrArgsNum
     Prealloc InstrPrealloc {..} ->
       return $ primitive Str.instrPrealloc <+> integer _preallocWordsNum
     AllocConstr tag -> (primitive Str.instrAlloc <+>) <$> ppConstrName tag

src/Juvix/Compiler/Asm/Transformation.hs

@@ -2,9 +2,11 @@ module Juvix.Compiler.Asm.Transformation
   ( module Juvix.Compiler.Asm.Transformation.StackUsage,
     module Juvix.Compiler.Asm.Transformation.Prealloc,
     module Juvix.Compiler.Asm.Transformation.Validate,
+    module Juvix.Compiler.Asm.Transformation.Apply,
   )
 where
 
+import Juvix.Compiler.Asm.Transformation.Apply
 import Juvix.Compiler.Asm.Transformation.Prealloc
 import Juvix.Compiler.Asm.Transformation.StackUsage
 import Juvix.Compiler.Asm.Transformation.Validate

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

@@ -0,0 +1,104 @@
+module Juvix.Compiler.Asm.Transformation.Apply where
+
+import Data.HashMap.Strict qualified as HashMap
+import Juvix.Compiler.Asm.Extra.Apply
+import Juvix.Compiler.Asm.Options
+import Juvix.Compiler.Asm.Transformation.Base
+
+computeFunctionApply :: (Member (Error AsmError) r) => ApplyBuiltins -> InfoTable -> FunctionInfo -> Sem r FunctionInfo
+computeFunctionApply blts tab fi = do
+  cs <- recurseS sig (fi ^. functionCode)
+  return fi {_functionCode = concat cs}
+  where
+    sig :: RecursorSig StackInfo r Code
+    sig =
+      RecursorSig
+        { _recursorInfoTable = tab,
+          _recurseInstr = \_ cmd -> goInstr cmd,
+          _recurseBranch = \_ cmd l r -> goBranch cmd l r,
+          _recurseCase = \_ cmd cs md -> goCase cmd cs md,
+          _recurseSave = \_ cmd b -> goSave cmd b
+        }
+
+    goInstr :: CmdInstr -> Sem r Code
+    goInstr cmd = case cmd ^. cmdInstrInstruction of
+      CallClosures InstrCallClosures {..} -> return $ goApply False _callClosuresArgsNum
+      TailCallClosures InstrCallClosures {..} -> return $ goApply True _callClosuresArgsNum
+      _ -> return [Instr cmd]
+
+    goApply :: Bool -> Int -> Code
+    goApply isTail n = replicate m (mkApply False (blts ^. applyBuiltinsNum)) ++ [mkApply isTail r]
+      where
+        (m, r) = n `divMod` (blts ^. applyBuiltinsNum)
+
+    mkApply :: Bool -> Int -> Command
+    mkApply isTail k =
+      Instr $
+        CmdInstr emptyInfo $
+          (if isTail then TailCall else Call)
+            InstrCall
+              { _callType = CallFun sym,
+                _callArgsNum = k + 1
+              }
+      where
+        sym = fromJust $ HashMap.lookup k (blts ^. applyBuiltinsMap)
+
+    goBranch :: CmdBranch -> [Code] -> [Code] -> Sem r Code
+    goBranch cmd l r =
+      return
+        [ Branch
+            cmd
+              { _cmdBranchTrue = concat l,
+                _cmdBranchFalse = concat r
+              }
+        ]
+
+    goCase :: CmdCase -> [[Code]] -> Maybe [Code] -> Sem r Code
+    goCase cmd cs md =
+      return
+        [ Case
+            cmd
+              { _cmdCaseBranches =
+                  zipWith
+                    (\br c -> CaseBranch (br ^. caseBranchTag) (concat c))
+                    (cmd ^. cmdCaseBranches)
+                    cs,
+                _cmdCaseDefault = fmap concat md
+              }
+        ]
+
+    goSave :: CmdSave -> [Code] -> Sem r Code
+    goSave cmd c = return [Save cmd {_cmdSaveCode = concat c}]
+
+computeApply :: (Member (Error AsmError) r) => InfoTable -> Sem r InfoTable
+computeApply tab = liftFunctionTransformation (computeFunctionApply blts tab') tab'
+  where
+    (blts, tab') = addApplyBuiltins tab
+
+checkNoCallClosures :: Options -> InfoTable -> Bool
+checkNoCallClosures opts tab =
+  case run $ runError $ runReader opts sb of
+    Left err -> error (show err)
+    Right b -> b
+  where
+    sb :: Sem '[Reader Options, Error AsmError] Bool
+    sb = allM (check . (^. functionCode)) (HashMap.elems (tab ^. infoFunctions))
+
+    check :: Code -> Sem '[Reader Options, Error AsmError] Bool
+    check c = foldS sig c True
+      where
+        sig =
+          FoldSig
+            { _foldInfoTable = tab,
+              _foldAdjust = id,
+              _foldInstr = \_ cmd b -> return $ b && goInstr (cmd ^. cmdInstrInstruction),
+              _foldBranch = \_ _ b1 b2 b3 -> return $ b1 && b2 && b3,
+              _foldCase = \_ _ bs bd b -> return $ and bs && fromMaybe True bd && b,
+              _foldSave = \_ _ b1 b2 -> return $ b1 && b2
+            }
+
+        goInstr :: Instruction -> Bool
+        goInstr = \case
+          CallClosures {} -> False
+          TailCallClosures {} -> False
+          _ -> True

src/Juvix/Compiler/Asm/Translation/FromSource.hs

@@ -1,6 +1,7 @@
 module Juvix.Compiler.Asm.Translation.FromSource
   ( module Juvix.Compiler.Asm.Translation.FromSource,
     module Juvix.Parser.Error,
+    BuilderState,
   )
 where
 
@@ -29,16 +30,22 @@ localS update a = do
 parseText :: Text -> Either MegaparsecError InfoTable
 parseText = runParser ""
 
+parseText' :: BuilderState -> Text -> Either MegaparsecError BuilderState
+parseText' bs = runParser' bs ""
+
 runParser :: FilePath -> Text -> Either MegaparsecError InfoTable
-runParser fileName input =
+runParser fileName input = (^. stateInfoTable) <$> runParser' emptyBuilderState fileName input
+
+runParser' :: BuilderState -> FilePath -> Text -> Either MegaparsecError BuilderState
+runParser' bs fileName input =
   case run $
     evalState @Index 0 $
       evalState @LocalNameMap mempty $
-        runInfoTableBuilder $
+        runInfoTableBuilder' bs $
           evalTopNameIdGen $
             P.runParserT parseToplevel fileName input of
     (_, Left err) -> Left (MegaparsecError err)
-    (tbl, Right ()) -> Right tbl
+    (bs', Right ()) -> Right bs'
 
 createBuiltinConstr ::
   Symbol ->
@@ -330,6 +337,8 @@ command = do
       return $ mkInstr' loc Dump
     "fail" ->
       return $ mkInstr' loc Failure
+    "argsnum" ->
+      return $ mkInstr' loc ArgsNum
     "alloc" ->
       mkInstr' loc . AllocConstr <$> constrTag
     "calloc" ->

src/Juvix/Compiler/Backend/C/Translation/FromReg.hs

@@ -232,6 +232,8 @@ fromRegInstr bNoStack info = \case
     return [StatementExpr $ macroVar "JUVIX_DUMP"]
   Reg.Failure Reg.InstrFailure {..} ->
     return [StatementExpr $ macroCall "JUVIX_FAILURE" [fromValue _instrFailureValue]]
+  Reg.ArgsNum Reg.InstrArgsNum {..} ->
+    return [StatementExpr $ macroCall "JUVIX_ARGS_NUM" [fromVarRef _instrArgsNumResult, fromValue _instrArgsNumValue]]
   Reg.Prealloc x ->
     return [fromPrealloc x]
   Reg.Alloc x ->

src/Juvix/Compiler/Reg/Extra.hs

@@ -24,6 +24,7 @@ computeMaxStackHeight lims = maximum . map go
       Trace {} -> 0
       Dump -> 0
       Failure {} -> 0
+      ArgsNum {} -> 0
       Prealloc InstrPrealloc {..} ->
         length _instrPreallocLiveVars
       Alloc {} -> 0
@@ -74,6 +75,7 @@ computeMaxCallClosuresArgsNum = maximum . map go
       Trace {} -> 0
       Dump -> 0
       Failure {} -> 0
+      ArgsNum {} -> 0
       Prealloc InstrPrealloc {} -> 0
       Alloc {} -> 0
       AllocClosure {} -> 0
@@ -118,6 +120,8 @@ computeStringMap strs = snd . run . execState (HashMap.size strs, strs) . mapM g
       Dump -> return ()
       Failure InstrFailure {..} ->
         goVal _instrFailureValue
+      ArgsNum InstrArgsNum {..} ->
+        goVal _instrArgsNumValue
       Prealloc {} -> return ()
       Alloc InstrAlloc {..} ->
         mapM_ goVal _instrAllocArgs

src/Juvix/Compiler/Reg/Language.hs

@@ -44,6 +44,7 @@ data Instruction
   | Trace InstrTrace
   | Dump
   | Failure InstrFailure
+  | ArgsNum InstrArgsNum
   | Prealloc InstrPrealloc
   | Alloc InstrAlloc
   | AllocClosure InstrAllocClosure
@@ -98,6 +99,11 @@ newtype InstrFailure = InstrFailure
   { _instrFailureValue :: Value
   }
 
+data InstrArgsNum = InstrArgsNum
+  { _instrArgsNumResult :: VarRef,
+    _instrArgsNumValue :: Value
+  }
+
 data InstrPrealloc = InstrPrealloc
   { _instrPreallocWordsNum :: Int,
     _instrPreallocLiveVars :: [VarRef]

src/Juvix/Compiler/Reg/Translation/FromAsm.hs

@@ -86,6 +86,7 @@ fromAsmInstr funInfo tab si Asm.CmdInstr {..} =
     Asm.Trace -> return $ Trace $ InstrTrace (VRef $ VarRef VarGroupStack n)
     Asm.Dump -> return Dump
     Asm.Failure -> return $ Failure $ InstrFailure (VRef $ VarRef VarGroupStack n)
+    Asm.ArgsNum -> return $ mkArgsNum (VarRef VarGroupStack n) (VRef $ VarRef VarGroupStack n)
     Asm.Prealloc x -> return $ mkPrealloc x
     Asm.AllocConstr tag -> return $ mkAlloc tag
     Asm.AllocClosure x -> return $ mkAllocClosure x
@@ -146,6 +147,9 @@ fromAsmInstr funInfo tab si Asm.CmdInstr {..} =
     mkAssign :: VarRef -> Value -> Instruction
     mkAssign tgt src = Assign (InstrAssign tgt src)
 
+    mkArgsNum :: VarRef -> Value -> Instruction
+    mkArgsNum tgt src = ArgsNum (InstrArgsNum tgt src)
+
     mkValue :: Asm.Value -> Value
     mkValue = \case
       Asm.ConstInt v -> ConstInt v

src/Juvix/Extra/Strings.hs

@@ -584,6 +584,9 @@ instrDump = "dump"
 instrFailure :: (IsString s) => s
 instrFailure = "fail"
 
+instrArgsNum :: (IsString s) => s
+instrArgsNum = "argsnum"
+
 instrPrealloc :: (IsString s) => s
 instrPrealloc = "prealloc"
 

test/Asm/Run/Positive.hs

@@ -24,6 +24,9 @@ testDescr PosTest {..} =
           _testAssertion = Steps $ asmRunAssertion file' expected' return (const (return ()))
         }
 
+filterTests :: [String] -> [PosTest] -> [PosTest]
+filterTests incl = filter (\PosTest {..} -> _name `elem` incl)
+
 allTests :: TestTree
 allTests =
   testGroup
@@ -216,5 +219,10 @@ tests =
       "Test037: String instructions"
       $(mkRelDir ".")
       $(mkRelFile "test037.jva")
-      $(mkRelFile "out/test037.out")
+      $(mkRelFile "out/test037.out"),
+    PosTest
+      "Test038: Apply & argsnum"
+      $(mkRelDir ".")
+      $(mkRelFile "test038.jva")
+      $(mkRelFile "out/test038.out")
   ]

test/Asm/Transformation.hs

@@ -1,7 +1,8 @@
 module Asm.Transformation where
 
+import Asm.Transformation.Apply qualified as Apply
 import Asm.Transformation.Prealloc qualified as Prealloc
 import Base
 
 allTests :: TestTree
-allTests = testGroup "JuvixAsm transformations" [Prealloc.allTests]
+allTests = testGroup "JuvixAsm transformations" [Prealloc.allTests, Apply.allTests]

test/Asm/Transformation/Apply.hs

@@ -0,0 +1,35 @@
+module Asm.Transformation.Apply (allTests) where
+
+import Asm.Run.Positive qualified as Run
+import Asm.Transformation.Base
+import Base
+import Juvix.Compiler.Asm.Options
+import Juvix.Compiler.Asm.Transformation
+import Juvix.Compiler.Asm.Transformation.Base
+
+allTests :: TestTree
+allTests =
+  testGroup "Apply" $
+    map liftTest $
+      Run.filterTests
+        [ "Test007: Higher-order functions",
+          "Test022: Self-application",
+          "Test025: Dynamic closure extension",
+          "Test032: Church numerals"
+        ]
+        Run.tests
+
+liftTest :: Run.PosTest -> TestTree
+liftTest _testEval =
+  fromTest
+    Test
+      { _testTransformation = runTransformation (runReader opts . computeApply),
+        _testAssertion = \tab -> unless (checkNoCallClosures opts tab) (error "check apply"),
+        _testEval
+      }
+  where
+    opts =
+      Options
+        { _optDebug = True,
+          _optLimits = getLimits TargetCWasm32Wasi True
+        }

tests/Asm/positive/out/test038.out

@@ -0,0 +1 @@
+5

tests/Asm/positive/test038.jva

@@ -0,0 +1,253 @@
+-- apply & argsnum
+
+function apply_1(*, *) : * {
+    push arg[0];
+    argsnum;
+    push 1;
+    eq;
+    br {
+        true: { -- argsnum = 1
+            push arg[1];
+            push arg[0];
+            tcall $ 1;
+        }
+        false: { -- argsnum > 1
+            push arg[1];
+            push arg[0];
+            cextend 1;
+            ret;
+        }
+    };
+}
+
+function apply_2(*, *, *) : * {
+    push arg[0];
+    argsnum;
+    tsave n {
+        push n;
+        push 2;
+        eq;
+        br {
+            true: { -- argsnum = 2
+                push arg[2];
+                push arg[1];
+                push arg[0];
+                tcall $ 2;
+            }
+            false: {
+                push n;
+                push 1;
+                eq;
+                br {
+                    true: { -- argsnum = 1
+                        push arg[2];
+                        push arg[1];
+                        push arg[0];
+                        call $ 1;
+                        tcall apply_1;
+                    }
+                    false: { -- argsnum > 2
+                        push arg[2];
+                        push arg[1];
+                        push arg[0];
+                        cextend 2;
+                        ret;
+                    }
+                };
+            }
+        };
+    };
+}
+
+function apply_3(*, *, *, *) : * {
+    push arg[0];
+    argsnum;
+    tsave n {
+        push n;
+        push 3;
+        eq;
+        br {
+            true: { -- argsnum = 3
+                push arg[3];
+                push arg[2];
+                push arg[1];
+                push arg[0];
+                tcall $ 3;
+            }
+            false: {
+                push n;
+                push 3;
+                lt;
+                br {
+                    true: { -- argsnum > 3
+                        push arg[3];
+                        push arg[2];
+                        push arg[1];
+                        push arg[0];
+                        cextend 3;
+                        ret;
+                    }
+                    false: { -- argsnum <= 2
+                        push n;
+                        push 2;
+                        eq;
+                        br {
+                            true: { -- argsnum = 2
+                                push arg[3];
+                                push arg[2];
+                                push arg[1];
+                                push arg[0];
+                                call $ 2;
+                                tcall apply_1;
+                            }
+                            false: { -- argsnum = 1
+                                push arg[3];
+                                push arg[2];
+                                push arg[1];
+                                push arg[0];
+                                call $ 1;
+                                tcall apply_2;
+                            }
+                        };
+                    }
+                };
+            }
+        };
+    };
+}
+
+function apply_4(*, *, *, *, *) : * {
+    push arg[0];
+    argsnum;
+    tsave n {
+        push n;
+        push 4;
+        eq;
+        br {
+            true: { -- argsnum = 4
+                push arg[4];
+                push arg[3];
+                push arg[2];
+                push arg[1];
+                push arg[0];
+                tcall $ 4;
+            }
+            false: {
+                push n;
+                push 4;
+                lt;
+                br {
+                    true: { -- argsnum > 4
+                        push arg[4];
+                        push arg[3];
+                        push arg[2];
+                        push arg[1];
+                        push arg[0];
+                        cextend 4;
+                        ret;
+                    }
+                    false: { -- argsnum <= 3
+                        push n;
+                        push 3;
+                        eq;
+                        br {
+                            true: { -- argsnum = 3
+                                push arg[4];
+                                push arg[3];
+                                push arg[2];
+                                push arg[1];
+                                push arg[0];
+                                call $ 3;
+                                tcall apply_1;
+                            }
+                            false: {
+                                push n;
+                                push 2;
+                                eq;
+                                br {
+                                    true: { -- argsnum = 2
+                                        push arg[4];
+                                        push arg[3];
+                                        push arg[2];
+                                        push arg[1];
+                                        push arg[0];
+                                        call $ 2;
+                                        tcall apply_2;
+                                    }
+                                    false: { -- argsnum = 1
+                                        push arg[4];
+                                        push arg[3];
+                                        push arg[2];
+                                        push arg[1];
+                                        push arg[0];
+                                        call $ 1;
+                                        tcall apply_3;
+                                    }
+                                };
+                            }
+                        };
+                    }
+                };
+            }
+        };
+    };
+}
+
+function S(*, *, *) {
+  push arg[2];
+  push arg[1];
+  call apply_1;
+  push arg[2];
+  push arg[0];
+  tcall apply_2;
+}
+
+function K(*, *) {
+  push arg[0];
+  ret;
+}
+
+function I(*) {
+  push arg[0];
+  calloc K 0;
+  push $;
+  tcall S;
+}
+
+function f3(integer, integer, integer) : integer {
+    push arg[2];
+    push arg[1];
+    push arg[0];
+    add;
+    mul;
+    ret;
+}
+
+function main() {
+    push 7;
+    push 1;
+    calloc I 0;
+    push $;
+    push $;
+    push $;
+    push $;
+    push $;
+    push $;
+    call apply_4;
+    call apply_3;
+    push 2;
+    calloc I 0;
+    push $;
+    push $;
+    call apply_3;
+    push 3;
+    calloc I 0;
+    push $;
+    push $;
+    call apply_3;
+    calloc f3 0;
+    call apply_3;
+    calloc K 0;
+    tcall apply_2;
+    -- result: 5
+}