RBS scan cache to improve overall performance

lib/rdoc/parser/rbs.rb

@@ -11,6 +11,8 @@ class RDoc::Parser::RBS < RDoc::Parser
   parse_files_matching RBS_FILE_EXTENSION
 
   def scan
+    @attributes_by_context = {}
+    @methods_by_context = {}
     _, _, decls = ::RBS::Parser.parse_signature(@content)
     decls.each do |decl|
       parse_decl decl, @top_level
@@ -90,12 +92,28 @@ def attr_rw_matches?(existing_rw, new_rw)
     existing_rw.each_char.any? { |rw| new_rw.include? rw }
   end
 
+  def attribute_index(context)
+    index = @attributes_by_context[context] ||= {}
+    context.attributes[index.length..].each do |attribute|
+      index[[attribute.name, attribute.singleton]] ||= attribute
+    end
+    index
+  end
+
+  def method_index(context)
+    index = @methods_by_context[context] ||= {}
+    context.method_list[index.length..].each do |method|
+      index[[method.name, !!method.singleton]] ||= method
+    end
+    index
+  end
+
   def merge_attribute_methods(context, name, rw, singleton, comment, type_signature_lines)
     method_names = []
     method_names << name if rw.include?('R')
     method_names << "#{name}=" if rw.include?('W')
 
-    methods = method_names.map { |method_name| context.find_method(method_name, singleton) }
+    methods = method_names.map { |method_name| method_index(context)[[method_name, singleton]] }
     methods.compact.each do |method|
       merge_documentation method, comment, type_signature_lines
     end
@@ -135,7 +153,7 @@ def parse_attr_decl(decl, context)
     type_signature_lines = [decl.type.to_s]
     name = decl.name.to_s
     singleton = decl.kind == :singleton
-    if attribute = context.find_attribute(name, singleton)
+    if attribute = attribute_index(context)[[name, singleton]]
       merge_documentation attribute, comment, type_signature_lines if
         attr_rw_matches? attribute.rw, rw
       return
@@ -245,7 +263,7 @@ def parse_method_decl(decl, context)
     singleton = rdoc_method_singleton?(decl)
     visibility = rdoc_method_visibility(decl)
 
-    if method = context.find_method(method_name, singleton)
+    if method = method_index(context)[[method_name, singleton]]
       merge_documentation method, comment, type_signature_lines
       return
     end

test/rdoc/parser/rbs_test.rb

@@ -143,6 +143,70 @@ def greet: () -> String
     assert_equal ['() -> String'], greet.type_signature_lines
   end
 
+  def test_scan_finds_existing_method_after_store_index_rebuild
+    ruby_top_level = @store.add_file 'sample.rb'
+    sample = ruby_top_level.add_class RDoc::NormalClass, 'Sample'
+
+    name = RDoc::AnyMethod.new 'name'
+    name.record_location ruby_top_level
+    sample.add_method name
+
+    name_writer = RDoc::Attr.new 'name', 'W', nil
+    name_writer.record_location ruby_top_level
+    sample.add_attribute name_writer
+
+    util_parser("class Sample\nend\n").scan
+    @store.clear_file_contributions @filename, keep_position: true
+    util_parser(<<~RBS).scan
+      class Sample
+        def name: () -> String
+      end
+    RBS
+
+    assert_equal ['() -> String'], name.type_signature_lines
+  end
+
+  def test_scan_indexes_forward_method_and_attribute_aliases
+    util_parser(<<~RBS).scan
+      class Sample
+        # Method alias docs.
+        alias salutation greet
+        alias display_name name
+
+        def greet: () -> String
+        # Base attribute docs.
+        attr_reader name: String
+
+        # Dedicated method docs.
+        def salutation: () -> String
+        # Dedicated attribute docs.
+        attr_reader display_name: String
+      end
+    RBS
+
+    sample = @store.find_class_named 'Sample'
+    salutation = sample.find_method 'salutation', false
+    display_name = sample.find_attribute 'display_name', false
+
+    assert_equal "Method alias docs.\n---\nDedicated method docs.", salutation.comment.to_s.strip
+    assert_equal "Base attribute docs.\n---\nDedicated attribute docs.", display_name.comment.to_s.strip
+  end
+
+  def test_scan_treats_legacy_nil_singleton_as_instance_method
+    ruby_top_level = @store.add_file 'sample.rb'
+    sample = ruby_top_level.add_class RDoc::NormalClass, 'Sample'
+    greet = RDoc::AnyMethod.new 'greet', singleton: nil
+    sample.add_method greet
+
+    util_parser(<<~RBS).scan
+      class Sample
+        def greet: () -> String
+      end
+    RBS
+
+    assert_equal ['() -> String'], greet.type_signature_lines
+  end
+
   def test_scan_preserves_rbs_markdown_when_extending_method_documentation
     ruby_top_level = @store.add_file 'sample.rb'
     sample = ruby_top_level.add_class RDoc::NormalClass, 'Sample'
@@ -266,6 +330,13 @@ class PrivateSample
     assert_equal :public, private_constructor.visibility
   end
 
+  def test_scan_member_lookup_linear_performance
+    assert_linear_performance([1, 10, 100]) do |factor|
+      members = Array.new(factor * 200) { |i| "  def m#{i}: () -> void\n  attr_reader a#{i}: String" }.join("\n")
+      util_parser("class C\n#{members}\nend\n").scan
+    end
+  end
+
   def util_parser(content)
     RDoc::Parser::RBS.new @top_level, content, @options, @stats
   end