Implement opaque closures

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/boot.jl

@@ -788,6 +788,10 @@ Unsigned(x::Union{Float16, Float32, Float64, Bool}) = UInt(x)
 Integer(x::Integer) = x
 Integer(x::Union{Float16, Float32, Float64}) = Int(x)
 
+# This is a utility that can be inserted by the optimizer. Eventually it may
+# need to be a builtin, but for now, it works ok
+getfield_opaque_env(o::OpaqueClosure, i::Int) = getfield(o.env, i)
+
 # Binding for the julia parser, called as
 #
 #    Core._parse(text, filename, offset, options)

base/compiler/abstractinterpretation.jl

@@ -100,7 +100,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
     edges = Any[]
     nonbot = 0  # the index of the only non-Bottom inference result if > 0
     seen = 0    # number of signatures actually inferred
-    istoplevel = sv.linfo.def isa Module
+    istoplevel = sv.linfo !== nothing && sv.linfo.def isa Module
     multiple_matches = napplicable > 1
 
     if f !== nothing && napplicable == 1 && is_method_pure(applicable[1]::MethodMatch)
@@ -203,6 +203,8 @@ function const_prop_profitable(@nospecialize(arg))
             isconstType(b) && return true
             const_prop_profitable(b) && return true
         end
+    elseif isa(arg, PartialOpaque)
+        return true
     elseif !isa(arg, Const) || (isa(arg.val, Symbol) || isa(arg.val, Type) || (!isa(arg.val, String) && !ismutable(arg.val)))
         # don't consider mutable values or Strings useful constants
         return true
@@ -349,7 +351,7 @@ function abstract_call_method(interp::AbstractInterpreter, method::Method, @nosp
     sv_method2 isa Method || (sv_method2 = nothing) # Union{Method, Nothing}
     while !(infstate === nothing)
         infstate = infstate::InferenceState
-        if method === infstate.linfo.def
+        if infstate.linfo !== nothing && method === infstate.linfo.def
             if infstate.linfo.specTypes == sig
                 # avoid widening when detecting self-recursion
                 # TODO: merge call cycle and return right away
@@ -1015,6 +1017,34 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
     return abstract_call_gf_by_type(interp, f, argtypes, atype, sv, max_methods)
 end
 
+function abstract_call_opaque_closure(interp::AbstractInterpreter, clos::PartialOpaque, argtypes::Vector{Any}, sv::InferenceState)
+    if isa(clos.ci, CodeInfo)
+        nargtypes = argtypes[2:end]
+        pushfirst!(nargtypes, clos.env)
+        result = InferenceResult(Core.OpaqueClosure, nargtypes)
+        state = InferenceState(result, copy(clos.ci), false, interp)
+        typeinf_local(interp, state)
+        finish(state, interp)
+        result.src.src.inferred = true
+        clos.ci = result.src
+        return CallMeta(result.result, false)
+    elseif isa(clos.ci, OptimizationState)
+        return CallMeta(clos.ci.src.rettype, nothing)
+    else
+        nargtypes = argtypes[2:end]
+        pushfirst!(nargtypes, clos.env)
+        sig = argtypes_to_type(nargtypes)
+        rt, edgecycle, edge = abstract_call_method(interp, clos.ci::Method, sig, Core.svec(), false, sv)
+        if !edgecycle
+            const_rettype = abstract_call_method_with_const_args(interp, rt, argtypes[1], nargtypes, MethodMatch(sig, Core.svec(), clos.ci::Method, false), sv, edgecycle)
+            if const_rettype ⊑ rt
+               rt = const_rettype
+            end
+        end
+        return CallMeta(rt, edge)
+    end
+end
+
 # call where the function is any lattice element
 function abstract_call(interp::AbstractInterpreter, fargs::Union{Nothing,Vector{Any}}, argtypes::Vector{Any},
                        sv::InferenceState, max_methods::Int = InferenceParams(interp).MAX_METHODS)
@@ -1026,6 +1056,8 @@ function abstract_call(interp::AbstractInterpreter, fargs::Union{Nothing,Vector{
         f = ft.parameters[1]
     elseif isa(ft, DataType) && isdefined(ft, :instance)
         f = ft.instance
+    elseif isa(ft, PartialOpaque)
+        return abstract_call_opaque_closure(interp, ft, argtypes, sv)
     else
         # non-constant function, but the number of arguments is known
         # and the ft is not a Builtin or IntrinsicFunction
@@ -1195,6 +1227,22 @@ function abstract_eval_statement(interp::AbstractInterpreter, @nospecialize(e),
                 t = PartialStruct(t, at.fields)
             end
         end
+    elseif e.head === :new_opaque_closure
+        t = Union{}
+        if length(e.args) >= 4
+            ea = e.args
+            n = length(ea)
+            argtypes = Vector{Any}(undef, n)
+            @inbounds for i = 1:n
+                ai = abstract_eval_value(interp, ea[i], vtypes, sv)
+                if ai === Bottom
+                    return Bottom
+                end
+                argtypes[i] = ai
+            end
+            t = _opaque_closure_tfunc(argtypes[1], argtypes[2], argtypes[3],
+                argtypes[4], argtypes[5:end], sv.linfo)
+        end
     elseif e.head === :foreigncall
         abstract_eval_value(interp, e.args[1], vtypes, sv)
         t = sp_type_rewrap(e.args[2], sv.linfo, true)
@@ -1339,14 +1387,14 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
             elseif isa(stmt, ReturnNode)
                 pc´ = n + 1
                 rt = widenconditional(abstract_eval_value(interp, stmt.val, s[pc], frame))
-                if !isa(rt, Const) && !isa(rt, Type) && !isa(rt, PartialStruct)
+                if !isa(rt, Const) && !isa(rt, Type) && !isa(rt, PartialStruct) && !isa(rt, PartialOpaque)
                     # only propagate information we know we can store
                     # and is valid inter-procedurally
                     rt = widenconst(rt)
                 end
                 if tchanged(rt, frame.bestguess)
                     # new (wider) return type for frame
-                    frame.bestguess = tmerge(frame.bestguess, rt)
+                    frame.bestguess = frame.bestguess === NOT_FOUND ? rt : tmerge(frame.bestguess, rt)
                     for (caller, caller_pc) in frame.cycle_backedges
                         # notify backedges of updated type information
                         typeassert(caller.stmt_types[caller_pc], VarTable) # we must have visited this statement before

base/compiler/inferenceresult.jl

@@ -3,7 +3,7 @@
 function is_argtype_match(@nospecialize(given_argtype),
                           @nospecialize(cache_argtype),
                           overridden_by_const::Bool)
-    if isa(given_argtype, Const) || isa(given_argtype, PartialStruct)
+    if isa(given_argtype, Const) || isa(given_argtype, PartialStruct) || isa(given_argtype, PartialOpaque)
         return is_lattice_equal(given_argtype, cache_argtype)
     end
     return !overridden_by_const

base/compiler/inferencestate.jl

@@ -5,7 +5,7 @@ const LineNum = Int
 mutable struct InferenceState
     params::InferenceParams
     result::InferenceResult # remember where to put the result
-    linfo::MethodInstance
+    linfo::Union{MethodInstance, Nothing}
     sptypes::Vector{Any}    # types of static parameter
     slottypes::Vector{Any}
     mod::Module
@@ -37,6 +37,8 @@ mutable struct InferenceState
     callers_in_cycle::Vector{InferenceState}
     parent::Union{Nothing, InferenceState}
 
+    has_opaque_closures::Bool
+
     # TODO: move these to InferenceResult / Params?
     cached::Bool
     limited::Bool
@@ -57,9 +59,23 @@ mutable struct InferenceState
                             cached::Bool, interp::AbstractInterpreter)
         linfo = result.linfo
         code = src.code::Array{Any,1}
-        toplevel = !isa(linfo.def, Method)
 
-        sp = sptypes_from_meth_instance(linfo::MethodInstance)
+        if !isa(linfo, Nothing)
+            toplevel = !isa(linfo.def, Method)
+            sp = sptypes_from_meth_instance(linfo::MethodInstance)
+            if !toplevel
+                meth = linfo.def
+                inmodule = meth.module
+            else
+                inmodule = linfo.def::Module
+            end
+        else
+            linfo = nothing
+            toplevel = true
+            inmodule = Core
+            sp = Any[]
+        end
+        code = src.code::Array{Any,1}
 
         nssavalues = src.ssavaluetypes::Int
         src.ssavaluetypes = Any[ NOT_FOUND for i = 1:nssavalues ]
@@ -93,13 +109,6 @@ mutable struct InferenceState
         W = BitSet()
         push!(W, 1) #initial pc to visit
 
-        if !toplevel
-            meth = linfo.def
-            inmodule = meth.module
-        else
-            inmodule = linfo.def::Module
-        end
-
         valid_worlds = WorldRange(src.min_world,
             src.max_world == typemax(UInt) ? get_world_counter() : src.max_world)
         frame = new(
@@ -112,7 +121,7 @@ mutable struct InferenceState
             ssavalue_uses, throw_blocks,
             Vector{Tuple{InferenceState,LineNum}}(), # cycle_backedges
             Vector{InferenceState}(), # callers_in_cycle
-            #=parent=#nothing,
+            #=parent=#nothing, #= has_opaque_closures =# false,
             cached, false, false, false,
             CachedMethodTable(method_table(interp)),
             interp)
@@ -244,6 +253,7 @@ end
 
 # temporarily accumulate our edges to later add as backedges in the callee
 function add_backedge!(li::MethodInstance, caller::InferenceState)
+    caller.linfo !== nothing || return # don't add backedges to opauqe closures
     isa(caller.linfo.def, Method) || return # don't add backedges to toplevel exprs
     if caller.stmt_edges[caller.currpc] === nothing
         caller.stmt_edges[caller.currpc] = []

base/compiler/optimize.jl

@@ -34,7 +34,7 @@ struct InliningState{S <: Union{EdgeTracker, Nothing}, T <: Union{InferenceCache
 end
 
 mutable struct OptimizationState
-    linfo::MethodInstance
+    linfo::Union{MethodInstance, Nothing}
     src::CodeInfo
     stmt_info::Vector{Any}
     mod::Module
@@ -313,6 +313,13 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}
                 ftyp = argextype(farg, src, sptypes, slottypes)
             end
         end
+        # Give calls to OpaqueClosures zero cost. The plan is for these to be a single
+        # indirect call so have very little cost. On the other hand, there
+        # is enormous benefit to inlining these into a function where we can
+        # see the definition of the OpaqueClosure. Perhaps this should even be negative
+        if widenconst(ftyp) <: Core.OpaqueClosure
+            return 0
+        end
         f = singleton_type(ftyp)
         if isa(f, IntrinsicFunction)
             iidx = Int(reinterpret(Int32, f::IntrinsicFunction)) + 1

base/compiler/ssair/driver.jl

@@ -34,15 +34,35 @@ function normalize(@nospecialize(stmt), meta::Vector{Any})
     return stmt
 end
 
-function convert_to_ircode(ci::CodeInfo, code::Vector{Any}, coverage::Bool, nargs::Int, sv::OptimizationState)
+function add_opaque_closure_argtypes!(argtypes::Vector{Any}, @nospecialize(t))
+    dt = unwrap_unionall(t)
+    dt1 = unwrap_unionall(dt.parameters[1])
+    if isa(dt1, TypeVar) || isa(dt1.parameters[1], TypeVar)
+        push!(argtypes, Any)
+    else
+        TT = dt1.parameters[1]
+        if isa(TT, Union)
+            TT = tuplemerge(TT.a, TT.b)
+        end
+        for p in TT.parameters
+            push!(argtypes, rewrap_unionall(p, t))
+        end
+    end
+end
+
+
+function convert_to_ircode(ci::CodeInfo, code::Vector{Any},
+                           coverage::Bool, nargs::Int, sv::OptimizationState,
+                           slottypes=sv.slottypes, stmtinfo=sv.stmt_info)
     # Go through and add an unreachable node after every
     # Union{} call. Then reindex labels.
     idx = 1
     oldidx = 1
     changemap = fill(0, length(code))
     labelmap = coverage ? fill(0, length(code)) : changemap
     prevloc = zero(eltype(ci.codelocs))
-    stmtinfo = sv.stmt_info
+    stmtinfo = copy(stmtinfo)
+    opaques = IRCode[]
     while idx <= length(code)
         codeloc = ci.codelocs[idx]
         if coverage && codeloc != prevloc && codeloc != 0
@@ -58,7 +78,21 @@ function convert_to_ircode(ci::CodeInfo, code::Vector{Any}, coverage::Bool, narg
             idx += 1
             prevloc = codeloc
         end
-        if code[idx] isa Expr && ci.ssavaluetypes[idx] === Union{}
+        stmt = code[idx]
+        if isexpr(stmt, :(=))
+            stmt = stmt.args[2]
+        end
+        ssat = ci.ssavaluetypes[idx]
+        if isa(ssat, PartialOpaque) && isexpr(stmt, :new_opaque_closure)
+            # Pre-convert any OpaqueClosure objects
+            if isa(ssat.ci, OptimizationState)
+                opaque_ir = make_ir(ssat.ci.src, 0, ssat.ci)
+                push!(opaques, opaque_ir)
+                stmt.head = :new_opaque_closure
+                stmt.args[4] = OpaqueClosureIdx(length(opaques))
+            end
+        end
+        if stmt isa Expr && ssat === Union{}
             if !(idx < length(code) && isa(code[idx + 1], ReturnNode) && !isdefined((code[idx + 1]::ReturnNode), :val))
                 # insert unreachable in the same basic block after the current instruction (splitting it)
                 insert!(code, idx + 1, ReturnNode())
@@ -106,7 +140,7 @@ function convert_to_ircode(ci::CodeInfo, code::Vector{Any}, coverage::Bool, narg
     cfg = compute_basic_blocks(code)
     types = Any[]
     stmts = InstructionStream(code, types, stmtinfo, ci.codelocs, flags)
-    ir = IRCode(stmts, cfg, collect(LineInfoNode, ci.linetable), sv.slottypes, meta, sv.sptypes)
+    ir = IRCode(stmts, cfg, collect(LineInfoNode, ci.linetable), slottypes, meta, sv.sptypes, opaques)
     return ir
 end
 
@@ -118,24 +152,37 @@ function slot2reg(ir::IRCode, ci::CodeInfo, nargs::Int, sv::OptimizationState)
     return ir
 end
 
-function run_passes(ci::CodeInfo, nargs::Int, sv::OptimizationState)
-    preserve_coverage = coverage_enabled(sv.mod)
-    ir = convert_to_ircode(ci, copy_exprargs(ci.code), preserve_coverage, nargs, sv)
+function compact_all!(ir::IRCode)
+    length(ir.stmts) == 0 && return ir
+    for i in 1:length(ir.opaques)
+        ir.opaques[i] = compact_all!(ir.opaques[i])
+    end
+    compact!(ir)
+end
+
+function make_ir(ci::CodeInfo, nargs::Int, sv::OptimizationState)
+    ir = convert_to_ircode(ci, copy_exprargs(ci.code), coverage_enabled(sv.mod), nargs, sv)
     ir = slot2reg(ir, ci, nargs, sv)
+    ir
+end
+
+function run_passes(ci::CodeInfo, nargs::Int, sv::OptimizationState)
+    ir = make_ir(ci, nargs, sv)
     #@Base.show ("after_construct", ir)
     # TODO: Domsorting can produce an updated domtree - no need to recompute here
     @timeit "compact 1" ir = compact!(ir)
     @timeit "Inlining" ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
     #@timeit "verify 2" verify_ir(ir)
-    ir = compact!(ir)
+    ir = compact_all!(ir)
     #@Base.show ("before_sroa", ir)
     @timeit "SROA" ir = getfield_elim_pass!(ir)
+    ir = opaque_closure_optim_pass!(ir)
     #@Base.show ir.new_nodes
     #@Base.show ("after_sroa", ir)
     ir = adce_pass!(ir)
     #@Base.show ("after_adce", ir)
     @timeit "type lift" ir = type_lift_pass!(ir)
-    @timeit "compact 3" ir = compact!(ir)
+    @timeit "compact 3" ir = compact_all!(ir)
     #@Base.show ir
     if JLOptions().debug_level == 2
         @timeit "verify 3" (verify_ir(ir); verify_linetable(ir.linetable))