From 87b1fa720076fcb6ca28eb905c43d7fe92a94657 Mon Sep 17 00:00:00 2001 From: Shuhei Kadowaki Date: Sat, 27 Sep 2025 03:56:54 +0900 Subject: [PATCH 1/2] Add primitive integration with JuliaLowering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This implements a very basic integration with JL. You can retrieve inferred trees in the following way: ```julia julia> using JET, JuliaSyntax julia> result = @report_call sum("julia") ═════ 2 possible errors found ═════ ┌ sum(a::String) @ Base ./reduce.jl:553 │┌ sum(a::String; kw::@Kwargs{}) @ Base ./reduce.jl:553 ││┌ sum(f::typeof(identity), a::String) @ Base ./reduce.jl:524 │││┌ sum(f::typeof(identity), a::String; kw::@Kwargs{}) @ Base ./reduce.jl:524 ││││┌ mapreduce(f::typeof(identity), op::typeof(Base.add_sum), itr::String) @ Base ./reduce.jl:299 │││││┌ mapreduce(f::typeof(identity), op::typeof(Base.add_sum), itr::String; kw::@Kwargs{}) @ Base ./reduce.jl:299 ││││││┌ mapfoldl(f::typeof(identity), op::typeof(Base.add_sum), itr::String) @ Base ./reduce.jl:167 │││││││┌ mapfoldl(f::typeof(identity), op::typeof(Base.add_sum), itr::String; init::Base._InitialValue) @ Base ./reduce.jl:167 ││││││││┌ mapfoldl_impl(f::typeof(identity), op::typeof(Base.add_sum), nt::Base._InitialValue, itr::String) @ Base ./:36 │││││││││┌ foldl_impl(op::Base.BottomRF{typeof(Base.add_sum)}, nt::Base._InitialValue, itr::String) @ Base ./:40 ││││││││││┌ _foldl_impl(op::Base.BottomRF{typeof(Base.add_sum)}, init::Base._InitialValue, itr::String) @ Base ./reduce.jl:54 │││││││││││┌ (::Base.BottomRF{typeof(Base.add_sum)})(acc::Char, x::Char) @ Base ./:78 ││││││││││││┌ add_sum(x::Char, y::Char) @ Base ./:16 │││││││││││││ no matching method found `+(::Char, ::Char)`: (x::Char + y::Char) ││││││││││││└──────────────────── │││││││││┌ foldl_impl(op::Base.BottomRF{typeof(Base.add_sum)}, nt::Base._InitialValue, itr::String) @ Base ./:41 ││││││││││┌ reduce_empty_iter(op::Base.BottomRF{typeof(Base.add_sum)}, itr::String) @ Base ./:372 │││││││││││┌ reduce_empty_iter(op::Base.BottomRF{typeof(Base.add_sum)}, itr::String, ::Base.HasEltype) @ Base ./:373 ││││││││││││┌ reduce_empty(op::Base.BottomRF{typeof(Base.add_sum)}, ::Type{Char}) @ Base ./:349 │││││││││││││┌ reduce_empty(::typeof(Base.add_sum), ::Type{Char}) @ Base ./:342 ││││││││││││││┌ reduce_empty(::typeof(+), ::Type{Char}) @ Base ./:335 │││││││││││││││ no matching method found `zero(::Type{Char})`: zero(T::Type{Char}) ││││││││││││││└──────────────────── julia> inferred = JET.get_trees(result, :add_sum) 1-element Vector{JuliaLowering.SyntaxTree}: (method core.nothing %₈ (code_info (block Base.+ (call %₁ slot₂/x slot₃/y) (return %₂)))) julia> JuliaSyntax.sourcetext(inferred[1]) "add_sum(x, y) = x + y" ``` There are several things to note: - Please run `Pkg.instantiate()` after checking out to this branch to install JL - I couldn't think of a good way to integrate this system with the existing cache system, so for now I disabled all the global cache for now on this branch. - In cases where retrieval of the `JL.SyntaxTree` corresponding to the target `CodeInfo` for inference fails, no inferred tree will exist. Currently, retrieval likely doesn't work well for keyword methods and methods with default arguments either --- Project.toml | 25 ++----- src/JETBase.jl | 82 +++++++++++++++++++++++ src/abstractinterpret/abstractanalyzer.jl | 11 ++- src/abstractinterpret/typeinfer.jl | 47 +++++++++++++ src/analyzers/jetanalyzer.jl | 12 ++-- src/analyzers/optanalyzer.jl | 3 +- 6 files changed, 153 insertions(+), 27 deletions(-) diff --git a/Project.toml b/Project.toml index 9ef8041b3..d5b615468 100644 --- a/Project.toml +++ b/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,12 +17,14 @@ 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] @@ -29,7 +32,6 @@ 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"] diff --git a/src/JETBase.jl b/src/JETBase.jl index b93c4322c..25af39b1b 100644 --- a/src/JETBase.jl +++ b/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,92 @@ function analyze_method_instance!(analyzer::AbstractAnalyzer, mi::MethodInstance return analyze_frame!(analyzer, frame) end +function find_target_tree(stlwr::JL.SyntaxTree, result::CC.InferenceResult) + 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 + + # 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) diff --git a/src/abstractinterpret/abstractanalyzer.jl b/src/abstractinterpret/abstractanalyzer.jl index 6907abcd9..f3a5844c3 100644 --- a/src/abstractinterpret/abstractanalyzer.jl +++ b/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 diff --git a/src/abstractinterpret/typeinfer.jl b/src/abstractinterpret/typeinfer.jl index 75875799f..7aca4e33a 100644 --- a/src/abstractinterpret/typeinfer.jl +++ b/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 # ================ diff --git a/src/analyzers/jetanalyzer.jl b/src/analyzers/jetanalyzer.jl index d102635f3..ab2fb468f 100644 --- a/src/analyzers/jetanalyzer.jl +++ b/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)) diff --git a/src/analyzers/optanalyzer.jl b/src/analyzers/optanalyzer.jl index 6d576d8e0..5a9949c93 100644 --- a/src/analyzers/optanalyzer.jl +++ b/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, From 9e9328454b4da2f107d6f868c12aac545bddac38 Mon Sep 17 00:00:00 2001 From: Shuhei Kadowaki Date: Mon, 13 Oct 2025 14:28:54 +0900 Subject: [PATCH 2/2] avoid to infer invalid `CodeInfo` produced by JL --- src/JETBase.jl | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/JETBase.jl b/src/JETBase.jl index 25af39b1b..175e731d0 100644 --- a/src/JETBase.jl +++ b/src/JETBase.jl @@ -704,7 +704,8 @@ function analyze_method_instance!(analyzer::AbstractAnalyzer, mi::MethodInstance return analyze_frame!(analyzer, frame) end -function find_target_tree(stlwr::JL.SyntaxTree, result::CC.InferenceResult) +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 @@ -756,6 +757,14 @@ function CC.InferenceState(result::InferenceResult, cache_mode::UInt8, analyzer @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]