Implement the dynamic dispatch loop in JuvixAsm (#2556)

* Closes #2555 * Depends on #2554
anoma · Dec 15, 2023 · 758d1cd · 758d1cd
1 parent 76548e4
commit 758d1cd
Show file tree

Hide file tree

Showing 23 changed files with 732 additions and 14 deletions.
diff --git a/app/Commands/Dev/Asm/Compile.hs b/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)

diff --git a/cntlines.sh b/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"
diff --git a/runtime/src/asm/apply.jva b/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;
+                                    }
+                                };
+                            }
+                        };
+                    }
+                };
+            }
+        };
+    };
+}
diff --git a/runtime/src/juvix/api.h b/runtime/src/juvix/api.h
@@ -112,6 +112,8 @@
         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)

diff --git a/src/Juvix/Compiler/Asm/Data/InfoTableBuilder.hs b/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

diff --git a/src/Juvix/Compiler/Asm/Extra/Apply.hs b/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
diff --git a/src/Juvix/Compiler/Asm/Extra/Recursors.hs b/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

diff --git a/src/Juvix/Compiler/Asm/Interpreter.hs b/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

diff --git a/src/Juvix/Compiler/Asm/Language.hs b/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

diff --git a/src/Juvix/Compiler/Asm/Pretty/Base.hs b/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

diff --git a/src/Juvix/Compiler/Asm/Transformation.hs b/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