mt-lang 0.2.0

data/lib/milk_tea/lsp/workspace/caches.rb

@@ -0,0 +1,260 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module LSP
+    class Workspace
+      module WorkspaceCaches
+        def get_content(uri)
+          @document_state_mutex.synchronize do
+            @open_documents[uri] || @indexed_documents[uri] || ''
+          end
+        end
+
+        def get_tokens(uri)
+          @tokens_cache[uri] ||= lex_document(uri)
+        end
+
+        def get_ast(uri)
+          @ast_cache[uri] ||= parse_document(uri)
+        end
+
+        def get_facts(uri, allow_last_good_fallback: true)
+          snapshot = get_tooling_snapshot(uri, allow_last_good_fallback:)
+          snapshot&.facts
+        end
+
+        def get_tooling_snapshot(uri, allow_last_good_fallback: true)
+          total_start = perf_logging? ? monotonic_time : nil
+          cache_state = 'miss'
+          lock_wait_ms = 0.0
+          analyze_ms = 0.0
+          snapshot = nil
+          last_good_snapshot = nil
+          generation = nil
+          @facts_cache_mutex.synchronize do
+            generation = @facts_generation[uri]
+            cached = @tooling_snapshot_cache[uri]
+            last_good_snapshot = @last_good_tooling_snapshot_cache[uri]
+            if cached
+              cache_state = 'hit'
+              snapshot = cached
+            end
+          end
+          return snapshot if snapshot
+
+          compute_snapshot = lambda do
+            @facts_cache_mutex.synchronize do
+              generation = @facts_generation[uri]
+              cached = @tooling_snapshot_cache[uri]
+              if cached
+                cache_state = 'hit'
+                snapshot = cached
+              end
+            end
+
+            unless snapshot
+              analyze_start = total_start ? monotonic_time : nil
+              snapshot = analyze_document(uri)
+              analyze_ms = elapsed_ms(analyze_start) if analyze_start
+              @facts_cache_mutex.synchronize do
+                if @facts_generation[uri] == generation
+                  @tooling_snapshot_cache[uri] = snapshot if snapshot
+                  if snapshot&.facts
+                    @facts_cache[uri] = snapshot.facts
+                    @last_good_facts_cache[uri] = snapshot.facts
+                    @last_good_tooling_snapshot_cache[uri] = snapshot
+                  end
+                  update_dependency_index(uri, snapshot&.facts) if snapshot
+                else
+                  cache_state = 'stale'
+                  snapshot = @tooling_snapshot_cache[uri] || @last_good_tooling_snapshot_cache[uri]
+                end
+              end
+            end
+          end
+
+          lock_wait_start = total_start ? monotonic_time : nil
+          if @facts_state_mutex.try_lock
+            begin
+              compute_snapshot.call
+            ensure
+              @facts_state_mutex.unlock
+            end
+          elsif allow_last_good_fallback && last_good_snapshot
+            cache_state = 'last_good'
+            snapshot = last_good_snapshot
+          else
+            @facts_state_mutex.synchronize do
+              lock_wait_ms = elapsed_ms(lock_wait_start) if lock_wait_start
+              compute_snapshot.call
+            end
+          end
+          snapshot
+        ensure
+          if total_start
+            result_state = snapshot&.facts.nil? ? 'nil' : 'ok'
+            log_perf_breakdown(
+              'workspace/get_tooling_snapshot',
+              elapsed_ms(total_start),
+              "uri=#{uri} cache=#{cache_state} result=#{result_state} stages_ms=lock_wait:#{lock_wait_ms},analyze:#{analyze_ms}",
+            )
+          end
+        end
+
+        def get_symbols(uri)
+          @symbols_cache[uri] ||= extract_symbols_from_tokens(uri)
+        end
+
+        def doc_comment_for_definition(uri, token)
+          return nil unless token
+
+          doc_comment_data_for_definition(uri, token)&.fetch(:raw_markdown, nil)
+        end
+
+        def doc_comment_data_for_definition(uri, token)
+          return nil unless token
+
+          docs_by_location = @doc_comments_cache[uri] ||= extract_doc_comments_for_definitions(uri)
+          docs_by_location[doc_comment_key(token.line, token.column)]
+        end
+
+        def position_to_offset(uri, line, char)
+          content = get_content(uri)
+          line_char_to_offset(content, line, char)
+        end
+
+        def all_documents
+          @document_state_mutex.synchronize do
+            (@indexed_documents.keys + @open_documents.keys).uniq
+          end
+        end
+
+        def module_name_for_uri(uri)
+          @document_module_names[uri]
+        end
+
+        # Return all identifier token locations matching +name+ across all known documents.
+        # Uses the identifier index for known documents, falling back to scanning.
+        def find_all_references(name)
+          refs = []
+          indexed_entries = @identifier_index_mutex.synchronize { @identifier_index[name]&.dup }
+          if indexed_entries
+            indexed_entries.each do |entry|
+              refs << {
+                uri: entry[:uri],
+                range: {
+                  start: { line: entry[:line] - 1, character: entry[:col] - 1 },
+                  end:   { line: entry[:line] - 1, character: entry[:col] - 1 + name.length },
+                },
+              }
+            end
+          end
+
+          unindexed = all_documents.reject { |uri| @indexed_uris.include?(uri) }
+          unless unindexed.empty?
+            refs.concat(scan_for_references(name, unindexed))
+          end
+          refs
+        end
+
+        def find_all_references_in(name, uris)
+          refs = []
+          uriset = uris.to_set
+          indexed_entries = @identifier_index_mutex.synchronize { @identifier_index[name]&.dup }
+          if indexed_entries
+            indexed_entries.each do |entry|
+              next unless uriset.include?(entry[:uri])
+
+              refs << {
+                uri: entry[:uri],
+                range: {
+                  start: { line: entry[:line] - 1, character: entry[:col] - 1 },
+                  end:   { line: entry[:line] - 1, character: entry[:col] - 1 + name.length },
+                },
+              }
+            end
+          end
+
+          unindexed = uris.reject { |uri| @indexed_uris.include?(uri) }
+          unless unindexed.empty?
+            refs.concat(scan_for_references(name, unindexed))
+          end
+          refs
+        end
+
+        def index_identifier_tokens(uri, tokens)
+          return unless tokens
+
+          @identifier_index_mutex.synchronize do
+            @indexed_uris << uri
+            tokens.each do |tok|
+              next unless tok.type == :identifier
+
+              (@identifier_index[tok.lexeme] ||= []) << { uri: uri, line: tok.line, col: tok.column }
+            end
+          end
+        end
+
+        def remove_identifier_index_entries(uri)
+          @identifier_index_mutex.synchronize do
+            @indexed_uris.delete(uri)
+            @identifier_index.each_value do |entries|
+              entries.reject! { |entry| entry[:uri] == uri }
+            end
+            @identifier_index.delete_if { |_name, entries| entries.empty? }
+          end
+        end
+
+        private
+
+        def scan_for_references(name, uris)
+          results = []
+          uris.each do |doc_uri|
+            toks = get_tokens(doc_uri)
+            next unless toks
+
+            toks.each do |tok|
+              next unless tok.type == :identifier && tok.lexeme == name
+
+              results << {
+                uri:   doc_uri,
+                range: {
+                  start: { line: tok.line - 1, character: tok.column - 1 },
+                  end:   { line: tok.line - 1, character: tok.column - 1 + name.length },
+                },
+              }
+            end
+          end
+          results
+        end
+
+        # ── Cache management ────────────────────────────────────────────────────
+
+        def invalidate_cache(uri, clear_last_good: false)
+          @tokens_cache.delete(uri)
+          @ast_cache.delete(uri)
+          @symbols_cache.delete(uri)
+          @doc_comments_cache.delete(uri)
+          @facts_cache_mutex.synchronize do
+            @facts_generation[uri] += 1
+            @facts_cache.delete(uri)
+            @tooling_snapshot_cache.delete(uri)
+            @diagnostics_cache.delete(uri)
+            @last_good_facts_cache.delete(uri) if clear_last_good
+            @last_good_tooling_snapshot_cache.delete(uri) if clear_last_good
+            @last_good_tokens_cache.delete(uri) if clear_last_good
+            @document_module_names.delete(uri) if clear_last_good
+            delete_dependency_index(uri) if clear_last_good
+          end
+          remove_identifier_index_entries(uri)
+          @definition_cache_mutex.synchronize do
+            @definition_index.each_value { |entries| entries.delete_if { |e| e[:uri] == uri } }
+            @definition_index.delete_if { |_k, v| v.empty? }
+            remove_definition_name_candidates_for_uri(uri)
+            @definition_miss_cache.clear
+          end
+        end
+      end
+    end
+  end
+end

data/lib/milk_tea/lsp/workspace/collection.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module LSP
+    class Workspace
+      module WorkspaceCollection
+        # Return cached diagnostics for +uri+, re-collecting when content or lint tier changes.
+        def collect_diagnostics(uri, lint_tier: :full)
+          total_start = perf_logging? ? monotonic_time : nil
+          content = get_content(uri)
+          hash = content.hash
+          normalized_lint_tier = Linter.normalize_lint_tier(lint_tier)
+          cache_state = 'miss'
+          lock_wait_ms = 0.0
+          collect_ms = 0.0
+          diagnostics = nil
+          generation = nil
+          cached_snapshot = nil
+          @facts_cache_mutex.synchronize do
+            generation = @facts_generation[uri]
+            cached_snapshot = @tooling_snapshot_cache[uri]
+            entry = @diagnostics_cache[uri]
+            if entry && entry[:content_hash] == hash && entry[:lint_tier] == normalized_lint_tier
+              cache_state = 'hit'
+              diagnostics = entry[:diagnostics]
+            end
+          end
+          return diagnostics if diagnostics
+
+          lock_wait_start = total_start ? monotonic_time : nil
+          @facts_state_mutex.synchronize do
+            lock_wait_ms = elapsed_ms(lock_wait_start) if lock_wait_start
+            @facts_cache_mutex.synchronize do
+              generation = @facts_generation[uri]
+              cached_snapshot = @tooling_snapshot_cache[uri]
+              entry = @diagnostics_cache[uri]
+              if entry && entry[:content_hash] == hash && entry[:lint_tier] == normalized_lint_tier
+                cache_state = 'hit'
+                diagnostics = entry[:diagnostics]
+              end
+            end
+
+            unless diagnostics
+              collect_start = total_start ? monotonic_time : nil
+              result = Diagnostics.collect(
+                uri,
+                content,
+                shared_module_cache: @shared_module_cache,
+                source_overrides: file_backed_source_overrides,
+                workspace_root_path: @workspace_root_path,
+                dependency_resolution_mode: @dependency_resolution_mode,
+                platform_override: @platform_override,
+                sema_snapshot: cached_snapshot,
+                strict_current_root_diagnostics: @strict_current_root_diagnostics_enabled,
+                lint_tier: normalized_lint_tier,
+              )
+              collect_ms = elapsed_ms(collect_start) if collect_start
+              diagnostics = result[:diagnostics]
+              facts = result[:facts]
+              snapshot = result[:sema_snapshot]
+              @facts_cache_mutex.synchronize do
+                if @facts_generation[uri] == generation
+                  @tooling_snapshot_cache[uri] = snapshot if snapshot
+                  @last_good_tooling_snapshot_cache[uri] = snapshot if snapshot&.facts
+                  @facts_cache[uri] = facts if facts
+                  @last_good_facts_cache[uri] = facts if facts
+                  @document_module_names[uri] = facts.module_name if facts&.module_name
+                  update_dependency_index(uri, facts)
+                  @diagnostics_cache[uri] = {
+                    content_hash: hash,
+                    lint_tier: normalized_lint_tier,
+                    diagnostics: diagnostics,
+                  }
+                else
+                  cache_state = 'stale'
+                end
+              end
+            end
+          end
+          diagnostics
+        rescue StandardError => e
+          cache_state = 'error'
+          warn "  #{e.backtrace.first(6).join("\n  ")}" if e.backtrace
+          []
+        ensure
+          if total_start
+            result_count = diagnostics ? diagnostics.length : 0
+            log_perf_breakdown(
+              'workspace/collect_diagnostics',
+              elapsed_ms(total_start),
+              "uri=#{uri} mode=#{normalized_lint_tier} cache=#{cache_state} diagnostics=#{result_count} stages_ms=lock_wait:#{lock_wait_ms},collect:#{collect_ms}",
+            )
+          end
+        end
+      end
+    end
+  end
+end

data/lib/milk_tea/lsp/workspace/definition_index.rb

@@ -0,0 +1,184 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module LSP
+    class Workspace
+      module WorkspaceDefinitionIndex
+        # Find a definition token across the whole workspace, preferring +preferred_uri+.
+        # Returns { uri:, token: } or nil.
+        # Uses demand-driven name lookup to avoid cold full-workspace indexing.
+        def find_definition_token_global(name, preferred_uri: nil, before_line: nil, before_char: nil)
+          # Check preferred_uri first without going through the index.
+          if preferred_uri
+            token = find_definition_token(preferred_uri, name, before_line:, before_char:)
+            return { uri: preferred_uri, token: token } if token
+          end
+
+          entries = nil
+          miss = false
+          @definition_cache_mutex.synchronize do
+            entries = @definition_index[name]&.dup
+            miss = @definition_miss_cache.include?(name)
+          end
+          if entries
+            entries.each do |entry|
+              next if entry[:uri] == preferred_uri
+
+              return entry
+            end
+          end
+
+          return nil if miss
+
+          candidate_definition_uris(name, exclude_uri: preferred_uri).each do |doc_uri|
+            token = find_definition_token(doc_uri, name)
+            next unless token
+
+            entry = { uri: doc_uri, token: token }
+            cache_definition_entry(name, entry)
+            return entry
+          end
+
+          @definition_cache_mutex.synchronize do
+            @definition_miss_cache << name
+          end
+          nil
+        end
+
+        private
+
+        def candidate_definition_uris(name, exclude_uri: nil)
+          matcher = definition_line_matcher(name)
+          open_uris, indexed_uris = @document_state_mutex.synchronize do
+            [@open_documents.keys, @indexed_documents.keys]
+          end
+
+          ordered_uris = (open_uris + indexed_uris).uniq
+          warmed_candidates = nil
+          warmed_uris = nil
+          @definition_cache_mutex.synchronize do
+            warmed_candidates = @definition_candidate_uris[name].dup
+            warmed_uris = @definition_names_by_uri.keys.to_set
+          end
+
+          matches = warmed_candidates.to_a.filter_map do |doc_uri|
+            next if doc_uri == exclude_uri
+
+            doc_uri
+          end
+
+          ordered_uris.each do |doc_uri|
+            next if doc_uri == exclude_uri
+            next if warmed_uris.include?(doc_uri)
+
+            content = get_content(doc_uri)
+            next if content.empty?
+            next unless content.match?(matcher)
+
+            warm_definition_candidates_for_uri(doc_uri, content)
+            matches << doc_uri
+          end
+
+          matches.uniq
+        end
+
+        def definition_line_matcher(name)
+          /#{DEFINITION_LINE_PREFIX}#{Regexp.escape(name)}\b/
+        end
+
+        def cache_definition_entry(name, entry)
+          @definition_cache_mutex.synchronize do
+            entries = (@definition_index[name] ||= [])
+            return if entries.any? { |existing| existing[:uri] == entry[:uri] && existing[:token].line == entry[:token].line && existing[:token].column == entry[:token].column }
+
+            entries << entry
+          end
+        end
+
+        def start_definition_warmup
+          return if @definition_warmup_thread&.alive?
+
+          @definition_warmup_thread = Thread.new do
+            Thread.current.name = 'mt-lsp-def-warmup' if Thread.current.respond_to?(:name=)
+
+            loop do
+              uri = @definition_warmup_queue.pop
+              break if uri == :__stop__
+
+              content = get_content(uri)
+              warm_definition_candidates_for_uri(uri, content)
+            end
+          rescue StandardError => e
+            warn "LSP definition warmup error: #{e.message}"
+          end
+        end
+
+        def stop_definition_warmup
+          worker = @definition_warmup_thread
+          return unless worker
+
+          @definition_warmup_queue << :__stop__
+          worker.join(1.0)
+          if worker.alive?
+            worker.kill
+            worker.join
+          end
+          @definition_warmup_thread = nil
+        rescue StandardError => e
+          warn "LSP definition warmup shutdown error: #{e.message}"
+        end
+
+        def enqueue_definition_warmup(uri)
+          return if uri.nil?
+
+          start_definition_warmup unless @definition_warmup_thread&.alive?
+
+          should_enqueue = false
+          @definition_cache_mutex.synchronize do
+            unless @definition_warmup_enqueued.include?(uri)
+              @definition_warmup_enqueued << uri
+              should_enqueue = true
+            end
+          end
+
+          @definition_warmup_queue << uri if should_enqueue
+        end
+
+        def warm_definition_candidates_for_uri(uri, content)
+          names = extract_definition_names(content)
+
+          @definition_cache_mutex.synchronize do
+            remove_definition_name_candidates_for_uri(uri)
+            @definition_names_by_uri[uri] = names
+            names.each { |name| @definition_candidate_uris[name] << uri }
+            @definition_warmup_enqueued.delete(uri)
+            @definition_miss_cache.clear unless names.empty?
+          end
+        end
+
+        def remove_definition_name_candidates_for_uri(uri)
+          names = @definition_names_by_uri.delete(uri)
+          return unless names
+
+          names.each do |name|
+            next unless @definition_candidate_uris.key?(name)
+
+            @definition_candidate_uris[name].delete(uri)
+            @definition_candidate_uris.delete(name) if @definition_candidate_uris[name].empty?
+          end
+        end
+
+        def extract_definition_names(content)
+          return Set.new if content.empty?
+
+          names = Set.new
+          content.each_line do |line|
+            match = line.match(DEFINITION_NAME_REGEX)
+            names << match[1] if match
+          end
+          names
+        end
+      end
+    end
+  end
+end