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..175e731d0 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,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) 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,