From 4fab12b95a5676d00b22596620e80c1b6ffc6591 Mon Sep 17 00:00:00 2001
From: Keno Fischer <keno@juliahub.com>
Date: Wed, 27 Mar 2024 01:18:30 +0000
Subject: [PATCH] opaque_closure: Lookup optimized oc source inside code
 instance

This is an alternative to #53852. I don't think it's semantically legal
to put optimized source into the :source field of a method, but it
should be fine to just look it up from the CodeInstance. That said,
this is a bit of an unusual configuration. In particular it wasn't
even reachable with the surface APIs, which assumed that inferred
IR was always supposed to be compiled.
---
 base/opaque_closure.jl | 12 ++++++------
 src/interpreter.c      | 14 +++++++++++++-
 test/opaque_closure.jl |  5 +++++
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/base/opaque_closure.jl b/base/opaque_closure.jl
index 83a364646ad39..b00955e7a0ca0 100644
--- a/base/opaque_closure.jl
+++ b/base/opaque_closure.jl
@@ -59,7 +59,7 @@ end
 function Core.OpaqueClosure(ir::IRCode, @nospecialize env...;
                             isva::Bool = false,
                             slotnames::Union{Nothing,Vector{Symbol}}=nothing,
-                            do_compile::Bool = true)
+                            kwargs...)
     # NOTE: we need ir.argtypes[1] == typeof(env)
     ir = Core.Compiler.copy(ir)
     # if the user didn't specify a definition MethodInstance or filename Symbol to use for the debuginfo, set a filename now
@@ -78,11 +78,11 @@ function Core.OpaqueClosure(ir::IRCode, @nospecialize env...;
     src.slotflags = fill(zero(UInt8), nargtypes)
     src.slottypes = copy(ir.argtypes)
     src = Core.Compiler.ir_to_codeinf!(src, ir)
-    return generate_opaque_closure(sig, Union{}, rt, src, nargs, isva, env...; do_compile)
+    return generate_opaque_closure(sig, Union{}, rt, src, nargs, isva, env...; kwargs...)
 end
 
-function Core.OpaqueClosure(src::CodeInfo, @nospecialize env...; rettype, sig, nargs, isva=false)
-    return generate_opaque_closure(sig, Union{}, rettype, src, nargs, isva, env...)
+function Core.OpaqueClosure(src::CodeInfo, @nospecialize env...; rettype, sig, nargs, isva=false, kwargs...)
+    return generate_opaque_closure(sig, Union{}, rettype, src, nargs, isva, env...; kwargs...)
 end
 
 function generate_opaque_closure(@nospecialize(sig), @nospecialize(rt_lb), @nospecialize(rt_ub),
@@ -90,8 +90,8 @@ function generate_opaque_closure(@nospecialize(sig), @nospecialize(rt_lb), @nosp
                                  mod::Module=@__MODULE__,
                                  lineno::Int=0,
                                  file::Union{Nothing,Symbol}=nothing,
-                                 isinferred::Bool=true,
-                                 do_compile::Bool=true)
+                                 do_compile::Bool=true,
+                                 isinferred::Bool=true)
     return ccall(:jl_new_opaque_closure_from_code_info, Any, (Any, Any, Any, Any, Any, Cint, Any, Cint, Cint, Any, Cint, Cint),
         sig, rt_lb, rt_ub, mod, src, lineno, file, nargs, isva, env, do_compile, isinferred)
 end
diff --git a/src/interpreter.c b/src/interpreter.c
index 2fb4b91927496..c7d7fa88ea447 100644
--- a/src/interpreter.c
+++ b/src/interpreter.c
@@ -760,7 +760,19 @@ JL_DLLEXPORT const jl_callptr_t jl_fptr_interpret_call_addr = &jl_fptr_interpret
 jl_value_t *jl_interpret_opaque_closure(jl_opaque_closure_t *oc, jl_value_t **args, size_t nargs)
 {
     jl_method_t *source = oc->source;
-    jl_code_info_t *code = jl_uncompress_ir(source, NULL, (jl_value_t*)source->source);
+    jl_code_info_t *code = NULL;
+    if (source->source) {
+        code = jl_uncompress_ir(source, NULL, (jl_value_t*)source->source);
+    }
+    else {
+        // OC constructed from optimized IR. It'll have a single specialization with optimized code
+        // in it that we'll try to interpret.
+        jl_svec_t *specializations = (jl_svec_t*)jl_atomic_load_relaxed(&source->specializations);
+        assert(jl_is_method_instance(specializations));
+        jl_method_instance_t *mi = (jl_method_instance_t *)specializations;
+        jl_code_instance_t *ci = jl_atomic_load_relaxed(&mi->cache);
+        code = jl_uncompress_ir(source, ci, jl_atomic_load_relaxed(&ci->inferred));
+    }
     interpreter_state *s;
     unsigned nroots = jl_source_nslots(code) + jl_source_nssavalues(code) + 2;
     jl_task_t *ct = jl_current_task;
diff --git a/test/opaque_closure.jl b/test/opaque_closure.jl
index 8378d8ae0d3dd..cb7576ee656be 100644
--- a/test/opaque_closure.jl
+++ b/test/opaque_closure.jl
@@ -358,3 +358,8 @@ let ir = first(only(Base.code_ircode(sin, (Int,))))
     oc = Core.OpaqueClosure(ir)
     @test (Base.show_method(IOBuffer(), oc.source::Method); true)
 end
+
+let ir = first(only(Base.code_ircode(sin, (Int,))))
+    oc = Core.OpaqueClosure(ir; do_compile=false)
+    @test oc(1) == sin(1)
+end