Add a "context" hook for generic calls during codegen (#36398)

This adds a new context field that rewrites generic calls like
`apply_generic(f, (args...,))` to `apply_generic(context, (f, args...))`
during codegen. The intention here is to allow external AbstractInterpreters
to provide custom implementations of apply_generic (usually recursing
analysis using the same interpreter, but other behavior may be
desired). This is a bit of a stopgap solution. I think in the fullness
of time, we'll probably want completely custom codegen for generic
callsites, to avoid the potential of a double-dispatch impact, but
for the moment this allows prototyping.

Author: Keno

base/reflection.jl

@@ -983,18 +983,22 @@ struct CodegenParams
 
     lookup::Ptr{Cvoid}
 
+    generic_context::Any
+
     function CodegenParams(; track_allocations::Bool=true, code_coverage::Bool=true,
                    static_alloc::Bool=true, prefer_specsig::Bool=false,
                    gnu_pubnames=true, debug_info_kind::Cint = default_debug_info_kind(),
                    module_setup=nothing, module_activation=nothing, raise_exception=nothing,
                    emit_function=nothing, emitted_function=nothing,
-                   lookup::Ptr{Cvoid}=cglobal(:jl_rettype_inferred))
+                   lookup::Ptr{Cvoid}=cglobal(:jl_rettype_inferred),
+                   generic_context = nothing)
         return new(
             Cint(track_allocations), Cint(code_coverage),
             Cint(static_alloc), Cint(prefer_specsig),
             Cint(gnu_pubnames), debug_info_kind,
             module_setup, module_activation, raise_exception,
-            emit_function, emitted_function, lookup)
+            emit_function, emitted_function, lookup,
+            generic_context)
     end
 end
 

src/cgutils.cpp

@@ -2872,5 +2872,6 @@ static int compare_cgparams(const jl_cgparams_t *a, const jl_cgparams_t *b)
            // hooks
            (a->module_setup == b->module_setup) &&
            (a->module_activation == b->module_activation) &&
-           (a->raise_exception == b->raise_exception);
+           (a->raise_exception == b->raise_exception) &&
+           (a->generic_context == b->generic_context);
 }

src/codegen.cpp

@@ -821,7 +821,7 @@ extern "C" {
         1,
 #endif
         jl_default_debug_info_kind, NULL, NULL, NULL, NULL, NULL,
-        jl_rettype_inferred };
+        jl_rettype_inferred, NULL };
 }
 
 template<typename T>
@@ -3379,7 +3379,15 @@ static jl_cgval_t emit_call(jl_codectx_t &ctx, jl_expr_t *ex, jl_value_t *rt)
         return emit_intrinsic(ctx, fi, args, nargs - 1);
     }
 
-    jl_cgval_t *argv = (jl_cgval_t*)alloca(sizeof(jl_cgval_t) * nargs);
+    jl_value_t *context = ctx.params->generic_context == jl_nothing ? nullptr : ctx.params->generic_context;
+    size_t n_generic_args = nargs + (context ? 1 : 0);
+
+    jl_cgval_t *generic_argv = (jl_cgval_t*)alloca(sizeof(jl_cgval_t) * n_generic_args);
+    jl_cgval_t *argv = generic_argv;
+    if (context) {
+        generic_argv[0] = mark_julia_const(context);
+        argv = &generic_argv[1];
+    }
     argv[0] = f;
     for (size_t i = 1; i < nargs; ++i) {
         argv[i] = emit_expr(ctx, args[i]);
@@ -3405,7 +3413,7 @@ static jl_cgval_t emit_call(jl_codectx_t &ctx, jl_expr_t *ex, jl_value_t *rt)
     }
 
     // emit function and arguments
-    Value *callval = emit_jlcall(ctx, jlapplygeneric_func, nullptr, argv, nargs, JLCALL_F_CC);
+    Value *callval = emit_jlcall(ctx, jlapplygeneric_func, nullptr, generic_argv, n_generic_args, JLCALL_F_CC);
     return mark_julia_type(ctx, callval, true, rt);
 }
 
@@ -7425,6 +7433,7 @@ extern "C" void jl_init_llvm(void)
     jl_default_cgparams.raise_exception = jl_nothing;
     jl_default_cgparams.emit_function = jl_nothing;
     jl_default_cgparams.emitted_function = jl_nothing;
+    jl_default_cgparams.generic_context = jl_nothing;
     jl_init_debuginfo();
 
     InitializeNativeTarget();

src/julia.h

@@ -2100,6 +2100,10 @@ typedef struct {
 
     // Cache access. Default: jl_rettype_inferred.
     jl_codeinstance_lookup_t lookup;
+
+    // If not `nothing`, rewrite all generic calls to call
+    // generic_context(f, args...) instead of f(args...).
+    jl_value_t *generic_context;
 } jl_cgparams_t;
 extern JL_DLLEXPORT jl_cgparams_t jl_default_cgparams;
 extern JL_DLLEXPORT int jl_default_debug_info_kind;