rework on src.slottypes and ir.argtypes (JuliaLang#49113)

This commit reworks our handling of `(src::CodeInfo).slottypes` and
`(ir::IRCode).argtypes`. Currently `ir.argtypes` contains types for call
arguments and local slots even after the SSA conversion (`slot2reg`),
which can be quite confusing. Similarly, `src.slot[names|flags|types]`
contains information about local slots regardless of whether `src` is
optimized or not, even though local slots never appear within `src`.

With this commit, we resize `ir.argtypes` so that it only contains
information about call arguments after `slot2reg`, as well as resize
`src.slot[names|flags|types]` after optimization.

This commit removes unnecessary information upon appropriate timings and
allows us to use `Core.OpaqueClosure(ir::IRCode)` constructor for such
`ir` that is retrieved by `Base.code_ircode`:

```julia
let ir = first(only(Base.code_ircode(sin, (Int,))))
    @test OpaqueClosure(ir)(42) == sin(42)
    ir = first(only(Base.code_ircode(sin, (Float64,))))
    @test OpaqueClosure(ir)(42.) == sin(42.)
end
```

Co-authored-by: Jameson Nash <[email protected]>

base/compiler/abstractinterpretation.jl

@@ -2060,7 +2060,7 @@ function most_general_argtypes(closure::PartialOpaque)
     if !isa(argt, DataType) || argt.name !== typename(Tuple)
         argt = Tuple
     end
-    return most_general_argtypes(closure.source, argt, false)
+    return most_general_argtypes(closure.source, argt, #=withfirst=#false)
 end
 
 # call where the function is any lattice element

base/compiler/optimize.jl

@@ -188,13 +188,16 @@ end
 
 function ir_to_codeinf!(opt::OptimizationState)
     (; linfo, src) = opt
-    optdef = linfo.def
-    replace_code_newstyle!(src, opt.ir::IRCode, isa(optdef, Method) ? Int(optdef.nargs) : 0)
+    src = ir_to_codeinf!(src, opt.ir::IRCode)
     opt.ir = nothing
+    validate_code_in_debug_mode(linfo, src, "optimized")
+    return src
+end
+
+function ir_to_codeinf!(src::CodeInfo, ir::IRCode)
+    replace_code_newstyle!(src, ir)
     widen_all_consts!(src)
     src.inferred = true
-    # finish updating the result struct
-    validate_code_in_debug_mode(linfo, src, "optimized")
     return src
 end
 
@@ -612,7 +615,11 @@ function convert_to_ircode(ci::CodeInfo, sv::OptimizationState)
     if cfg === nothing
         cfg = compute_basic_blocks(code)
     end
-    return IRCode(stmts, cfg, linetable, sv.slottypes, meta, sv.sptypes)
+    # NOTE this `argtypes` contains types of slots yet: it will be modified to contain the
+    # types of call arguments only once `slot2reg` converts this `IRCode` to the SSA form
+    # and eliminates slots (see below)
+    argtypes = sv.slottypes
+    return IRCode(stmts, cfg, linetable, argtypes, meta, sv.sptypes)
 end
 
 function process_meta!(meta::Vector{Expr}, @nospecialize stmt)
@@ -631,6 +638,9 @@ function slot2reg(ir::IRCode, ci::CodeInfo, sv::OptimizationState)
     defuse_insts = scan_slot_def_use(nargs, ci, ir.stmts.inst)
     𝕃ₒ = optimizer_lattice(sv.inlining.interp)
     @timeit "construct_ssa" ir = construct_ssa!(ci, ir, domtree, defuse_insts, sv.slottypes, 𝕃ₒ) # consumes `ir`
+    # NOTE now we have converted `ir` to the SSA form and eliminated slots
+    # let's resize `argtypes` now and remove unnecessary types for the eliminated slots
+    resize!(ir.argtypes, nargs)
     return ir
 end
 

base/compiler/ssair/legacy.jl

@@ -10,11 +10,7 @@ the original `ci::CodeInfo` are modified.
 """
 function inflate_ir!(ci::CodeInfo, linfo::MethodInstance)
     sptypes = sptypes_from_meth_instance(linfo)
-    if ci.inferred
-        argtypes, _ = matching_cache_argtypes(fallback_lattice, linfo)
-    else
-        argtypes = Any[ Any for i = 1:length(ci.slotflags) ]
-    end
+    argtypes, _ = matching_cache_argtypes(fallback_lattice, linfo)
     return inflate_ir!(ci, sptypes, argtypes)
 end
 function inflate_ir!(ci::CodeInfo, sptypes::Vector{VarState}, argtypes::Vector{Any})
@@ -60,15 +56,21 @@ Mainly used for testing or interactive use.
 inflate_ir(ci::CodeInfo, linfo::MethodInstance) = inflate_ir!(copy(ci), linfo)
 inflate_ir(ci::CodeInfo, sptypes::Vector{VarState}, argtypes::Vector{Any}) = inflate_ir!(copy(ci), sptypes, argtypes)
 function inflate_ir(ci::CodeInfo)
-    slottypes = ci.slottypes
-    argtypes = Any[ slottypes === nothing ? Any : slottypes[i] for i = 1:length(ci.slotflags) ]
+    parent = ci.parent
+    isa(parent, MethodInstance) && return inflate_ir(ci, parent)
+    # XXX the length of `ci.slotflags` may be different from the actual number of call
+    # arguments, but we really don't know that information in this case
+    argtypes = Any[ Any for i = 1:length(ci.slotflags) ]
     return inflate_ir(ci, VarState[], argtypes)
 end
 
-function replace_code_newstyle!(ci::CodeInfo, ir::IRCode, nargs::Int)
+function replace_code_newstyle!(ci::CodeInfo, ir::IRCode)
     @assert isempty(ir.new_nodes)
     # All but the first `nargs` slots will now be unused
+    nargs = length(ir.argtypes)
+    resize!(ci.slotnames, nargs)
     resize!(ci.slotflags, nargs)
+    resize!(ci.slottypes, nargs)
     stmts = ir.stmts
     code = ci.code = stmts.inst
     ssavaluetypes = ci.ssavaluetypes = stmts.type

base/compiler/typeinfer.jl

@@ -1004,13 +1004,13 @@ function typeinf_ext(interp::AbstractInterpreter, mi::MethodInstance)
             tree.ssavaluetypes = 1
             tree.codelocs = Int32[1]
             tree.linetable = LineInfoNode[LineInfoNode(method.module, method.name, method.file, method.line, Int32(0))]
-            tree.inferred = true
             tree.ssaflags = UInt8[0]
             set_inlineable!(tree, true)
             tree.parent = mi
             tree.rettype = Core.Typeof(rettype_const)
             tree.min_world = code.min_world
             tree.max_world = code.max_world
+            tree.inferred = true
             return tree
         elseif isa(inf, CodeInfo)
             ccall(:jl_typeinf_timing_end, Cvoid, (UInt64,), start_time)

base/compiler/types.jl

@@ -69,6 +69,9 @@ mutable struct InferenceResult
     argescapes               # ::ArgEscapeCache if optimized, nothing otherwise
     must_be_codeinf::Bool    # if this must come out as CodeInfo or leaving it as IRCode is ok
     function InferenceResult(linfo::MethodInstance, cache_argtypes::Vector{Any}, overridden_by_const::BitVector)
+        # def = linfo.def
+        # nargs = def isa Method ? Int(def.nargs) : 0
+        # @assert length(cache_argtypes) == nargs
         return new(linfo, cache_argtypes, overridden_by_const, nothing, nothing,
             WorldRange(), Effects(), Effects(), nothing, true)
     end

base/opaque_closure.jl

@@ -40,10 +40,11 @@ function compute_ir_rettype(ir::IRCode)
     return Core.Compiler.widenconst(rt)
 end
 
-function compute_oc_argtypes(ir, nargs, isva)
-    argtypes = ir.argtypes[2:end]
-    @assert nargs == length(argtypes)
-    argtypes = Core.Compiler.anymap(Core.Compiler.widenconst, argtypes)
+function compute_oc_signature(ir::IRCode, nargs::Int, isva::Bool)
+    argtypes = Vector{Any}(undef, nargs)
+    for i = 1:nargs
+        argtypes[i] = Core.Compiler.widenconst(ir.argtypes[i+1])
+    end
     if isva
         lastarg = pop!(argtypes)
         if lastarg <: Tuple
@@ -52,35 +53,42 @@ function compute_oc_argtypes(ir, nargs, isva)
             push!(argtypes, Vararg{Any})
         end
     end
-    Tuple{argtypes...}
+    return Tuple{argtypes...}
 end
 
-function Core.OpaqueClosure(ir::IRCode, env...;
-        nargs::Int = length(ir.argtypes)-1,
-        isva::Bool = false,
-        rt = compute_ir_rettype(ir),
-        do_compile::Bool = true)
-    if (isva && nargs > length(ir.argtypes)) || (!isva && nargs != length(ir.argtypes)-1)
-        throw(ArgumentError("invalid argument count"))
-    end
+function Core.OpaqueClosure(ir::IRCode, @nospecialize env...;
+                            isva::Bool = false,
+                            do_compile::Bool = true)
+    # NOTE: we need ir.argtypes[1] == typeof(env)
     ir = Core.Compiler.copy(ir)
+    nargs = length(ir.argtypes)-1
+    sig = compute_oc_signature(ir, nargs, isva)
+    rt = compute_ir_rettype(ir)
     src = ccall(:jl_new_code_info_uninit, Ref{CodeInfo}, ())
-    src.slotflags = UInt8[]
     src.slotnames = fill(:none, nargs+1)
+    src.slotflags = fill(zero(UInt8), length(ir.argtypes))
     src.slottypes = copy(ir.argtypes)
-    Core.Compiler.replace_code_newstyle!(src, ir, nargs+1)
-    Core.Compiler.widen_all_consts!(src)
-    src.inferred = true
-    # NOTE: we need ir.argtypes[1] == typeof(env)
-
-    ccall(:jl_new_opaque_closure_from_code_info, Any, (Any, Any, Any, Any, Any, Cint, Any, Cint, Cint, Any, Cint),
-        compute_oc_argtypes(ir, nargs, isva), Union{}, rt, @__MODULE__, src, 0, nothing, nargs, isva, env, do_compile)
+    src.rettype = rt
+    src = Core.Compiler.ir_to_codeinf!(src, ir)
+    return generate_opaque_closure(sig, Union{}, rt, src, nargs, isva, env...; do_compile)
 end
 
-function Core.OpaqueClosure(src::CodeInfo, env...)
-    M = src.parent.def
-    sig = Base.tuple_type_tail(src.parent.specTypes)
+function Core.OpaqueClosure(src::CodeInfo, @nospecialize env...)
+    src.inferred || throw(ArgumentError("Expected inferred src::CodeInfo"))
+    mi = src.parent::Core.MethodInstance
+    sig = Base.tuple_type_tail(mi.specTypes)
+    method = mi.def::Method
+    nargs = method.nargs-1
+    isva = method.isva
+    return generate_opaque_closure(sig, Union{}, src.rettype, src, nargs, isva, env...)
+end
 
-    ccall(:jl_new_opaque_closure_from_code_info, Any, (Any, Any, Any, Any, Any, Cint, Any, Cint, Cint, Any, Cint),
-          sig, Union{}, src.rettype, @__MODULE__, src, 0, nothing, M.nargs - 1, M.isva, env, true)
+function generate_opaque_closure(@nospecialize(sig), @nospecialize(rt_lb), @nospecialize(rt_ub),
+                                 src::CodeInfo, nargs::Int, isva::Bool, @nospecialize env...;
+                                 mod::Module=@__MODULE__,
+                                 lineno::Int=0,
+                                 file::Union{Nothing,Symbol}=nothing,
+                                 do_compile::Bool=true)
+    return ccall(:jl_new_opaque_closure_from_code_info, Any, (Any, Any, Any, Any, Any, Cint, Any, Cint, Cint, Any, Cint),
+        sig, rt_lb, rt_ub, mod, src, lineno, file, nargs, isva, env, do_compile)
 end

test/compiler/inference.jl

@@ -4099,18 +4099,18 @@ let # Test the presence of PhiNodes in lowered IR by taking the above function,
     mi = Core.Compiler.specialize_method(first(methods(f_convert_me_to_ir)),
         Tuple{Bool, Float64}, Core.svec())
     ci = Base.uncompressed_ast(mi.def)
+    ci.slottypes = Any[ Any for i = 1:length(ci.slotflags) ]
     ci.ssavaluetypes = Any[Any for i = 1:ci.ssavaluetypes]
     sv = Core.Compiler.OptimizationState(mi, Core.Compiler.NativeInterpreter())
     ir = Core.Compiler.convert_to_ircode(ci, sv)
     ir = Core.Compiler.slot2reg(ir, ci, sv)
     ir = Core.Compiler.compact!(ir)
-    Core.Compiler.replace_code_newstyle!(ci, ir, 4)
-    ci.ssavaluetypes = length(ci.code)
+    Core.Compiler.replace_code_newstyle!(ci, ir)
+    ci.ssavaluetypes = length(ci.ssavaluetypes)
     @test any(x->isa(x, Core.PhiNode), ci.code)
     oc = @eval b->$(Expr(:new_opaque_closure, Tuple{Bool, Float64}, Any, Any,
         Expr(:opaque_closure_method, nothing, 2, false, LineNumberNode(0, nothing), ci)))(b, 1.0)
     @test Base.return_types(oc, Tuple{Bool}) == Any[Float64]
-
     oc = @eval ()->$(Expr(:new_opaque_closure, Tuple{Bool, Float64}, Any, Any,
         Expr(:opaque_closure_method, nothing, 2, false, LineNumberNode(0, nothing), ci)))(true, 1.0)
     @test Base.return_types(oc, Tuple{}) == Any[Float64]

test/compiler/ssair.jl

@@ -358,6 +358,25 @@ end
     @test first(only(Base.code_ircode(demo; optimize_until = "SROA"))) isa Compiler.IRCode
 end
 
+# slots after SSA conversion
+function f_with_slots(a, b)
+    # `c` and `d` are local variables
+    c = a + b
+    d = c > 0
+    return (c, d)
+end
+let # #self#, a, b, c, d
+    unopt = code_typed1(f_with_slots, (Int,Int); optimize=false)
+    @test length(unopt.slotnames) == length(unopt.slotflags) == length(unopt.slottypes) == 5
+    ir_withslots = first(only(Base.code_ircode(f_with_slots, (Int,Int); optimize_until="convert")))
+    @test length(ir_withslots.argtypes) == 5
+    # #self#, a, b
+    opt = code_typed1(f_with_slots, (Int,Int); optimize=true)
+    @test length(opt.slotnames) == length(opt.slotflags) == length(opt.slottypes) == 3
+    ir_ssa = first(only(Base.code_ircode(f_with_slots, (Int,Int); optimize_until="slot2reg")))
+    @test length(ir_ssa.argtypes) == 3
+end
+
 let
     function test_useref(stmt, v, op)
         if isa(stmt, Expr)

test/opaque_closure.jl

@@ -256,7 +256,14 @@ let src = first(only(code_typed(+, (Int, Int))))
     @test oc(40, 2) == 42
     @test isa(oc, OpaqueClosure{Tuple{Int,Int}, Int})
     @test_throws TypeError oc("40", 2)
-    @test OpaqueClosure(ir)(40, 2) == 42 # OpaqueClosure constructor should be non-destructive
+    @test OpaqueClosure(ir)(40, 2) == 42 # the `OpaqueClosure(::IRCode)` constructor should be non-destructive
+end
+let ir = first(only(Base.code_ircode(sin, (Int,))))
+    @test OpaqueClosure(ir)(42) == sin(42)
+    @test OpaqueClosure(ir)(42) == sin(42) # the `OpaqueClosure(::IRCode)` constructor should be non-destructive
+    ir = first(only(Base.code_ircode(sin, (Float64,))))
+    @test OpaqueClosure(ir)(42.) == sin(42.)
+    @test OpaqueClosure(ir)(42.) == sin(42.) # the `OpaqueClosure(::IRCode)` constructor should be non-destructive
 end
 
 # variadic arguments