LSP: auto-require completion for unrequired constants

.gitignore

@@ -2,4 +2,7 @@
 /pkg/
 dog_bench.stackprof.dump
 dog_bench.pf2profile
-dog_bench.json
+dog_bench.json
+
+rbs_collection*.yaml
+.gem_rbs_collection

lib/typeprof/core/service.rb

@@ -7,21 +7,20 @@ def initialize(options)
       @rbs_text_nodes = {}
 
       @genv = GlobalEnv.new
-      @genv.load_core_rbs(load_rbs_declarations(@options[:rbs_collection]).declarations, @options[:position_encoding])
+      @genv.load_core_rbs(load_rbs_declarations.declarations, @options[:position_encoding])
 
       Builtin.new(genv).deploy
+
+      @constant_catalog = options[:constant_catalog]
+
+      @rbs_loader = RBS::EnvironmentLoader.new(core_root: nil)
+      @rbs_loader.repository.add(@options[:rbs_collection].fullpath) if @options[:rbs_collection]
     end
 
-    def load_rbs_declarations(rbs_collection)
-      if rbs_collection
-        loader = RBS::EnvironmentLoader.new
-        loader.add_collection(rbs_collection)
-        RBS::Environment.from_loader(loader)
-      else
-        return $raw_rbs_env if defined?($raw_rbs_env)
-        loader = RBS::EnvironmentLoader.new
-        $raw_rbs_env = RBS::Environment.from_loader(loader)
-      end
+    def load_rbs_declarations
+      return $raw_rbs_env if defined?($raw_rbs_env)
+      loader = RBS::EnvironmentLoader.new
+      $raw_rbs_env = RBS::Environment.from_loader(loader)
     end
 
     attr_reader :genv
@@ -58,6 +57,7 @@ def update_rb_file(path, code)
       prev_node = @rb_text_nodes[path]
 
       code = File.read(path) unless code
+      load_libraries_for_requires(code)
       node = AST.parse_rb(path, code, @options[:position_encoding])
       return false unless node
 
@@ -114,6 +114,33 @@ def update_rb_file(path, code)
       return true
     end
 
+    def load_libraries_for_requires(code)
+      code.each_line do |line|
+        if md = line.match(/\A\s*require\s*\(?\s*['"]([^'"]+)['"]/)
+          load_library_for_require(md[1])
+        end
+      end
+    end
+
+    def load_library_for_require(require_name)
+      lib_name = require_name.tr("/", "-")
+      return if @rbs_loader.libs.any? { |l| l.name == lib_name }
+
+      prev_libs = @rbs_loader.libs.dup
+      begin
+        @rbs_loader.add(library: lib_name, version: nil)
+        new_libs = @rbs_loader.libs - prev_libs
+        @rbs_loader.each_signature do |source, path, _buffer, _decls, _dirs|
+          next unless source.is_a?(RBS::EnvironmentLoader::Library)
+          next unless new_libs.include?(source)
+          next if @rbs_text_nodes.key?(path.to_s)
+          update_rbs_file(path.to_s, path.read)
+        end
+      rescue RBS::EnvironmentLoader::UnknownLibraryError, RuntimeError, Gem::LoadError, LoadError
+        @rbs_loader.libs.replace(prev_libs)
+      end
+    end
+
     def update_rbs_file(path, code)
       prev_decls = @rbs_text_nodes[path]
 
@@ -386,6 +413,76 @@ def code_lens(path)
       end
     end
 
+    def each_const_completion(path, pos, &blk)
+      const_node = find_constant_at(path, pos) || (pos.column > 0 ? find_constant_at(path, pos.left) : nil)
+      return unless const_node
+
+      prefix = const_node.cname.to_s
+      seen = ::Set.new
+
+      if const_node.cbase
+        parent_cpath = static_cpath_of(const_node.cbase)
+        return unless parent_cpath
+        mod = @genv.resolve_cpath(parent_cpath)
+        if mod && mod.exist?
+          mod.consts.each do |cname, cdef|
+            next unless cdef.exist?
+            next unless cname.to_s.start_with?(prefix)
+            yield cname.to_s, nil if seen.add?(cname)
+          end
+        end
+        @constant_catalog&.each_match(parent_cpath, prefix) do |cname, require_name|
+          yield cname.to_s, require_name if seen.add?(cname)
+        end
+      else
+        each_visible_const_with_prefix(const_node, prefix) do |cname|
+          yield cname.to_s, nil if seen.add?(cname)
+        end
+        @constant_catalog&.each_match([], prefix) do |cname, require_name|
+          yield cname.to_s, require_name if seen.add?(cname)
+        end
+      end
+    end
+
+    def find_constant_at(path, pos)
+      result = nil
+      @rb_text_nodes[path]&.retrieve_at(pos) do |node|
+        next if result
+        next unless node.is_a?(AST::ConstantReadNode)
+        cr = node.cname_code_range
+        result = node if cr.first <= pos && pos <= cr.last
+      end
+      result
+    end
+
+    def static_cpath_of(node)
+      parts = []
+      current = node
+      while current.is_a?(AST::ConstantReadNode)
+        parts.unshift(current.cname)
+        current = current.cbase
+      end
+      current.nil? ? parts : nil
+    end
+
+    def each_visible_const_with_prefix(const_node, prefix)
+      cref = const_node.lenv.cref
+      search_ancestors = !const_node.strict_const_scope
+      while cref
+        mod = @genv.resolve_cpath(cref.cpath)
+        @genv.each_superclass(mod, false) do |m, _singleton|
+          break if m == @genv.mod_object && cref.outer
+          m.consts.each do |cname, cdef|
+            next unless cdef.exist?
+            yield cname if cname.to_s.start_with?(prefix)
+          end
+          break unless search_ancestors
+        end
+        search_ancestors = false
+        cref = cref.outer
+      end
+    end
+
     def completion(path, trigger, pos)
       @rb_text_nodes[path]&.retrieve_at(pos) do |node|
         if node.code_range.last == pos.right

lib/typeprof/lsp.rb

@@ -2,6 +2,7 @@
 require "json"
 
 require_relative "lsp/text"
+require_relative "lsp/constant_catalog"
 require_relative "lsp/messages"
 require_relative "lsp/server"
 require_relative "lsp/util"

lib/typeprof/lsp/constant_catalog.rb

@@ -0,0 +1,95 @@
+module TypeProf::LSP
+  # Lightweight catalog of constants discovered in stdlib + rbs_collection RBS,
+  # without loading them into the type environment. Used by completion to suggest
+  # constants that need a `require` to be available, paired with the require name
+  # so the editor can auto-insert the require line via additionalTextEdits.
+  class ConstantCatalog
+    def initialize(rbs_collection: nil)
+      @entries = {}
+      build_from_stdlib
+      build_from_collection(rbs_collection) if rbs_collection
+    end
+
+    def each_match(parent_cpath, prefix)
+      @entries.each do |cpath, require_name|
+        next unless cpath.size == parent_cpath.size + 1
+        next unless cpath[0...-1] == parent_cpath
+        next unless cpath.last.to_s.start_with?(prefix)
+        yield cpath.last, require_name
+      end
+    end
+
+    def require_name_for(cpath)
+      @entries[cpath]
+    end
+
+    private
+
+    def build_from_stdlib
+      stdlib_root = Pathname(RBS::Repository::DEFAULT_STDLIB_ROOT)
+      return unless stdlib_root.directory?
+      stdlib_root.each_child do |lib_dir|
+        next unless lib_dir.directory?
+        register_dir(lib_dir, resolve_require_name(lib_dir.basename.to_s))
+      end
+    end
+
+    def build_from_collection(lockfile)
+      loader = RBS::EnvironmentLoader.new(core_root: nil)
+      loader.add_collection(lockfile)
+      loader.each_signature do |source, path, _buffer, _decls, _dirs|
+        next unless source.is_a?(RBS::EnvironmentLoader::Library)
+        register_file(path, resolve_require_name(source.name))
+      end
+    end
+
+    def resolve_require_name(gem_name)
+      # 1. dirname as-is — e.g. `open-uri`
+      return gem_name if Gem.find_files(gem_name).any?
+      # 2. dash → slash — e.g. `net-http` → `net/http`
+      slash_form = gem_name.tr("-", "/")
+      return slash_form if Gem.find_files(slash_form).any?
+      # 3. fallback (gem not installed in the LSP process); slash form is the
+      # more common require convention so use it as a best guess.
+      slash_form
+    end
+
+    def register_dir(dir, require_name)
+      Pathname.glob(dir + "**/*.rbs").each do |path|
+        register_file(path, require_name)
+      end
+    end
+
+    def register_file(path, require_name)
+      content = path.read
+      buf = RBS::Buffer.new(name: path.to_s, content: content)
+      _, _, decls = RBS::Parser.parse_signature(buf)
+      walk_decls(decls, [], require_name)
+    rescue StandardError, RBS::ParsingError
+      # Skip files that fail to parse; the catalog is best-effort.
+    end
+
+    def walk_decls(decls, prefix, require_name)
+      decls.each do |d|
+        case d
+        when RBS::AST::Declarations::Class, RBS::AST::Declarations::Module, RBS::AST::Declarations::Interface
+          cpath = compute_cpath(prefix, d.name)
+          @entries[cpath] ||= require_name
+          walk_decls(d.members.grep(RBS::AST::Declarations::Base), cpath, require_name)
+        when RBS::AST::Declarations::Constant
+          cpath = compute_cpath(prefix, d.name)
+          @entries[cpath] ||= require_name
+        end
+      end
+    end
+
+    def compute_cpath(prefix, type_name)
+      ns_path = type_name.namespace.path
+      if type_name.namespace.absolute?
+        ns_path + [type_name.name]
+      else
+        prefix + ns_path + [type_name.name]
+      end
+    end
+  end
+end

lib/typeprof/lsp/messages.rb

@@ -327,6 +327,7 @@ def run
       end
       items = []
       sort = "aaaa"
+      original_pos = pos
       text.modify_for_completion(text, pos) do |string, trigger, pos|
         @server.update_file(text.path, string)
         pos = TypeProf::CodePosition.from_lsp(pos)
@@ -339,13 +340,68 @@ def run
           }
           sort = sort.succ
         end
+
+        next unless trigger.nil?
+
+        const_sort = "z000"
+        @server.each_const_completion(text.path, TypeProf::CodePosition.from_lsp(original_pos)) do |label, require_name|
+          item = {
+            label: label,
+            kind: 21, # Constant
+            sortText: const_sort,
+          }
+          if require_name
+            insert = compute_require_insert(text.string, require_name)
+            edits = insert && build_require_text_edits(require_name, insert, original_pos)
+            if edits
+              item[:detail] = "from '#{ require_name }'"
+              item[:additionalTextEdits] = edits
+            end
+          end
+          items << item
+          const_sort = const_sort.succ
+        end
       end
       respond(
         isIncomplete: false,
         items: items,
       )
       @server.update_file(text.path, text.string)
     end
+
+    def compute_require_insert(source, require_name)
+      lines = source.lines
+      insert_line = 0
+      lines.each_with_index do |line, i|
+        s = line.strip
+        if (i == 0 && s.start_with?("#!")) || s.match?(/\A\s*#\s*(?:encoding|frozen_string_literal|coding|warn_indent):/)
+          insert_line = i + 1
+        else
+          break
+        end
+      end
+      while (line = lines[insert_line])
+        md = line.match(/\A\s*require(?:_relative)?\s*\(?\s*['"]([^'"]+)['"]/)
+        break unless md
+        return nil if md[1] == require_name
+        insert_line += 1
+      end
+      eol = source.include?("\r\n") ? "\r\n" : "\n"
+      next_line = lines[insert_line]
+      suffix = next_line.nil? || next_line.strip.empty? ? "" : eol
+      { line: insert_line, eol: eol, suffix: suffix }
+    end
+
+    # nil when the insert position would overlap with the cursor (LSP forbids
+    # overlap between additionalTextEdits and the primary edit at the cursor).
+    def build_require_text_edits(require_name, require_insert, completion_pos)
+      return nil if require_insert[:line] == completion_pos[:line] && completion_pos[:character] == 0
+      pos = { line: require_insert[:line], character: 0 }
+      [{
+        range: { start: pos, end: pos },
+        newText: "require '#{ require_name }'#{ require_insert[:eol] }#{ require_insert[:suffix] }",
+      }]
+    end
   end
 
   # textDocument/signatureHelp request