wip: inlay-hint: prototype inferred return types hint

src/analysis/Interpreter.jl

@@ -2,10 +2,9 @@ module Interpreter
 
 export LSInterpreter
 
-using JuliaSyntax: JuliaSyntax as JS
 using JET: CC, JET
 using ..JETLS:
-    AnalysisEntry, FullAnalysisInfo, SavedFileInfo, Server,
+    JS, AnalysisEntry, FullAnalysisInfo, SavedFileInfo, Server,
     JETLS_DEV_MODE, build_tree!, get_saved_file_info, yield_to_endpoint, send
 using ..JETLS.URIs2
 using ..JETLS.LSP
@@ -18,29 +17,40 @@ Counter() = Counter(0)
 increment!(counter::Counter) = counter.count += 1
 Base.getindex(counter::Counter) = counter.count
 
+struct SigInferenceResult
+    filename::String
+    node::JS.SyntaxNode
+    result
+    SigInferenceResult(filename::String, node::JS.SyntaxNode) = new(filename, node)
+    SigInferenceResult(filename::String, node::JS.SyntaxNode, @nospecialize(result)) = new(filename, node, result)
+end
+
 struct LSInterpreter{S<:Server, I<:FullAnalysisInfo} <: JET.ConcreteInterpreter
     server::S
     info::I
     analyzer::LSAnalyzer
     counter::Counter
+    curnode::Base.RefValue{JS.SyntaxNode}
+    results::Vector{SigInferenceResult}
     state::JET.InterpretationState
-    function LSInterpreter(server::S, info::I, analyzer::LSAnalyzer, counter::Counter) where S<:Server where I<:FullAnalysisInfo
-        return new{S,I}(server, info, analyzer, counter)
+    function LSInterpreter(server::S, info::I, analyzer::LSAnalyzer, counter::Counter, curnode::Base.RefValue{JS.SyntaxNode}, results::Vector{SigInferenceResult}) where S<:Server where I<:FullAnalysisInfo
+        return new{S,I}(server, info, analyzer, counter, curnode, results)
     end
-    function LSInterpreter(server::S, info::I, analyzer::LSAnalyzer, counter::Counter, state::JET.InterpretationState) where S<:Server where I<:FullAnalysisInfo
-        return new{S,I}(server, info, analyzer, counter, state)
+    function LSInterpreter(server::S, info::I, analyzer::LSAnalyzer, counter::Counter, curnode::Base.RefValue{JS.SyntaxNode}, results::Vector{SigInferenceResult}, state::JET.InterpretationState) where S<:Server where I<:FullAnalysisInfo
+        return new{S,I}(server, info, analyzer, counter, curnode, results, state)
     end
 end
 
 # The main constructor
-LSInterpreter(server::Server, info::FullAnalysisInfo) = LSInterpreter(server, info, LSAnalyzer(info.entry), Counter())
+LSInterpreter(server::Server, info::FullAnalysisInfo) =
+    LSInterpreter(server, info, LSAnalyzer(info.entry), Counter(), Ref{JS.SyntaxNode}(), SigInferenceResult[])
 
 # `JET.ConcreteInterpreter` interface
 JET.InterpretationState(interp::LSInterpreter) = interp.state
 function JET.ConcreteInterpreter(interp::LSInterpreter, state::JET.InterpretationState)
     # add `state` to `interp`, and update `interp.analyzer.cache`
     initialize_cache!(interp.analyzer, state.res.analyzed_files)
-    return LSInterpreter(interp.server, interp.info, interp.analyzer, interp.counter, state)
+    return LSInterpreter(interp.server, interp.info, interp.analyzer, interp.counter, interp.curnode, interp.results, state)
 end
 JET.ToplevelAbstractAnalyzer(interp::LSInterpreter) = interp.analyzer
 
@@ -93,6 +103,7 @@ function JET.analyze_from_definitions!(interp::LSInterpreter, config::JET.Toplev
                 match.method, match.spec_types, match.sparams)
             reports = JET.get_reports(analyzer, result)
             append!(res.inference_error_reports, reports)
+            interp.results[i] = SigInferenceResult(interp.results[i].filename, interp.results[i].node, result)
         else
             # something went wrong
             if JETLS_DEV_MODE
@@ -130,6 +141,21 @@ function JET.virtual_process!(interp::LSInterpreter,
     return res
 end
 
+function JET.lower_with_err_handling(interp::LSInterpreter, node::JS.SyntaxNode, x::Expr)
+    interp.curnode[] = node
+    @invoke JET.lower_with_err_handling(interp::JET.ConcreteInterpreter, node::JS.SyntaxNode, x::Expr)
+end
+
+function JET.collect_toplevel_signature!(interp::LSInterpreter, frame::JET.JuliaInterpreter.Frame, @nospecialize(node))
+    res = @invoke JET.collect_toplevel_signature!(interp::JET.ConcreteInterpreter, frame::JET.JuliaInterpreter.Frame, node::Any)
+    if !isnothing(res)
+        filename = JET.InterpretationState(interp).filename
+        @assert length(res) == length(interp.results) + 1
+        push!(interp.results, SigInferenceResult(filename, interp.curnode[]))
+    end
+    return res
+end
+
 function JET.try_read_file(interp::LSInterpreter, include_context::Module, filename::AbstractString)
     uri = filename2uri(filename)
     fi = get_saved_file_info(interp.server.state, uri)

src/analysis/full-analysis.jl

@@ -75,17 +75,19 @@ function analyze_parsed_if_exist(server::Server, info::FullAnalysisInfo, args...
         @assert !isnothing(filename) lazy"Unsupported URI: $uri"
         parsed = build_tree!(JS.SyntaxNode, fi; filename)
         begin_full_analysis_progress(server, info)
+        interp = LSInterpreter(server, info)
         try
-            return JET.analyze_and_report_expr!(LSInterpreter(server, info), parsed, filename, args...; jetconfigs...)
+            return interp, JET.analyze_and_report_expr!(interp, parsed, filename, args...; jetconfigs...)
         finally
             end_full_analysis_progress(server, info)
         end
     else
         filepath = uri2filepath(uri)
         @assert filepath !== nothing lazy"Unsupported URI: $uri"
         begin_full_analysis_progress(server, info)
+        interp = LSInterpreter(server, info)
         try
-            return JET.analyze_and_report_file!(LSInterpreter(server, info), filepath, args...; jetconfigs...)
+            return interp, JET.analyze_and_report_file!(interp, filepath, args...; jetconfigs...)
         finally
             end_full_analysis_progress(server, info)
         end
@@ -103,7 +105,7 @@ function update_analyzer_world(analyzer::LSAnalyzer)
     return JET.AbstractAnalyzer(analyzer, newstate)
 end
 
-function new_analysis_unit(entry::AnalysisEntry, result)
+function new_analysis_unit(entry::AnalysisEntry, (interp, result))
     analyzed_file_infos = Dict{URI,JET.AnalyzedFileInfo}(
         # `filepath` is an absolute path (since `path` is specified as absolute)
         filename2uri(filepath) => analyzed_file_info for (filepath, analyzed_file_info) in result.res.analyzed_files)
@@ -112,12 +114,12 @@ function new_analysis_unit(entry::AnalysisEntry, result)
     successfully_analyzed_file_infos = copy(analyzed_file_infos)
     is_full_analysis_successful(result) || empty!(successfully_analyzed_file_infos)
     analysis_result = FullAnalysisResult(
-        #=staled=#false, result.res.actual2virtual::JET.Actual2Virtual, update_analyzer_world(result.analyzer),
+        #=staled=#false, result.res.actual2virtual::JET.Actual2Virtual, update_analyzer_world(result.analyzer), interp,
         uri2diagnostics, analyzed_file_infos, successfully_analyzed_file_infos)
     return AnalysisUnit(entry, analysis_result)
 end
 
-function update_analysis_unit!(analysis_unit::AnalysisUnit, result)
+function update_analysis_unit!(analysis_unit::AnalysisUnit, (interp, result))
     uri2diagnostics = analysis_unit.result.uri2diagnostics
     cached_analyzed_file_infos = analysis_unit.result.analyzed_file_infos
     cached_successfully_analyzed_file_infos = analysis_unit.result.successfully_analyzed_file_infos
@@ -143,6 +145,7 @@ function update_analysis_unit!(analysis_unit::AnalysisUnit, result)
     if is_full_analysis_successful(result)
         analysis_unit.result.actual2virtual = result.res.actual2virtual
         analysis_unit.result.analyzer = update_analyzer_world(result.analyzer)
+        analysis_unit.result.interp = interp
     end
 end
 

src/inlay-hint.jl

@@ -36,6 +36,7 @@ function handle_InlayHintRequest(server::Server, msg::InlayHintRequest)
 
     inlay_hints = InlayHint[]
     syntactic_inlay_hints!(inlay_hints, fi, range)
+    return_type_inlay_hints!(inlay_hints, server.state, uri, fi, range)
 
     return send(server, InlayHintResponse(;
         id = msg.id,
@@ -88,3 +89,48 @@ function syntactic_inlay_hints!(inlay_hints::Vector{InlayHint}, fi::FileInfo, ra
     return inlay_hints
 end
 syntactic_inlay_hints(args...) = syntactic_inlay_hints!(InlayHint[], args...) # used by tests
+
+function return_type_inlay_hints!(inlay_hints::Vector{InlayHint}, state::ServerState, uri::URI, fi::FileInfo, range::Range)
+    sfi = @something get_saved_file_info(state, uri) return nothing
+    if JS.sourcetext(fi.parsed_stream) ≠ JS.sourcetext(sfi.parsed_stream)
+        return nothing
+    end
+
+    filename = uri2filename(uri)
+    analysis_unit = find_analysis_unit_for_uri(state, uri)
+    interp = get_context_interpreter(analysis_unit)
+    postprocessor = get_post_processor(analysis_unit)
+    if !isnothing(interp)
+        for result in interp.results
+            filename == result.filename || continue
+            isdefined(result, :result) || continue
+            methodnode = result.node
+            if JS.kind(methodnode) === JS.K"doc"
+                JS.numchildren(methodnode) ≥ 2 || continue
+                methodnode = methodnode[2]
+            end
+            JS.kind(methodnode) === JS.K"function" || continue
+            JS.numchildren(methodnode) ≥ 2 || continue
+            arglist = methodnode[1]
+            JS.kind(arglist) === JS.K"call" || continue
+            overlap(get_source_range(arglist), range) || continue
+            position = offset_to_xy(fi, JS.last_byte(arglist)+1)
+            label = "::" * postprocessor(string(CC.widenconst(result.result.result)))
+            tooltip = let codeinfostr =
+                postprocessor(sprint(io->show(io, result.result.src; debuginfo=:none)))
+                value = "```\n" * codeinfostr * "\n```"
+                MarkupContent(;
+                    kind = MarkupKind.Markdown,
+                    value)
+            end
+            push!(inlay_hints, InlayHint(;
+                position,
+                label,
+                tooltip,
+                kind = InlayHintKind.Type,
+                paddingLeft = true))
+        end
+    end
+    return inlay_hints
+end
+return_type_inlay_hints(args...) = return_type_inlay_hints!(InlayHint[], args...) # used by tests

src/types.jl

@@ -110,6 +110,7 @@ mutable struct FullAnalysisResult
     staled::Bool
     actual2virtual::JET.Actual2Virtual
     analyzer::LSAnalyzer
+    interp
     const uri2diagnostics::URI2Diagnostics
     const analyzed_file_infos::Dict{URI,JET.AnalyzedFileInfo}
     const successfully_analyzed_file_infos::Dict{URI,JET.AnalyzedFileInfo}

src/utils/server.jl

@@ -72,6 +72,10 @@ get_context_analyzer(::Nothing, uri::URI) = LSAnalyzer(uri)
 get_context_analyzer(::OutOfScope, uri::URI) = LSAnalyzer(uri)
 get_context_analyzer(analysis_unit::AnalysisUnit, ::URI) = analysis_unit.result.analyzer
 
+get_context_interpreter(::Nothing) = nothing
+get_context_interpreter(::OutOfScope) = nothing
+get_context_interpreter(analysis_unit::AnalysisUnit) = analysis_unit.result.interp
+
 get_post_processor(::Nothing) = LSPostProcessor(JET.PostProcessor())
 get_post_processor(::OutOfScope) = LSPostProcessor(JET.PostProcessor())
 get_post_processor(analysis_unit::AnalysisUnit) = LSPostProcessor(JET.PostProcessor(analysis_unit.result.actual2virtual))