improve virtual global var assignement

src/TypeProfiler.jl

@@ -32,7 +32,7 @@ import Core.Compiler:
     InternalCodeCache, CodeInstance, WorldRange,
     MethodInstance, Bottom, NOT_FOUND, MethodMatchInfo, UnionSplitInfo, MethodLookupResult,
     Const, VarTable, SSAValue, SlotNumber, Slot, slot_id, GlobalRef, GotoIfNot, ReturnNode,
-    widenconst, isconstType, typeintersect, ⊑, Builtin, CallMeta, is_throw_call,
+    widenconst, isconstType, typeintersect, ⊑, Builtin, CallMeta, is_throw_call, tmerge,
     argtypes_to_type, abstract_eval_ssavalue, _methods_by_ftype, specialize_method, typeinf
 
 import Base:

src/abstractinterpretation.jl

@@ -127,48 +127,50 @@ end
 #     return ret
 # end
 
-# works as the native (i.e. "make as much progress on `frame` as possible (without handling
-# cycles)", but when `interp` is profiling on virtual toplevel lambda and `frame` is in
-# virtual toplevel (i.e. when `isroot(frame) === true`), keep the traced types of  `SlotNumber`s
-# (which are originally global variables) in `TPInterpreter.virtual_globalvar_table` so that
-# they can be referred across profilings on different virtual (toplevel) functions
-# NOTE:
-# virtual global assignments should happen here because `SlotNumber`s can be optimized away
-# after the optimization happens
 function typeinf_local(interp::TPInterpreter, frame::InferenceState)
     set_current_frame!(interp, frame)
 
     ret = invoke_native(typeinf_local, interp, frame)
 
-    # assign virtual global variable for toplevel frames
-    if istoplevel(interp) && isroot(frame)
-        for (pc, stmt) in enumerate(frame.src.code)
-            isexpr(stmt, :(=)) && set_virtual_globalvar!(interp, frame, pc, stmt)
-        end
+    # assign virtual global variable
+    # TODO: maybe move this into the `typeinf` overload ?
+    # currently this needs to be done here since `set_virtual_globalvar!` assumes we can
+    # access to the type of lhs via `frame.src.ssavaluetypes`, which doesn't hold true after
+    # optimization
+    for (pc, stmt) in enumerate(frame.src.code)
+        is_global_assign(stmt) && set_virtual_globalvar!(interp, frame, pc, stmt)
     end
 
     return ret
 end
 
-istoplevel(interp::TPInterpreter) = interp.istoplevel
+is_global_assign(@nospecialize(_)) = false
+is_global_assign(ex::Expr)         = isexpr(ex, :(=)) && first(ex.args) isa GlobalRef
 
 function get_virtual_globalvar(interp, mod, sym)
-    haskey(interp.virtual_globalvar_table, mod) || return nothing
-    return get(interp.virtual_globalvar_table[mod], sym, nothing)
+    vgvt4mod = get(interp.virtual_globalvar_table, mod, nothing)
+    isnothing(vgvt4mod) && return nothing
+    id2vgv = get(vgvt4mod, sym, nothing)
+    isnothing(id2vgv) && return nothing
+    return last(id2vgv)
 end
 
 function set_virtual_globalvar!(interp, frame, pc, stmt)
-    mod = frame.mod
-    haskey(interp.virtual_globalvar_table, mod) || (interp.virtual_globalvar_table[mod] = Dict())
-
-    slt = first(stmt.args)::Slot
-    lhs = frame.src.slotnames[slt.id]::Symbol
-    rhs = frame.src.ssavaluetypes[pc]
-    if rhs === NOT_FOUND
-        rhs = Bottom
+    gr = first(stmt.args)::GlobalRef
+    vgvt4mod = get!(interp.virtual_globalvar_table, gr.mod, Dict())
+
+    sym = gr.name
+    id = get_id(interp)
+    prev_id, prev_typ = get!(vgvt4mod, sym, id => Bottom)
+
+    typ = frame.src.ssavaluetypes[pc]
+    if typ === NOT_FOUND
+        typ = Bottom
     end
 
-    interp.virtual_globalvar_table[mod][lhs] = rhs
+    vgvt4mod[sym] = id => id === prev_id ?
+                          tmerge(prev_typ, typ) :
+                          typ
 end
 
 """

src/abstractinterpreterinterface.jl

@@ -11,21 +11,20 @@ struct TPInterpreter <: AbstractInterpreter
 
     #= TypeProfiler.jl specific =#
 
+    # for escaping reporting duplicated cached reports, sequential assignment of virtual global variable
+    id::Symbol
+
     # caching
     # -------
-    # for escaping reporting duplicated cached reports
-    id::Symbol
     # for constructing virtual stack frame from cached reports
     current_frame::Ref{InferenceState}
     # keeping reports from frame that always `throw`
     exception_reports::Vector{Pair{Int,ExceptionReport}}
 
     # virtual toplevel execution
     # --------------------------
-    # tracks this the currrent frame is toplevel (and should assign virtual global variables)
-    istoplevel::Bool
-    # keeps virtual global variables
-    virtual_globalvar_table::Dict{Module,Dict{Symbol,Any}} # maybe we don't need this nested dicts
+    # module -> (sym -> (id => type))
+    virtual_globalvar_table::Dict{Module,Dict{Symbol,Pair{Symbol,Any}}}
 
     # profiling
     # ---------
@@ -40,7 +39,6 @@ struct TPInterpreter <: AbstractInterpreter
                            optimize::Bool = true,
                            compress::Bool = false,
                            discard_trees::Bool = false,
-                           istoplevel::Bool = false,
                            virtual_globalvar_table::AbstractDict = Dict(),
                            filter_native_remarks::Bool = true,
                            )
@@ -55,7 +53,6 @@ struct TPInterpreter <: AbstractInterpreter
                    id,
                    Ref{InferenceState}(),
                    [],
-                   istoplevel,
                    virtual_globalvar_table,
                    filter_native_remarks,
                    [],

src/virtualprocess.jl

@@ -50,8 +50,11 @@ the _transform_ includes:
 - try to expand macros in a context of `virtualmod`
 - handle `include` by recursively calling this function on the `include`d file
 - replace self-reference of `actualmodsym` with that of `virtualmod`
-- remove `const` annotations so that remaining code block can be wrapped into a virtual
-  function (they are not allowed in a function body)
+- tweak assignments
+  * remove `const` annotations so that remaining code block can be wrapped into a virtual
+    function (they are not allowed in a function body)
+  * annotate regular assignments with `global`, which virtual global variable assignment looks
+    for during abstract interpretation (see [`typeinf_local(interp::TPInterpreter, frame::InferenceState)`](@ref))
 
 !!! warning
     this approach involves following limitations:
@@ -142,11 +145,15 @@ function virtual_process!(toplevelex, filename, actualmodsym, virtualmod, interp
         :quote in scope && return x
 
         # expand macro
+        # ------------
+
         if isexpr(x, :macrocall)
             x = macroexpand_with_err_handling(virtualmod, x)
         end
 
-        # hoist and evaluate these toplevel expressions
+        # hoist and evaluate toplevel expressions
+        # ---------------------------------------
+
         # these shouldn't happen when in function scope, since then they can be wrongly
         # hoisted and evaluated when they're wrapped in closures, while they actually cause
         # syntax errors; if they're in other "toplevel definition"s, they will just cause
@@ -162,20 +169,30 @@ function virtual_process!(toplevelex, filename, actualmodsym, virtualmod, interp
         end
 
         # evaluate toplevel function definitions
+        # --------------------------------------
+
         !islocalscope(scope) && isfuncdef(x) && return eval_with_err_handling(virtualmod, x)
 
-        # remove `const` annotation
+        # tweak assignments
+        # -----------------
+
+        # remove `const` annotation, annotate `global` instead
         if isexpr(x, :const)
             return if !islocalscope(scope)
-                first(x.args)
+                Expr(:global, first(x.args))
             else
                 report = SyntaxErrorReport("unsupported `const` declaration on local variable", filename, line)
                 push!(ret.toplevel_error_reports, report)
                 nothing
             end
         end
 
-        # handle static `include` call
+        # annotate `global`s for regular assignments in global scope
+        !(islocalscope(scope) || is_already_scoped(scope)) && is_assignment(x) && return Expr(:global, x)
+
+        # handle static `include` calls
+        # -----------------------------
+
         if isinclude(x)
             # TODO: maybe find a way to handle two args `include` calls
             include_file = eval_with_err_handling(virtualmod, last(x.args))
@@ -262,7 +279,6 @@ function virtual_process!(toplevelex, filename, actualmodsym, virtualmod, interp
                                optimize                = may_optimize(interp),
                                compress                = may_compress(interp),
                                discard_trees           = may_discard_trees(interp),
-                               istoplevel              = !isa(x, Symbol) && !islocalscope(x), # disable virtual global variable assignment when profiling non-toplevel blocks
                                virtual_globalvar_table = interp.virtual_globalvar_table, # pass on virtual global variable table
                                filter_native_remarks   = interp.filter_native_remarks,
                                )
@@ -354,6 +370,9 @@ function isfuncdef(ex)
     return false
 end
 
+is_already_scoped(scope) = last(scope) in (:local, :global)
+is_assignment(x) = isexpr(x, :(=)) && Base.isidentifier(first(x.args))
+
 isinclude(ex) = isexpr(ex, :call) && first(ex.args) === :include
 
 islnn(@nospecialize(_)) = false

test/runtests.jl

@@ -20,6 +20,9 @@ end
 import Base:
     Fix1, Fix2
 
+import Core.Compiler:
+    ⊑
+
 const CC = Core.Compiler
 
 const FIXTURE_DIR = normpath(@__DIR__, "fixtures")