ics: directly call astInterpretFunc for calls up to 2 args

src/codegen/ast_interpreter.cpp

@@ -68,8 +68,7 @@ class ASTInterpreter {
 public:
     ASTInterpreter(CLFunction* clfunc, Box** vregs);
 
-    void initArguments(int nargs, BoxedClosure* closure, BoxedGenerator* generator, Box* arg1, Box* arg2, Box* arg3,
-                       Box** args);
+    void initArguments(BoxedClosure* closure, BoxedGenerator* generator, Box* arg1, Box* arg2, Box* arg3, Box** args);
 
     static Box* execute(ASTInterpreter& interpreter, CFGBlock* start_block = NULL, AST_stmt* start_at = NULL);
     static Box* executeInner(ASTInterpreter& interpreter, CFGBlock* start_block, AST_stmt* start_at);
@@ -251,8 +250,8 @@ ASTInterpreter::ASTInterpreter(CLFunction* clfunc, Box** vregs)
     assert(scope_info);
 }
 
-void ASTInterpreter::initArguments(int nargs, BoxedClosure* _closure, BoxedGenerator* _generator, Box* arg1, Box* arg2,
-                                   Box* arg3, Box** args) {
+void ASTInterpreter::initArguments(BoxedClosure* _closure, BoxedGenerator* _generator, Box* arg1, Box* arg2, Box* arg3,
+                                   Box** args) {
     passed_closure = _closure;
     generator = _generator;
 
@@ -280,8 +279,7 @@ void ASTInterpreter::initArguments(int nargs, BoxedClosure* _closure, BoxedGener
             val = createDict();
         doStore(param_names.kwarg_name, Value(val, 0));
     }
-
-    assert(nargs == i);
+    assert(i == param_names.totalParameters());
 }
 
 void ASTInterpreter::startJITing(CFGBlock* block, int exit_offset) {
@@ -1722,8 +1720,8 @@ static int calculateNumVRegs(CLFunction* clfunc) {
     return cfg->sym_vreg_map.size();
 }
 
-Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* generator, Box* globals, Box* arg1,
-                          Box* arg2, Box* arg3, Box** args) {
+Box* astInterpretFunction(CLFunction* clfunc, Box* closure, Box* generator, Box* globals, Box* arg1, Box* arg2,
+                          Box* arg3, Box** args) {
     UNAVOIDABLE_STAT_TIMER(t0, "us_timer_in_interpreter");
 
     SourceInfo* source_info = clfunc->source.get();
@@ -1744,7 +1742,7 @@ Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* gene
             new_effort = EffortLevel::MAXIMAL;
 
         std::vector<ConcreteCompilerType*> arg_types;
-        for (int i = 0; i < nargs; i++) {
+        for (int i = 0; i < clfunc->param_names.totalParameters(); i++) {
             Box* arg = getArg(i, arg1, arg2, arg3, args);
 
             assert(arg || i == clfunc->param_names.kwargsIndex()); // only builtin functions can pass NULL args
@@ -1813,7 +1811,7 @@ Box* astInterpretFunction(CLFunction* clfunc, int nargs, Box* closure, Box* gene
         interpreter.setGlobals(source_info->parent_module);
     }
 
-    interpreter.initArguments(nargs, (BoxedClosure*)closure, (BoxedGenerator*)generator, arg1, arg2, arg3, args);
+    interpreter.initArguments((BoxedClosure*)closure, (BoxedGenerator*)generator, arg1, arg2, arg3, args);
     Box* v = ASTInterpreter::execute(interpreter);
     return v ? v : None;
 }
@@ -1829,7 +1827,7 @@ Box* astInterpretFunctionEval(CLFunction* clfunc, Box* globals, Box* boxedLocals
     }
 
     ASTInterpreter interpreter(clfunc, vregs);
-    interpreter.initArguments(0, NULL, NULL, NULL, NULL, NULL, NULL);
+    interpreter.initArguments(NULL, NULL, NULL, NULL, NULL, NULL);
     interpreter.setBoxedLocals(boxedLocals);
 
     assert(!clfunc->source->scoping->areGlobalsFromModule());

src/codegen/ast_interpreter.h

@@ -70,8 +70,8 @@ struct Value {
     Value(Box* o, RewriterVar* var) : o(o), var(var) {}
 };
 
-Box* astInterpretFunction(CLFunction* f, int nargs, Box* closure, Box* generator, Box* globals, Box* arg1, Box* arg2,
-                          Box* arg3, Box** args);
+Box* astInterpretFunction(CLFunction* f, Box* closure, Box* generator, Box* globals, Box* arg1, Box* arg2, Box* arg3,
+                          Box** args);
 Box* astInterpretFunctionEval(CLFunction* cf, Box* globals, Box* boxedLocals);
 Box* astInterpretDeopt(CLFunction* cf, AST_expr* after_expr, AST_stmt* enclosing_stmt, Box* expr_val,
                        FrameStackState frame_state);

src/codegen/irgen/hooks.cpp

@@ -326,7 +326,7 @@ void compileAndRunModule(AST_Module* m, BoxedModule* bm) {
     }
 
     UNAVOIDABLE_STAT_TIMER(t0, "us_timer_interpreted_module_toplevel");
-    Box* r = astInterpretFunction(clfunc, 0, NULL, NULL, NULL, NULL, NULL, NULL, NULL);
+    Box* r = astInterpretFunction(clfunc, NULL, NULL, NULL, NULL, NULL, NULL, NULL);
     assert(r == None);
 }
 

src/runtime/objmodel.cpp

@@ -3936,20 +3936,30 @@ static Box* callChosenCF(CompiledFunction* chosen_cf, BoxedClosure* closure, Box
 // This function exists for the rewriter: astInterpretFunction takes 9 args, but the rewriter
 // only supports calling functions with at most 6 since it can currently only pass arguments
 // in registers.
-static Box* astInterpretHelper(CLFunction* f, int num_args, BoxedClosure* closure, BoxedGenerator* generator,
-                               Box* globals, Box** _args) {
+static Box* astInterpretHelper(CLFunction* f, BoxedClosure* closure, BoxedGenerator* generator, Box* globals,
+                               Box** _args) {
     Box* arg1 = _args[0];
     Box* arg2 = _args[1];
     Box* arg3 = _args[2];
     Box* args = _args[3];
 
-    return astInterpretFunction(f, num_args, closure, generator, globals, arg1, arg2, arg3, (Box**)args);
+    return astInterpretFunction(f, closure, generator, globals, arg1, arg2, arg3, (Box**)args);
 }
 
-static Box* astInterpretHelperCapi(CLFunction* f, int num_args, BoxedClosure* closure, BoxedGenerator* generator,
-                                   Box* globals, Box** _args) noexcept {
+static Box* astInterpretHelperCapi(CLFunction* f, BoxedClosure* closure, BoxedGenerator* generator, Box* globals,
+                                   Box** _args) noexcept {
     try {
-        return astInterpretHelper(f, num_args, closure, generator, globals, _args);
+        return astInterpretHelper(f, closure, generator, globals, _args);
+    } catch (ExcInfo e) {
+        setCAPIException(e);
+        return NULL;
+    }
+}
+
+static Box* astInterpretHelper2ArgsCapi(CLFunction* f, BoxedClosure* closure, BoxedGenerator* generator, Box* globals,
+                                        Box* arg1, Box* arg2) noexcept {
+    try {
+        return astInterpretFunction(f, closure, generator, globals, arg1, arg2, NULL, NULL);
     } catch (ExcInfo e) {
         setCAPIException(e);
         return NULL;
@@ -3988,42 +3998,53 @@ Box* callCLFunc(CLFunction* f, CallRewriteArgs* rewrite_args, int num_output_arg
             // TODO this kind of embedded reference needs to be tracked by the GC somehow?
             // Or maybe it's ok, since we've guarded on the function object?
             arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)f, Location::forArg(0)));
-            arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)num_output_args, Location::forArg(1)));
-            arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)closure, Location::forArg(2)));
-            arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)generator, Location::forArg(3)));
-            arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)globals, Location::forArg(4)));
-
-            // Hacky workaround: the rewriter can only pass arguments in registers, so use this helper function
-            // to unpack some of the additional arguments:
-            // TODO if there's only one arg we could just pass it normally
-            RewriterVar* arg_array = rewrite_args->rewriter->allocate(4);
-            arg_vec.push_back(arg_array);
-            if (num_output_args >= 1)
-                arg_array->setAttr(0, rewrite_args->arg1);
-            if (num_output_args >= 2)
-                arg_array->setAttr(8, rewrite_args->arg2);
-            if (num_output_args >= 3)
-                arg_array->setAttr(16, rewrite_args->arg3);
-            if (num_output_args >= 4)
-                arg_array->setAttr(24, rewrite_args->args);
+            arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)closure, Location::forArg(1)));
+            arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)generator, Location::forArg(2)));
+            arg_vec.push_back(rewrite_args->rewriter->loadConst((intptr_t)globals, Location::forArg(3)));
+
+            if (num_output_args <= 2) {
+                if (num_output_args >= 1)
+                    arg_vec.push_back(rewrite_args->arg1);
+                if (num_output_args >= 2)
+                    arg_vec.push_back(rewrite_args->arg2);
+
+                if (S == CXX)
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretFunction, arg_vec);
+                else
+                    rewrite_args->out_rtn
+                        = rewrite_args->rewriter->call(true, (void*)astInterpretHelper2ArgsCapi, arg_vec);
+            } else {
+                // Hacky workaround: the rewriter can only pass arguments in registers, so use this helper function
+                // to unpack some of the additional arguments:
+                RewriterVar* arg_array = rewrite_args->rewriter->allocate(4);
+                arg_vec.push_back(arg_array);
+                if (num_output_args >= 1)
+                    arg_array->setAttr(0, rewrite_args->arg1);
+                if (num_output_args >= 2)
+                    arg_array->setAttr(8, rewrite_args->arg2);
+                if (num_output_args >= 3)
+                    arg_array->setAttr(16, rewrite_args->arg3);
+                if (num_output_args >= 4)
+                    arg_array->setAttr(24, rewrite_args->args);
+
+                if (S == CXX)
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelper, arg_vec);
+                else
+                    rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelperCapi, arg_vec);
+            }
 
-            if (S == CXX)
-                rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelper, arg_vec);
-            else
-                rewrite_args->out_rtn = rewrite_args->rewriter->call(true, (void*)astInterpretHelperCapi, arg_vec);
             rewrite_args->out_success = true;
         }
 
         if (S == CAPI) {
             try {
-                return astInterpretFunction(f, num_output_args, closure, generator, globals, oarg1, oarg2, oarg3,
-                                            oargs);
+                return astInterpretFunction(f, closure, generator, globals, oarg1, oarg2, oarg3, oargs);
             } catch (ExcInfo e) {
                 setCAPIException(e);
                 return NULL;
             }
         } else {
-            return astInterpretFunction(f, num_output_args, closure, generator, globals, oarg1, oarg2, oarg3, oargs);
+            return astInterpretFunction(f, closure, generator, globals, oarg1, oarg2, oarg3, oargs);
         }
     }