Merge pull request #874 from Shopify/vs/add_indexer

Add indexer

.rubocop.yml

@@ -32,6 +32,7 @@ Sorbet/TrueSigil:
   Enabled: true
   Include:
     - "test/**/*.rb"
+    - "lib/ruby_indexer/test/**/*.rb"
   Exclude:
     - "**/*.rake"
     - "lib/**/*.rb"
@@ -43,6 +44,7 @@ Sorbet/StrictSigil:
   Exclude:
     - "**/*.rake"
     - "test/**/*.rb"
+    - "lib/ruby_indexer/test/**/*.rb"
     - "lib/ruby-lsp.rb"
 
 Style/StderrPuts:

Gemfile.lock

@@ -5,6 +5,7 @@ PATH
       language_server-protocol (~> 3.17.0)
       sorbet-runtime
       syntax_tree (>= 6.1.1, < 7)
+      yarp (~> 0.6.0)
 
 GEM
   remote: https://rubygems.org/
@@ -167,6 +168,7 @@ GEM
     yard-sorbet (0.8.1)
       sorbet-runtime (>= 0.5)
       yard (>= 0.9)
+    yarp (0.6.0)
     zeitwerk (2.6.8)
 
 PLATFORMS

Rakefile

@@ -8,7 +8,7 @@ require "ruby_lsp/check_docs"
 Rake::TestTask.new(:test) do |t|
   t.libs << "test"
   t.libs << "lib"
-  t.test_files = FileList["test/**/*_test.rb"]
+  t.test_files = FileList["test/**/*_test.rb", "lib/ruby_indexer/test/**/*_test.rb"]
 end
 
 RDoc::Task.new do |rdoc|

lib/ruby_indexer/lib/ruby_indexer/configuration.rb

@@ -0,0 +1,20 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyIndexer
+  class Configuration
+    extend T::Sig
+
+    sig { returns(T::Array[String]) }
+    attr_accessor :files_to_index
+
+    sig { void }
+    def initialize
+      files_to_index = $LOAD_PATH.flat_map { |p| Dir.glob("#{p}/**/*.rb", base: p) }
+      files_to_index.concat(Dir.glob("#{Dir.pwd}/**/*.rb"))
+      files_to_index.reject! { |path| path.end_with?("_test.rb") }
+
+      @files_to_index = T.let(files_to_index, T::Array[String])
+    end
+  end
+end

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -0,0 +1,115 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyIndexer
+  class Index
+    extend T::Sig
+
+    sig { returns(T::Hash[String, T::Array[Entry]]) }
+    attr_reader :entries
+
+    sig { void }
+    def initialize
+      # Holds all entries in the index using the following format:
+      # {
+      #  "Foo" => [#<Entry::Class>, #<Entry::Class>],
+      #  "Foo::Bar" => [#<Entry::Class>],
+      # }
+      @entries = T.let({}, T::Hash[String, T::Array[Entry]])
+
+      # Holds references to where entries where discovered so that we can easily delete them
+      # {
+      #  "/my/project/foo.rb" => [#<Entry::Class>, #<Entry::Class>],
+      #  "/my/project/bar.rb" => [#<Entry::Class>],
+      # }
+      @files_to_entries = T.let({}, T::Hash[String, T::Array[Entry]])
+    end
+
+    sig { params(path: String).void }
+    def delete(path)
+      # For each constant discovered in `path`, delete the associated entry from the index. If there are no entries
+      # left, delete the constant from the index.
+      @files_to_entries[path]&.each do |entry|
+        entries = @entries[entry.name]
+        next unless entries
+
+        # Delete the specific entry from the list for this name
+        entries.delete(entry)
+        # If all entries were deleted, then remove the name from the hash
+        @entries.delete(entry.name) if entries.empty?
+      end
+
+      @files_to_entries.delete(path)
+    end
+
+    sig { params(entry: Entry).void }
+    def <<(entry)
+      (@entries[entry.name] ||= []) << entry
+      (@files_to_entries[entry.file_path] ||= []) << entry
+    end
+
+    sig { params(fully_qualified_name: String).returns(T.nilable(T::Array[Entry])) }
+    def [](fully_qualified_name)
+      @entries[fully_qualified_name.delete_prefix("::")]
+    end
+
+    # Try to find the entry based on the nesting from the most specific to the least specific. For example, if we have
+    # the nesting as ["Foo", "Bar"] and the name as "Baz", we will try to find it in this order:
+    # 1. Foo::Bar::Baz
+    # 2. Foo::Baz
+    # 3. Baz
+    sig { params(name: String, nesting: T::Array[String]).returns(T.nilable(T::Array[Entry])) }
+    def resolve(name, nesting)
+      (nesting.length + 1).downto(0).each do |i|
+        prefix = T.must(nesting[0...i]).join("::")
+        full_name = prefix.empty? ? name : "#{prefix}::#{name}"
+        entries = @entries[full_name]
+        return entries if entries
+      end
+
+      nil
+    end
+
+    sig { void }
+    def clear
+      @entries.clear
+    end
+
+    class Entry
+      extend T::Sig
+
+      sig { returns(String) }
+      attr_reader :name
+
+      sig { returns(String) }
+      attr_reader :file_path
+
+      sig { returns(YARP::Location) }
+      attr_reader :location
+
+      sig { returns(T::Array[String]) }
+      attr_reader :comments
+
+      sig { params(name: String, file_path: String, location: YARP::Location, comments: T::Array[String]).void }
+      def initialize(name, file_path, location, comments)
+        @name = name
+        @file_path = file_path
+        @location = location
+        @comments = comments
+      end
+
+      class Namespace < Entry
+        sig { returns(String) }
+        def short_name
+          T.must(@name.split("::").last)
+        end
+      end
+
+      class Module < Namespace
+      end
+
+      class Class < Namespace
+      end
+    end
+  end
+end

lib/ruby_indexer/lib/ruby_indexer/visitor.rb

@@ -0,0 +1,91 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyIndexer
+  class IndexVisitor < YARP::Visitor
+    extend T::Sig
+
+    sig { params(index: Index, parse_result: YARP::ParseResult, file_path: String).void }
+    def initialize(index, parse_result, file_path)
+      @index = index
+      @parse_result = parse_result
+      @file_path = file_path
+      @stack = T.let([], T::Array[String])
+      @comments_by_line = T.let(
+        parse_result.comments.to_h do |c|
+          [c.location.start_line, c]
+        end,
+        T::Hash[Integer, YARP::Comment],
+      )
+
+      super()
+    end
+
+    sig { void }
+    def run
+      visit(@parse_result.value)
+    end
+
+    sig { params(node: T.nilable(YARP::Node)).void }
+    def visit(node)
+      case node
+      when YARP::ProgramNode, YARP::StatementsNode
+        visit_child_nodes(node)
+      when YARP::ClassNode
+        add_index_entry(node, Index::Entry::Class)
+      when YARP::ModuleNode
+        add_index_entry(node, Index::Entry::Module)
+      end
+    end
+
+    # Override to avoid using `map` instead of `each`
+    sig { params(nodes: T::Array[T.nilable(YARP::Node)]).void }
+    def visit_all(nodes)
+      nodes.each { |node| visit(node) }
+    end
+
+    private
+
+    sig { params(node: T.any(YARP::ClassNode, YARP::ModuleNode), klass: T.class_of(Index::Entry)).void }
+    def add_index_entry(node, klass)
+      name = node.constant_path.location.slice
+
+      unless /^[A-Z:]/.match?(name)
+        return visit_child_nodes(node)
+      end
+
+      fully_qualified_name = name.start_with?("::") ? name : fully_qualify_name(name)
+      name.delete_prefix!("::")
+
+      comments = collect_comments(node)
+      @index << klass.new(fully_qualified_name, @file_path, node.location, comments)
+      @stack << name
+      visit_child_nodes(node)
+      @stack.pop
+    end
+
+    sig { params(node: YARP::Node).returns(T::Array[String]) }
+    def collect_comments(node)
+      comments = []
+
+      start_line = node.location.start_line - 1
+      start_line -= 1 unless @comments_by_line.key?(start_line)
+
+      start_line.downto(1) do |line|
+        comment = @comments_by_line[line]
+        break unless comment
+
+        comments.unshift(comment.location.slice)
+      end
+
+      comments
+    end
+
+    sig { params(name: String).returns(String) }
+    def fully_qualify_name(name)
+      return name if @stack.empty?
+
+      "#{@stack.join("::")}::#{name}"
+    end
+  end
+end

lib/ruby_indexer/ruby_indexer.rb

@@ -0,0 +1,36 @@
+# typed: strict
+# frozen_string_literal: true
+
+require "ruby_indexer/lib/ruby_indexer/visitor"
+require "ruby_indexer/lib/ruby_indexer/index"
+require "ruby_indexer/lib/ruby_indexer/configuration"
+
+module RubyIndexer
+  class << self
+    extend T::Sig
+
+    sig { params(block: T.proc.params(configuration: Configuration).void).void }
+    def configure(&block)
+      block.call(configuration)
+    end
+
+    sig { returns(Configuration) }
+    def configuration
+      @configuration ||= T.let(Configuration.new, T.nilable(Configuration))
+    end
+
+    sig { params(paths: T::Array[String]).returns(Index) }
+    def index(paths = configuration.files_to_index)
+      index = Index.new
+      paths.each { |path| index_single(index, path) }
+      index
+    end
+
+    sig { params(index: Index, path: String, source: T.nilable(String)).void }
+    def index_single(index, path, source = nil)
+      content = source || File.read(path)
+      visitor = IndexVisitor.new(index, YARP.parse(content), path)
+      visitor.run
+    end
+  end
+end