OpaqueClosure - Runtime & Codegen only

This is #37849 modulo inference/optimization support.
This is self-contained and the functionality tests
pass, though the tests testing for various optimizations
of course don't (they're marked as test_broken).
I believe I have addressed the review in the
original PR relevant to the files extracted here.

Author: Keno

base/Base.jl

@@ -232,6 +232,9 @@ using .Libc: getpid, gethostname, time
 
 include("env.jl")
 
+# OpaqueClosure
+include("opaque_closure.jl")
+
 # Concurrency
 include("linked_list.jl")
 include("condition.jl")

base/compiler/ssair/ir.jl

@@ -375,7 +375,8 @@ function is_relevant_expr(e::Expr)
                       :gc_preserve_begin, :gc_preserve_end,
                       :foreigncall, :isdefined, :copyast,
                       :undefcheck, :throw_undef_if_not,
-                      :cfunction, :method, :pop_exception)
+                      :cfunction, :method, :pop_exception,
+                      :new_opaque_closure)
 end
 
 function setindex!(x::UseRef, @nospecialize(v))

base/compiler/validation.jl

@@ -30,7 +30,8 @@ const VALID_EXPR_HEADS = IdDict{Any,Any}(
     :thunk => 1:1,
     :throw_undef_if_not => 2:2,
     :aliasscope => 0:0,
-    :popaliasscope => 0:0
+    :popaliasscope => 0:0,
+    :new_opaque_closure => 4:typemax(Int)
 )
 
 # @enum isn't defined yet, otherwise I'd use it for this

base/opaque_closure.jl

@@ -0,0 +1,23 @@
+function show(io::IO, oc::Core.OpaqueClosure{A, R}) where {A, R}
+    show_tuple_as_call(io, Symbol(""), A; hasfirst=false)
+    print(io, "::", R)
+    print(io, "->◌")
+end
+
+function show(io::IO, ::MIME"text/plain", oc::Core.OpaqueClosure{A, R}) where {A, R}
+    show(io, oc)
+end
+
+"""
+    @opaque (args...)->...
+
+Marks a given closure as "opaque". Opaque closures capture the
+world age of their creation (as opposed to their invocation).
+This allows for more aggressive optimization of the capture
+list, but trades off against the ability to inline opaque
+closures at the call site, if their creation is not statically
+visible.
+"""
+macro opaque(ex)
+    esc(Expr(:opaque_closure, ex))
+end

base/reflection.jl

@@ -1140,6 +1140,9 @@ function code_typed(@nospecialize(f), @nospecialize(types=Tuple);
     if isa(f, Core.Builtin)
         throw(ArgumentError("argument is not a generic function"))
     end
+    if isa(f, Core.OpaqueClosure)
+        return code_typed_opaque_closure(f, types; optimize, debuginfo, world, interp)
+    end
     ft = Core.Typeof(f)
     if isa(types, Type)
         u = unwrap_unionall(types)
@@ -1186,6 +1189,24 @@ function code_typed_by_type(@nospecialize(tt::Type);
     return asts
 end
 
+function code_typed_opaque_closure(closure::Core.OpaqueClosure, @nospecialize(types=Tuple);
+        optimize=true,
+        debuginfo::Symbol=:default,
+        world = get_world_counter(),
+        interp = Core.Compiler.NativeInterpreter(world))
+    ccall(:jl_is_in_pure_context, Bool, ()) && error("code reflection cannot be used from generated functions")
+    if isa(closure.ci, CodeInfo)
+        code = copy(closure.ci)
+        debuginfo === :none && remove_linenums!(code)
+        return Any[Pair{CodeInfo,Any}(code, code.rettype)]
+    else
+        mi = Core.Compiler.specialize_method(closure.ci::Method, signature_type(closure.env, types), Core.svec())
+        (code, ty) = Core.Compiler.typeinf_code(interp, mi, optimize)
+        debuginfo === :none && remove_linenums!(code)
+        return Any[Pair{CodeInfo,Any}(code, ty)]
+    end
+end
+
 function return_types(@nospecialize(f), @nospecialize(types=Tuple), interp=Core.Compiler.NativeInterpreter())
     ccall(:jl_is_in_pure_context, Bool, ()) && error("code reflection cannot be used from generated functions")
     if isa(f, Core.Builtin)

src/Makefile

@@ -45,7 +45,7 @@ RUNTIME_SRCS := \
 	simplevector runtime_intrinsics precompile \
 	threading partr stackwalk gc gc-debug gc-pages gc-stacks method \
 	jlapi signal-handling safepoint timing subtype \
-	crc32c APInt-C processor ircode
+	crc32c APInt-C processor ircode opaque_closure
 SRCS := jloptions runtime_ccall rtutils
 
 LLVMLINK :=

src/ast.c

@@ -43,6 +43,7 @@ jl_sym_t *pop_exception_sym;
 jl_sym_t *exc_sym;     jl_sym_t *error_sym;
 jl_sym_t *new_sym;     jl_sym_t *using_sym;
 jl_sym_t *splatnew_sym;
+jl_sym_t *new_opaque_closure_sym;
 jl_sym_t *const_sym;   jl_sym_t *thunk_sym;
 jl_sym_t *foreigncall_sym; jl_sym_t *as_sym;
 jl_sym_t *global_sym; jl_sym_t *list_sym;
@@ -359,6 +360,7 @@ void jl_init_common_symbols(void)
     pop_exception_sym = jl_symbol("pop_exception");
     new_sym = jl_symbol("new");
     splatnew_sym = jl_symbol("splatnew");
+    new_opaque_closure_sym = jl_symbol("new_opaque_closure");
     const_sym = jl_symbol("const");
     global_sym = jl_symbol("global");
     thunk_sym = jl_symbol("thunk");

src/builtins.c

@@ -1503,6 +1503,8 @@ void jl_init_intrinsic_functions(void) JL_GC_DISABLED
     inm->parent = jl_core_module;
     jl_set_const(jl_core_module, jl_symbol("Intrinsics"), (jl_value_t*)inm);
     jl_mk_builtin_func(jl_intrinsic_type, "IntrinsicFunction", jl_f_intrinsic_call);
+    jl_mk_builtin_func(jl_unwrap_unionall(jl_opaque_closure_type),
+        "OpaqueClosure", jl_f_opaque_closure_call);
 
 #define ADD_I(name, nargs) add_intrinsic(inm, #name, name);
 #define ADD_HIDDEN(name, nargs)
@@ -1618,6 +1620,7 @@ void jl_init_primitives(void) JL_GC_DISABLED
     add_builtin("Ptr", (jl_value_t*)jl_pointer_type);
     add_builtin("LLVMPtr", (jl_value_t*)jl_llvmpointer_type);
     add_builtin("Task", (jl_value_t*)jl_task_type);
+    add_builtin("OpaqueClosure", (jl_value_t*)jl_opaque_closure_type);
 
     add_builtin("AbstractArray", (jl_value_t*)jl_abstractarray_type);
     add_builtin("DenseArray", (jl_value_t*)jl_densearray_type);

src/clangsa/GCChecker.cpp

@@ -753,6 +753,7 @@ bool GCChecker::isGCTrackedType(QualType QT) {
                    Name.endswith_lower("jl_uniontype_t") ||
                    Name.endswith_lower("jl_method_match_t") ||
                    Name.endswith_lower("jl_vararg_t") ||
+                   Name.endswith_lower("jl_opaque_closure_t") ||
                    // Probably not technically true for these, but let's allow
                    // it
                    Name.endswith_lower("typemap_intersection_env") ||

src/codegen.cpp

@@ -862,6 +862,8 @@ static const std::map<jl_fptr_args_t, JuliaFunction*> builtin_func_map = {
     { &jl_f_apply_type,         new JuliaFunction{"jl_f_apply_type", get_func_sig, get_func_attrs} },
 };
 
+static const auto jl_new_opaque_closure_jlcall_func = new JuliaFunction{"jl_new_opaque_closure_jlcall", get_func_sig, get_func_attrs};
+
 static int globalUnique = 0;
 
 // --- code generation ---
@@ -2624,6 +2626,13 @@ static Value *emit_f_is(jl_codectx_t &ctx, const jl_cgval_t &arg1, const jl_cgva
     return emit_box_compare(ctx, arg1, arg2, nullcheck1, nullcheck2);
 }
 
+static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
+    emit_function(
+        jl_method_instance_t *lam,
+        jl_code_info_t *src,
+        jl_value_t *jlrettype,
+        jl_codegen_params_t &params);
+
 static bool emit_builtin_call(jl_codectx_t &ctx, jl_cgval_t *ret, jl_value_t *f,
                               const jl_cgval_t *argv, size_t nargs, jl_value_t *rt,
                               jl_expr_t *ex)
@@ -3228,6 +3237,7 @@ static bool emit_builtin_call(jl_codectx_t &ctx, jl_cgval_t *ret, jl_value_t *f,
         }
         return true;
     }
+
     return false;
 }
 
@@ -4477,6 +4487,98 @@ static jl_cgval_t emit_expr(jl_codectx_t &ctx, jl_value_t *expr, ssize_t ssaval)
         // it to the inferred type.
         return mark_julia_type(ctx, val, true, (jl_value_t*)jl_any_type);
     }
+    else if (head == new_opaque_closure_sym) {
+        size_t nargs = jl_array_len(ex->args);
+        assert(nargs >= 4 && "Not enough arguments in new_opaque_closure");
+        jl_cgval_t *argv = (jl_cgval_t*)alloca(sizeof(jl_cgval_t) * nargs);
+        for (size_t i = 0; i < nargs; ++i) {
+            argv[i] = emit_expr(ctx, args[i]);
+        }
+        const jl_cgval_t &argt = argv[0];
+        const jl_cgval_t &lb = argv[1];
+        const jl_cgval_t &ub = argv[2];
+        const jl_cgval_t &source = argv[3];
+        if (source.constant == NULL || argt.constant == NULL || lb.constant == NULL || ub.constant == NULL) {
+            // For now, we require non-constant source to be created using
+            // eval - we may lift this restriction if necessary.
+            emit_error(ctx, "Opaque closure source be constant");
+            return jl_cgval_t();
+        }
+        bool valid = jl_is_type(argt.constant) && jl_is_type(lb.constant) &&
+            jl_is_type(ub.constant);
+        if (valid && jl_is_code_info(source.constant) && jl_is_concrete_immutable(argt.constant)) {
+            // TODO: Emit this inline and outline it late using LLVM's coroutine
+            // support.
+            jl_code_info_t *closure_src = (jl_code_info_t *)source.constant;
+            std::unique_ptr<Module> closure_m;
+            jl_llvm_functions_t closure_decls;
+
+            jl_method_instance_t *li;
+            jl_value_t *closure_t;
+            jl_tupletype_t *env_t;
+            jl_svec_t *sig_args;
+            JL_GC_PUSH4(&li, &closure_t, &env_t, &sig_args);
+
+            li = jl_new_method_instance_uninit();
+            li->uninferred = (jl_value_t*)closure_src;
+            jl_tupletype_t *argt_typ = (jl_tupletype_t *)argt.constant;
+
+            closure_t = jl_apply_type2((jl_value_t*)jl_opaque_closure_type, (jl_value_t*)argt_typ, ub.constant);
+
+            size_t nsig = 1 + jl_svec_len(argt_typ->parameters);
+            sig_args = jl_alloc_svec_uninit(nsig);
+            jl_svecset(sig_args, 0, closure_t);
+            for (size_t i = 0; i < jl_svec_len(argt_typ->parameters); ++i) {
+                jl_svecset(sig_args, 1+i, jl_svecref(argt_typ->parameters, i));
+            }
+            li->specTypes = (jl_value_t*)jl_apply_tuple_type_v(jl_svec_data(sig_args), nsig);
+            li->def.module = ctx.module;
+            std::tie(closure_m, closure_decls) = emit_function(li, closure_src,
+                ub.constant, ctx.emission_context);
+            jl_merge_module(ctx.f->getParent(), std::move(closure_m));
+
+            jl_value_t **env_component_ts = (jl_value_t**)alloca(sizeof(jl_value_t*) * (nargs-4));
+            for (size_t i = 0; i < nargs - 4; ++i) {
+                env_component_ts[i] = argv[4+i].typ;
+            }
+
+            env_t = jl_apply_tuple_type_v(env_component_ts, nargs-4);
+
+            // TODO: Inline the env at the end of the opaque closure and generate a descriptor for GC
+            jl_cgval_t env = emit_new_struct(ctx, (jl_value_t*)env_t, nargs-4, &argv[4]);
+
+            Function *specptr = ctx.f->getParent()->getFunction(closure_decls.specFunctionObject);
+            jl_cgval_t fptr = mark_julia_type(ctx,
+                specptr ? (llvm::Value*)specptr : (llvm::Value*)ConstantPointerNull::get((llvm::PointerType*)T_pvoidfunc),
+                false, jl_voidpointer_type);
+
+            Value *jlptr;
+            if (closure_decls.functionObject == "jl_fptr_args" ||
+                closure_decls.functionObject == "jl_fptr_sparam") {
+                jlptr = specptr;
+            } else {
+                jlptr = ctx.f->getParent()->getFunction(closure_decls.functionObject);
+            }
+            jl_cgval_t jlcall_ptr = mark_julia_type(ctx,
+                jlptr ? (llvm::Value*)jlptr : (llvm::Value*)ConstantPointerNull::get((llvm::PointerType*)T_pvoidfunc),
+                false, jl_voidpointer_type);
+
+            jl_cgval_t closure_fields[4] =  {
+                env,
+                source,
+                jlcall_ptr,
+                fptr
+            };
+
+            jl_cgval_t ret = emit_new_struct(ctx, closure_t, 4, closure_fields);
+            JL_GC_POP();
+            return ret;
+        }
+
+        return mark_julia_type(ctx,
+                emit_jlcall(ctx, jl_new_opaque_closure_jlcall_func, V_rnull, argv, nargs, JLCALL_F_CC),
+                true, jl_any_type);
+    }
     else if (head == exc_sym) {
         return mark_julia_type(ctx,
                 ctx.builder.CreateCall(prepare_call(jl_current_exception_func)),
@@ -5742,15 +5844,29 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
     ctx.code = src->code;
 
     std::map<int, BasicBlock*> labels;
+    bool toplevel = false;
     ctx.module = jl_is_method(lam->def.method) ? lam->def.method->module : lam->def.module;
     ctx.linfo = lam;
+    ctx.name = name_from_method_instance(lam);
+    bool is_opaque_closure = false;
+    if (jl_is_method(lam->def.method)) {
+        ctx.nargs = lam->def.method->nargs;
+    }
+    else {
+        ctx.nargs = 0;
+        // This is an opaque closure
+        jl_datatype_t *closure = jl_first_argument_datatype(lam->specTypes);
+        if (closure && jl_is_opaque_closure_type(closure)) {
+            assert(jl_is_concrete_immutable(jl_tparam0(closure)));
+            ctx.nargs = 1 + jl_datatype_nfields(jl_tparam0(closure));
+            is_opaque_closure = true;
+        }
+    }
+    toplevel = !jl_is_method(lam->def.method);
     ctx.rettype = jlrettype;
     ctx.source = src;
-    ctx.name = name_from_method_instance(lam);
     ctx.funcName = ctx.name;
     ctx.spvals_ptr = NULL;
-    ctx.nargs = jl_is_method(lam->def.method) ? lam->def.method->nargs : 0;
-    bool toplevel = !jl_is_method(lam->def.method);
     jl_array_t *stmts = ctx.code;
     size_t stmtslen = jl_array_dim0(stmts);
 
@@ -5764,7 +5880,7 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
 
     StringRef dbgFuncName = ctx.name;
     int toplineno = -1;
-    if (jl_is_method(lam->def.method)) {
+    if (lam && jl_is_method(lam->def.method)) {
         toplineno = lam->def.method->line;
         ctx.file = jl_symbol_name(lam->def.method->file);
     }
@@ -5793,7 +5909,7 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
 
     assert(lam->specTypes); // the specTypes field should always be assigned
 
-    if (nreq > 0 && lam->def.method->isva) {
+    if (nreq > 0 && jl_is_method(lam->def.value) && lam->def.method->isva) {
         nreq--;
         va = 1;
         jl_sym_t *vn = (jl_sym_t*)jl_array_ptr_ref(src->slotnames, ctx.nargs - 1);
@@ -5820,6 +5936,15 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
         if (argname == unused_sym)
             continue;
         jl_value_t *ty = jl_nth_slot_type(lam->specTypes, i);
+        // OpaqueClosure implicitly loads the env
+        if (i == 0 && is_opaque_closure) {
+            if (jl_is_array(src->slottypes)) {
+                ty = jl_arrayref((jl_array_t*)src->slottypes, i);
+            }
+            else {
+                ty = (jl_value_t*)jl_any_type;
+            }
+        }
         varinfo.value = mark_julia_type(ctx, (Value*)NULL, false, ty);
     }
     if (va && ctx.vaSlot != -1) {
@@ -6302,6 +6427,13 @@ static std::pair<std::unique_ptr<Module>, jl_llvm_functions_t>
                 }
             }
 
+            // If this is an opaque closure, implicitly load the env
+            if (i == 0 && is_opaque_closure) {
+                theArg = convert_julia_type(ctx,
+                    emit_getfield_knownidx(ctx, theArg, 0, (jl_datatype_t*)argType),
+                    vi.value.typ);
+            }
+
             if (vi.boxroot == NULL) {
                 assert(vi.value.V == NULL && "unexpected variable slot created for argument");
                 // keep track of original (possibly boxed) value to avoid re-boxing or moving