Merge pull request #648 from aviatesk/avi/per-uri-cache

sync: Group per-file cache invalidations into a helper

Author: aviatesk

src/analysis/occurrence-analysis.jl

@@ -361,9 +361,10 @@ end
 function get_binding_occurrences!(
         state::ServerState, uri::URI, fi::FileInfo, st0::JS.SyntaxTree; kwargs...
     )
+    cache_uri = canonical_cache_uri(state, uri)
     range_key = JS.byte_range(st0)
     return store!(state.binding_occurrences_cache) do cache::BindingOccurrencesCacheData
-        file_cache = get(cache, uri, nothing)
+        file_cache = get(cache, cache_uri, nothing)
         if file_cache !== nothing && haskey(file_cache, range_key)
             return cache, file_cache[range_key]
         end
@@ -383,7 +384,7 @@ function get_binding_occurrences!(
         else
             file_cache = BindingOccurrencesCacheEntry(file_cache, range_key => cache_result)
         end
-        return BindingOccurrencesCacheData(cache, uri => file_cache), cache_result
+        return BindingOccurrencesCacheData(cache, cache_uri => file_cache), cache_result
     end
 end
 
@@ -603,9 +604,10 @@ function collect_inert_global_occurrences!(
 end
 
 function invalidate_binding_occurrences_cache!(state::ServerState, uri::URI)
+    cache_uri = canonical_cache_uri(state, uri)
     store!(state.binding_occurrences_cache) do cache::BindingOccurrencesCacheData
-        if haskey(cache, uri)
-            Base.delete(cache, uri), nothing
+        if haskey(cache, cache_uri)
+            Base.delete(cache, cache_uri), nothing
         else
             cache, nothing
         end

src/diagnostic.jl

@@ -1507,11 +1507,7 @@ function get_lowering_diagnostics!(
         cancel_flag::CancelFlag;
         lookup_func = nothing
     )
-    # For notebooks, every cell URI in a notebook resolves to the same
-    # concat-source `FileInfo`, so cache under the notebook URI to share
-    # entries across cells and to match the invalidation site
-    # (`cache_notebook_file_info!` invalidates the notebook URI).
-    cache_uri = @something get_notebook_uri_for_cell(server.state, uri) uri
+    cache_uri = canonical_cache_uri(server.state, uri)
     return store!(server.state.lowering_diagnostics_cache) do cache::LoweringDiagnosticsCacheData
         if haskey(cache, cache_uri)
             return cache, cache[cache_uri]
@@ -1526,9 +1522,10 @@ function get_lowering_diagnostics!(
 end
 
 function invalidate_lowering_diagnostics_cache!(state::ServerState, uri::URI)
+    cache_uri = canonical_cache_uri(state, uri)
     store!(state.lowering_diagnostics_cache) do cache::LoweringDiagnosticsCacheData
-        if haskey(cache, uri)
-            Base.delete(cache, uri), nothing
+        if haskey(cache, cache_uri)
+            Base.delete(cache, cache_uri), nothing
         else
             cache, nothing
         end

src/did-change-watched-files.jl

@@ -149,9 +149,7 @@ function handle_jl_file_change!(server::Server, change::FileEvent)
             @assert change.type == FileChangeType.Deleted
             invalidate_unsynced_file_cache!(state, uri)
         end
-        invalidate_document_symbol_cache!(state, uri)
-        invalidate_binding_occurrences_cache!(state, uri)
-        invalidate_lowering_diagnostics_cache!(state, uri)
+        invalidate_per_file_caches!(state, uri)
     end
     request_diagnostic_refresh!(server)
 end

src/document-symbol.jl

@@ -65,23 +65,25 @@ function strip_name_from_detail(sym::DocumentSymbol)
 end
 
 function get_document_symbols!(state::ServerState, uri::URI, fi::FileInfo)
+    cache_uri = canonical_cache_uri(state, uri)
     return store!(state.document_symbol_cache) do cache::DocumentSymbolCacheData
-        if haskey(cache, uri)
-            symbols = cache[uri]
+        if haskey(cache, cache_uri)
+            symbols = cache[cache_uri]
             return cache, symbols
         end
         st0 = build_syntax_tree(fi)
         pos = Position(; line=0, character=0)
         (; mod) = get_context_info(state, uri, pos)
         symbols = extract_document_symbols(st0, fi, mod)
-        return DocumentSymbolCacheData(cache, uri => symbols), symbols
+        return DocumentSymbolCacheData(cache, cache_uri => symbols), symbols
     end
 end
 
 function invalidate_document_symbol_cache!(state::ServerState, uri::URI)
+    cache_uri = canonical_cache_uri(state, uri)
     store!(state.document_symbol_cache) do cache::DocumentSymbolCacheData
-        if haskey(cache, uri)
-            Base.delete(cache, uri), nothing
+        if haskey(cache, cache_uri)
+            Base.delete(cache, cache_uri), nothing
         else
             cache, nothing
         end

src/document-synchronization.jl

@@ -4,6 +4,17 @@ function ParseStream!(s::Union{AbstractString,Vector{UInt8}})
     return stream
 end
 
+# Drop every per-file cache entry for `uri`. Called whenever a file's content
+# changes (didChange/didOpen, notebook cell edits, watched-file events) or its
+# module context changes (full-analysis updates). `clear_*_cache!` (e.g. for
+# diagnostic config changes) is separate because not all caches share that
+# invalidation trigger.
+function invalidate_per_file_caches!(state::ServerState, uri::URI)
+    invalidate_document_symbol_cache!(state, uri)
+    invalidate_binding_occurrences_cache!(state, uri)
+    invalidate_lowering_diagnostics_cache!(state, uri)
+end
+
 """
     cache_file_info!(server::Server, uri::URI, version::Int, text::Union{AbstractString,Vector{UInt8}})
     cache_file_info!(server::Server, uri::URI, version::Int, parsed_stream::JS.ParseStream)
@@ -30,9 +41,7 @@ function cache_file_info!(
         Base.PersistentDict(cache, uri => fi), nothing
     end
 
-    invalidate_document_symbol_cache!(state, uri)
-    invalidate_binding_occurrences_cache!(state, uri)
-    invalidate_lowering_diagnostics_cache!(state, uri)
+    invalidate_per_file_caches!(state, uri)
 
     if !state.suppress_notifications && any_deleted
         notify_diagnostics!(server; ensure_cleared=uri)

src/notebook.jl

@@ -7,6 +7,15 @@ get_notebook_info(state::ServerState, uri::URI, default=nothing) =
 get_notebook_uri_for_cell(state::ServerState, cell_uri::URI, default=nothing) =
     get(load(state.cell_to_notebook), cell_uri, default)
 
+# Per-file caches (`document_symbol_cache`, `binding_occurrences_cache`,
+# `lowering_diagnostics_cache`) key on a single canonical URI per logical file.
+# For notebooks, every cell URI shares the same concat-source `FileInfo`, so the
+# canonical key is the notebook URI. Cache accessors and their invalidations
+# both run input through this helper so reads, writes, and invalidations always
+# meet on the same key.
+canonical_cache_uri(state::ServerState, uri::URI) =
+    get_notebook_uri_for_cell(state, uri, uri)
+
 function concatenate_cells(cells::Vector{NotebookCellInfo})
     source = ""
     cell_ranges = CellRange[]
@@ -28,9 +37,7 @@ function cache_notebook_file_info!(server::Server, notebook_uri::URI, notebook_i
     store!(state.file_cache) do cache
         Base.PersistentDict(cache, notebook_uri => fi), fi
     end
-    invalidate_document_symbol_cache!(state, notebook_uri)
-    invalidate_binding_occurrences_cache!(state, notebook_uri)
-    invalidate_lowering_diagnostics_cache!(state, notebook_uri)
+    invalidate_per_file_caches!(state, notebook_uri)
     return fi
 end
 

src/types.jl

@@ -683,22 +683,15 @@ mutable struct ServerState
     # Cache for files not synced via document-synchronization (unsynced files).
     # Populated on-demand by `get_unsynced_file_info!`, invalidated by `workspace/didChangeWatchedFiles`.
     const unsynced_file_cache::UnsyncedFileCache
-    # Document symbol cache for both synced and unsynced files.
-    # Uses LWContainer for concurrent writes from:
-    # - `get_document_symbols!` (on cache miss)
-    # - `textDocument/didChange` (invalidates synced files)
-    # - `workspace/didChangeWatchedFiles` (invalidates unsynced files)
+    # Per-file caches keyed on the canonical (notebook-aware) URI of a logical file.
+    # All three are dropped together via `invalidate_per_file_caches!` on content change
+    # (didChange/didOpen, notebook cell edits, watched-file events).
+    # `binding_occurrences_cache` and `lowering_diagnostics_cache` are additionally
+    # invalidated by `update_analysis_cache!` when full-analysis changes module context,
+    # since both embed `binfo.mod`. `lowering_diagnostics_cache` is also cleared wholesale
+    # on diagnostic-affecting config changes via `clear_lowering_diagnostics_cache!`.
     const document_symbol_cache::DocumentSymbolCache
-    # Binding occurrences cache for global binding analysis (references, rename).
-    # Same invalidation pattern as document_symbol_cache.
-    # TODO: This cache uses analysis context (module context from full-analysis).
-    # It should also be invalidated when full-analysis updates module context,
-    # but that is not yet implemented.
     const binding_occurrences_cache::BindingOccurrencesCache
-    # Per-file cache of the diagnostics produced by `lowering_diagnostics!` over a
-    # file's top-level statements. Lets cross-file `workspace/diagnostic`
-    # invalidations (see `compute_workspace_diagnostic_result_id`) skip re-running
-    # `jl_lower_for_scope_resolution` for files whose own content is unchanged.
     const lowering_diagnostics_cache::LoweringDiagnosticsCache
     const analysis_manager::AnalysisManager
     const extra_diagnostics::ExtraDiagnostics

src/utils/server.jl

@@ -297,7 +297,7 @@ Returns a named tuple containing:
   to recognize, which are caused by JET implementation details
 """
 function get_context_info(state::ServerState, uri::URI, pos::Position; lookup_func=nothing)
-    lookup_uri = @something get_notebook_uri_for_cell(state, uri) uri
+    lookup_uri = canonical_cache_uri(state, uri)
     if lookup_func !== nothing
         analysis_info = get_analysis_info(lookup_func, state.analysis_manager, lookup_uri)
     else
@@ -328,7 +328,7 @@ function get_context_module(analysis_result::AnalysisResult, uri::URI, pos::Posi
     return curmod
 end
 function get_context_module(state::ServerState, uri::URI, pos::Position; lookup_func=nothing)
-    lookup_uri = @something get_notebook_uri_for_cell(state, uri) uri
+    lookup_uri = canonical_cache_uri(state, uri)
     if lookup_func !== nothing
         analysis_info = get_analysis_info(lookup_func, state.analysis_manager, lookup_uri)
     else
@@ -346,7 +346,7 @@ get_post_processor(::OutOfScope) = LSPostProcessor(JET.PostProcessor())
 get_post_processor(analysis_result::AnalysisResult) = LSPostProcessor(JET.PostProcessor(analysis_result.actual2virtual))
 
 function has_analyzed_context(state::ServerState, uri::URI; lookup_func=nothing)
-    lookup_uri = @something get_notebook_uri_for_cell(state, uri) uri
+    lookup_uri = canonical_cache_uri(state, uri)
     if lookup_func !== nothing
         analysis_info = get_analysis_info(lookup_func, state.analysis_manager, lookup_uri)
     else