Add primitive integration with JuliaLowering

Project.toml

@@ -8,6 +8,7 @@ CodeTracking = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
 Compiler = "807dbc54-b67e-4c79-8afb-eafe4df6f2e1"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JuliaInterpreter = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"
+JuliaLowering = "f3c80556-a63f-4383-b822-37d64f81a311"
 JuliaSyntax = "70703baa-626e-46a2-a12c-08ffd08c73b4"
 LoweredCodeUtils = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
@@ -16,20 +17,21 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+[sources]
+JuliaLowering = {rev = "jetls-hacking", url = "https://github.com/mlechu/JuliaLowering.jl"}
+JuliaSyntax = {rev = "jetls-hacking", url = "https://github.com/JuliaLang/JuliaSyntax.jl"}
+
 [weakdeps]
-# Cthulhu = "f68482b8-f384-11e8-15f7-abe071a5a75f"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 
 [extensions]
-# JETCthulhuExt = "Cthulhu"
 JETReviseExt = "Revise"
 
 [compat]
 Aqua = "0.8.2"
 BenchmarkTools = "1.3.2"
 CodeTracking = "1, 2"
 Compiler = "0.1"
-# Cthulhu = "2.17"
 Downloads = "1.6.0"
 Example = "0.5.3"
 InteractiveUtils = "1.10"
@@ -53,7 +55,6 @@ julia = "1.12"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-# Cthulhu = "f68482b8-f384-11e8-15f7-abe071a5a75f"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 Example = "7876af07-990d-54b4-ab0e-23690620f79a"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
@@ -66,18 +67,4 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = [
-    "Aqua",
-    "BenchmarkTools",
-    # "Cthulhu",
-    "Downloads",
-    "Example",
-    "JSON3",
-    "Libdl",
-    "Logging",
-    "Random",
-    "Revise",
-    "StaticArrays",
-    "TOML",
-    "Test",
-]
+test = ["Aqua", "BenchmarkTools", "Downloads", "Example", "JSON3", "Libdl", "Logging", "Random", "Revise", "StaticArrays", "TOML", "Test"]

src/JETBase.jl

@@ -37,6 +37,7 @@ using Base.Meta: ParseError, isexpr, lower
 using Base.Experimental: @MethodTable, @overlay
 
 using JuliaSyntax: JuliaSyntax as JS
+using JuliaLowering: JuliaLowering as JL
 using .JS: @K_str
 
 using CodeTracking: CodeTracking
@@ -703,11 +704,101 @@ function analyze_method_instance!(analyzer::AbstractAnalyzer, mi::MethodInstance
     return analyze_frame!(analyzer, frame)
 end
 
+find_target_tree(stlwr::JL.SyntaxTree, ::InferenceResult) = find_target_tree(stlwr)
+function find_target_tree(stlwr::JL.SyntaxTree)
+    JS.kind(stlwr) === JS.K"code_info" || return nothing
+    JS.numchildren(stlwr) ≥ 1 || return nothing
+    target_tree = nothing
+    stlwr1 = stlwr[1]
+    for i = 1:JS.numchildren(stlwr1)
+        stlwrᵢ = stlwr1[i]
+        if JS.kind(stlwrᵢ) === JS.K"method" && JS.numchildren(stlwrᵢ) ≥ 3
+            mtree = stlwrᵢ[3]
+            if JS.kind(mtree) === JS.K"code_info"
+                if !isnothing(target_tree)
+                    return nothing # TODO kwfunc, func with defualt arguments
+                else
+                    target_tree = stlwrᵢ
+                end
+            end
+        end
+    end
+    return target_tree
+end
+
+prepare_type_attr(st::JL.SyntaxTree) = let g = JL.syntax_graph(st)
+    attrs = Dict(pairs(g.attributes))
+    attrs[:type] = Dict{Int, Any}()
+    return JL.SyntaxTree(JL.SyntaxGraph(g.edge_ranges, g.edges, attrs), st._id)
+end
+
+function ParseStream!(s::Union{AbstractString,Vector{UInt8}}; rule::Symbol=:all)
+    stream = JS.ParseStream(s)
+    JS.parse!(stream; rule)
+    return stream
+end
+
 function CC.InferenceState(result::InferenceResult, cache_mode::UInt8,  analyzer::AbstractAnalyzer)
     init_result!(analyzer, result)
+    m = result.linfo.def
+    if m isa Method
+        s, first_line = @something CodeTracking.definition(String, m) @goto fallback
+        ps = ParseStream!(s; rule=:statement)
+        st0 = JS.build_tree(JL.SyntaxTree, ps; first_line)
+        stlwr = try
+            JL.lower(m.module, st0)
+        catch
+            @goto fallback
+        end
+        tree = @something find_target_tree(stlwr, result) @goto fallback
+        src = try
+            JL.to_lowered_expr(tree[3])
+        catch
+            @goto fallback
+        end
+        src isa Core.CodeInfo || @goto fallback
+        src = resolve_toplevel_symbols!(src, m.module)
+
+        for i = 1:length(src.code)
+            stmt = src.code[i]
+            if stmt isa GlobalRef && stmt.name === :new
+                @goto fallback
+            end
+        end
+
+        # Validate the mapping between this tree and `CodeInfo`. Otherwise type annotation will fail.
+        tree_stmts = tree[3][1]
+        length(src.code) == JS.numchildren(tree_stmts) || @goto fallback
+        world = CC.get_inference_world(analyzer)
+        mi = result.linfo
+        orig_src = CC.retrieve_code_info(mi, world)
+        length(src.code) == length(orig_src.code) || @goto fallback
+
+        # XXX JL bug probably. Need to fix up things so that inference can actually use this `src` for infernece
+        src.nargs = orig_src.nargs
+        src.isva = orig_src.isva
+
+        analyzer.state.tree_cache[result] = prepare_type_attr(tree)
+        return CC.InferenceState(result, src, cache_mode, analyzer)
+    end
+    @label fallback
     return @invoke InferenceState(result::InferenceResult, cache_mode::UInt8, analyzer::AbstractInterpreter)
 end
 
+function get_trees(result::JETCallResult, name::Symbol)
+    results = collect(keys(result.analyzer.state.analysis_results))
+    idxs = findall(results) do r::InferenceResult
+        m = r.linfo.def
+        m isa Method || return false
+        m.name === name
+    end
+    trees = JL.SyntaxTree[]
+    for r in results[idxs]
+        push!(trees, @something get(result.analyzer.state.tree_cache, r, nothing) continue)
+    end
+    return trees
+end
+
 function analyze_frame!(analyzer::AbstractAnalyzer, frame::InferenceState)
     set_entry!(analyzer, frame.linfo)
     tworld = typeinf_world(analyzer)

src/abstractinterpret/abstractanalyzer.jl

@@ -141,6 +141,8 @@ mutable struct AnalyzerState
 
     # some `AbstractAnalyzer` may want to use this
     entry::Union{Nothing,MethodInstance}
+
+    const tree_cache::IdDict{CC.InferenceResult, JL.SyntaxTree}
 end
 
 # define shortcut getter/setter methods for `AbstractAnalyzer`s
@@ -165,7 +167,8 @@ function AnalyzerState(world::UInt  = get_world_counter();
                          #=cache_target::Union{Nothing,Pair{Symbol,InferenceState}}=# nothing,
                          #=concretized::BitVector=# non_toplevel_concretized,
                          #=binding_states::AbstractAbstractBindings=# AbstractAbstractBindings(),
-                         #=entry::Union{Nothing,MethodInstance}=# nothing)
+                         #=entry::Union{Nothing,MethodInstance}=# nothing,
+                         #=tree_cache=# IdDict{CC.InferenceResult, JL.SyntaxTree}())
 end
 
 # The constructor that inherits from existing `state::AnalyzerState`
@@ -175,7 +178,8 @@ function AnalyzerState(state::AnalyzerState, refresh_local_cache::Bool=true;
                        opt_params::OptimizationParams = state.opt_params,
                        concretized::BitVector = state.concretized,
                        binding_states::AbstractAbstractBindings = state.binding_states,
-                       entry::Union{Nothing,MethodInstance} = state.entry)
+                       entry::Union{Nothing,MethodInstance} = state.entry,
+                       tree_cache::IdDict{CC.InferenceResult, JL.SyntaxTree} = state.tree_cache)
     if refresh_local_cache
         inf_cache = InferenceResult[]
         analysis_results = IdDict{InferenceResult,AnalysisResult}()
@@ -191,7 +195,8 @@ function AnalyzerState(state::AnalyzerState, refresh_local_cache::Bool=true;
                          #=cache_target=# nothing,
                          concretized,
                          binding_states,
-                         entry)
+                         entry,
+                         tree_cache)
 end
 
 # dummies for non-toplevel analysis

src/abstractinterpret/typeinfer.jl

@@ -411,6 +411,53 @@ function CC.finish!(analyzer::AbstractAnalyzer, frame::InferenceState, validatio
     return @invoke CC.finish!(analyzer::AbstractInterpreter, frame::InferenceState, validation_world::UInt, time_before::UInt64)
 end
 
+function annotate_types!(citree::JL.SyntaxTree, frame::CC.InferenceState)
+    for i = 1:length(frame.src.code)
+        stmt = frame.src.code[i]
+        stmttype = frame.src.ssavaluetypes[i]
+        stmttree = citree[i]
+        if JS.kind(stmttree) in JS.KSet"newvar goto gotoifnot"
+            # The `ssavaluetype` corresponding to these nodes is always `Any`, and since
+            # the provenance information for these nodes is very broad, it's more convenient
+            # for the implementation of `get_type_for_range` to leave them untyped
+            continue
+        end
+        JL.setattr!(stmttree; type=stmttype)
+        if stmt isa Expr
+            stmt.head === :meta && continue
+            treeref = stmttree
+            if stmt.head === :(=)
+                lhs = stmt.args[1]
+                if lhs isa Core.SlotNumber
+                    JL.setattr!(treeref[1]; type=stmttype)
+                end
+                stmt = stmt.args[2]
+                stmt isa Expr || continue
+                treeref = treeref[2]
+            end
+            for i = 1:length(stmt.args)
+                arg = stmt.args[i]
+                if arg isa Core.SlotNumber
+                    argtyp = CC.argextype(arg, frame.src, frame.sptypes)
+                    JL.setattr!(treeref[i]; type=argtyp)
+                end
+            end
+        elseif stmt isa ReturnNode
+            rettyp = CC.argextype(stmt.val, frame.src, frame.sptypes)
+            JL.setattr!(stmttree; type=rettyp)
+        end
+    end
+end
+
+function CC.finishinfer!(frame::CC.InferenceState, analyzer::AbstractAnalyzer, cycleid::Int)
+    ret = @invoke CC.finishinfer!(frame::CC.InferenceState, analyzer::AbstractInterpreter, cycleid::Int)
+    if haskey(analyzer.state.tree_cache, frame.result)
+        tree = analyzer.state.tree_cache[frame.result]
+        annotate_types!(tree[3][1], frame)
+    end
+    return ret
+end
+
 # top-level bridge
 # ================
 

src/analyzers/jetanalyzer.jl

@@ -112,7 +112,8 @@ JETInterface.AnalyzerState(analyzer::JETAnalyzer) = analyzer.state
 function JETInterface.AbstractAnalyzer(analyzer::T, state::AnalyzerState) where T<:JETAnalyzer
     method_table = CachedMethodTable(OverlayMethodTable(state.world, JET_METHOD_TABLE))
     cache_key = compute_hash(state.inf_params, nameof(T), analyzer.config)
-    analysis_token = get!(AnalysisToken, JET_ANALYZER_CACHE, cache_key)
+    # analysis_token = get!(AnalysisToken, JET_ANALYZER_CACHE, cache_key)
+    analysis_token = AnalysisToken() # disable cache for now
     return T(state, analysis_token, method_table, analyzer.config)
 end
 JETInterface.AnalysisToken(analyzer::JETAnalyzer) = analyzer.analysis_token
@@ -1380,15 +1381,18 @@ function JETAnalyzer(world::UInt = Base.get_world_counter();
     # Create the appropriate analyzer type based on mode
     if mode === :basic
         cache_key = compute_hash(state.inf_params, BasicJETAnalyzer, config, __cache_hash__)
-        analysis_token = get!(AnalysisToken, JET_ANALYZER_CACHE, cache_key)
+        # analysis_token = get!(AnalysisToken, JET_ANALYZER_CACHE, cache_key)
+        analysis_token = AnalysisToken() # disable global cache for now
         return BasicJETAnalyzer(state, analysis_token, method_table, config)
     elseif mode === :sound
         cache_key = compute_hash(state.inf_params, SoundJETAnalyzer, config, __cache_hash__)
-        analysis_token = get!(AnalysisToken, JET_ANALYZER_CACHE, cache_key)
+        # analysis_token = get!(AnalysisToken, JET_ANALYZER_CACHE, cache_key)
+        analysis_token = AnalysisToken() # disable global cache for now
         return SoundJETAnalyzer(state, analysis_token, method_table, config)
     elseif mode === :typo
         cache_key = compute_hash(state.inf_params, TypoJETAnalyzer, config, __cache_hash__)
-        analysis_token = get!(AnalysisToken, JET_ANALYZER_CACHE, cache_key)
+        # analysis_token = get!(AnalysisToken, JET_ANALYZER_CACHE, cache_key)
+        analysis_token = AnalysisToken() # disable global cache for now
         return TypoJETAnalyzer(state, analysis_token, method_table, config)
     else
         throw(JETConfigError("`mode` configuration should be either of `:basic`, `:sound` or `:typo`", :mode, mode))

src/analyzers/optanalyzer.jl

@@ -361,7 +361,8 @@ function OptAnalyzer(world::UInt = Base.get_world_counter();
     cache_key = compute_hash(state.inf_params, state.opt_params, OptAnalyzer,
                              skip_noncompileable_calls, __cache_hash__)
     cache_key = @invoke hash(function_filter::Any, cache_key::UInt) # HACK avoid dynamic dispatch
-    analysis_token = get!(AnalysisToken, OPT_ANALYZER_CACHE, cache_key)
+    # analysis_token = get!(AnalysisToken, OPT_ANALYZER_CACHE, cache_key)
+    analysis_token = AnalysisToken() # disable cache for now
     return OptAnalyzer(
         state,
         analysis_token,