Allow indexing enhancements to create namespaces (#2857)

### Motivation

We realized that in order to support concerns `class_methods do...end`, indexing enhancements needed to be more powerful and have the liberty to create fake/temporary namespaces.

When trying to make that possible, we also realized that the API had many shortcomings that made enhancements harder than they had to be.

### Implementation

The idea is to instantiate enhancements with the declaration listener and then expose the API from it. That way, the declaration listener can hold important state like the code units cache and the current file path and enhancements only make flow adjustments and additions to the indexing process.

This PR also allows enhancements to create modules and classes, which allows us to support `class_methods do...end` in the Rails add-on and allows the RSpec add-on to support `let` and `subject` properly.

### Automated Tests

Added tests.

Author: vinistock

jekyll/add-ons.markdown

@@ -271,15 +271,11 @@ This is how you could write an enhancement to teach the Ruby LSP to understand t
 class MyIndexingEnhancement < RubyIndexer::Enhancement
   # This on call node handler is invoked any time during indexing when we find a method call. It can be used to insert
   # more entries into the index depending on the conditions
-  def on_call_node_enter(owner, node, file_path, code_units_cache)
-    return unless owner
+  def on_call_node_enter(node)
+    return unless @listener.current_owner
 
-    # Get the ancestors of the current class
-    ancestors = @index.linearized_ancestors_of(owner.name)
-
-    # Return early unless the method call is the one we want to handle and the class invoking the DSL inherits from
-    # our library's parent class
-    return unless node.name == :my_dsl_that_creates_methods && ancestors.include?("MyLibrary::ParentClass")
+    # Return early unless the method call is the one we want to handle
+    return unless node.name == :my_dsl_that_creates_methods
 
     # Create a new entry to be inserted in the index. This entry will represent the declaration that is created via
     # meta-programming. All entries are defined in the `entry.rb` file.
@@ -293,24 +289,16 @@ class MyIndexingEnhancement < RubyIndexer::Enhancement
       RubyIndexer::Entry::Signature.new([RubyIndexer::Entry::RequiredParameter.new(name: :a)])
     ]
 
-    new_entry = RubyIndexer::Entry::Method.new(
-      "new_method", # The name of the method that gets created via meta-programming
-      file_path,    # The file_path where the DSL call was found. This should always just be the file_path received
-      location,     # The Prism node location where the DSL call was found
-      location,     # The Prism node location for the DSL name location. May or not be the same
-      nil,          # The documentation for this DSL call. This should always be `nil` to ensure lazy fetching of docs
-      signatures,   # All signatures for this method (every way it can be invoked)
-      RubyIndexer::Entry::Visibility::PUBLIC, # The method's visibility
-      owner,        # The method's owner. This is almost always going to be the same owner received
+    @listener.add_method(
+      "new_method", # Name of the method
+      location,     # Prism location for the node defining this method
+      signatures    # Signatures available to invoke this method
     )
-
-    # Push the new entry to the index
-    @index.add(new_entry)
   end
 
   # This method is invoked when the parser has finished processing the method call node.
   # It can be used to perform cleanups like popping a stack...etc.
-  def on_call_node_leave(owner, node, file_path, code_units_cache); end
+  def on_call_node_leave(node); end
 end
 ```
 

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -18,13 +18,12 @@ class DeclarationListener
         parse_result: Prism::ParseResult,
         file_path: String,
         collect_comments: T::Boolean,
-        enhancements: T::Array[Enhancement],
       ).void
     end
-    def initialize(index, dispatcher, parse_result, file_path, collect_comments: false, enhancements: [])
+    def initialize(index, dispatcher, parse_result, file_path, collect_comments: false)
       @index = index
       @file_path = file_path
-      @enhancements = enhancements
+      @enhancements = T.let(Enhancement.all(self), T::Array[Enhancement])
       @visibility_stack = T.let([Entry::Visibility::PUBLIC], T::Array[Entry::Visibility])
       @comments_by_line = T.let(
         parse_result.comments.to_h do |c|
@@ -86,15 +85,9 @@ def initialize(index, dispatcher, parse_result, file_path, collect_comments: fal
 
     sig { params(node: Prism::ClassNode).void }
     def on_class_node_enter(node)
-      @visibility_stack.push(Entry::Visibility::PUBLIC)
       constant_path = node.constant_path
-      name = constant_path.slice
-
-      comments = collect_comments(node)
-
       superclass = node.superclass
-
-      nesting = actual_nesting(name)
+      nesting = actual_nesting(constant_path.slice)
 
       parent_class = case superclass
       when Prism::ConstantReadNode, Prism::ConstantPathNode
@@ -113,53 +106,29 @@ def on_class_node_enter(node)
         end
       end
 
-      entry = Entry::Class.new(
+      add_class(
         nesting,
-        @file_path,
-        Location.from_prism_location(node.location, @code_units_cache),
-        Location.from_prism_location(constant_path.location, @code_units_cache),
-        comments,
-        parent_class,
+        node.location,
+        constant_path.location,
+        parent_class_name: parent_class,
+        comments: collect_comments(node),
       )
-
-      @owner_stack << entry
-      @index.add(entry)
-      @stack << name
     end
 
     sig { params(node: Prism::ClassNode).void }
     def on_class_node_leave(node)
-      @stack.pop
-      @owner_stack.pop
-      @visibility_stack.pop
+      pop_namespace_stack
     end
 
     sig { params(node: Prism::ModuleNode).void }
     def on_module_node_enter(node)
-      @visibility_stack.push(Entry::Visibility::PUBLIC)
       constant_path = node.constant_path
-      name = constant_path.slice
-
-      comments = collect_comments(node)
-
-      entry = Entry::Module.new(
-        actual_nesting(name),
-        @file_path,
-        Location.from_prism_location(node.location, @code_units_cache),
-        Location.from_prism_location(constant_path.location, @code_units_cache),
-        comments,
-      )
-
-      @owner_stack << entry
-      @index.add(entry)
-      @stack << name
+      add_module(constant_path.slice, node.location, constant_path.location, comments: collect_comments(node))
     end
 
     sig { params(node: Prism::ModuleNode).void }
     def on_module_node_leave(node)
-      @stack.pop
-      @owner_stack.pop
-      @visibility_stack.pop
+      pop_namespace_stack
     end
 
     sig { params(node: Prism::SingletonClassNode).void }
@@ -201,9 +170,7 @@ def on_singleton_class_node_enter(node)
 
     sig { params(node: Prism::SingletonClassNode).void }
     def on_singleton_class_node_leave(node)
-      @stack.pop
-      @owner_stack.pop
-      @visibility_stack.pop
+      pop_namespace_stack
     end
 
     sig { params(node: Prism::MultiWriteNode).void }
@@ -318,7 +285,7 @@ def on_call_node_enter(node)
       end
 
       @enhancements.each do |enhancement|
-        enhancement.on_call_node_enter(@owner_stack.last, node, @file_path, @code_units_cache)
+        enhancement.on_call_node_enter(node)
       rescue StandardError => e
         @indexing_errors << <<~MSG
           Indexing error in #{@file_path} with '#{enhancement.class.name}' on call node enter enhancement: #{e.message}
@@ -339,7 +306,7 @@ def on_call_node_leave(node)
       end
 
       @enhancements.each do |enhancement|
-        enhancement.on_call_node_leave(@owner_stack.last, node, @file_path, @code_units_cache)
+        enhancement.on_call_node_leave(node)
       rescue StandardError => e
         @indexing_errors << <<~MSG
           Indexing error in #{@file_path} with '#{enhancement.class.name}' on call node leave enhancement: #{e.message}
@@ -464,6 +431,98 @@ def on_alias_method_node_enter(node)
       )
     end
 
+    sig do
+      params(
+        name: String,
+        node_location: Prism::Location,
+        signatures: T::Array[Entry::Signature],
+        visibility: Entry::Visibility,
+        comments: T.nilable(String),
+      ).void
+    end
+    def add_method(name, node_location, signatures, visibility: Entry::Visibility::PUBLIC, comments: nil)
+      location = Location.from_prism_location(node_location, @code_units_cache)
+
+      @index.add(Entry::Method.new(
+        name,
+        @file_path,
+        location,
+        location,
+        comments,
+        signatures,
+        visibility,
+        @owner_stack.last,
+      ))
+    end
+
+    sig do
+      params(
+        name: String,
+        full_location: Prism::Location,
+        name_location: Prism::Location,
+        comments: T.nilable(String),
+      ).void
+    end
+    def add_module(name, full_location, name_location, comments: nil)
+      location = Location.from_prism_location(full_location, @code_units_cache)
+      name_loc = Location.from_prism_location(name_location, @code_units_cache)
+
+      entry = Entry::Module.new(
+        actual_nesting(name),
+        @file_path,
+        location,
+        name_loc,
+        comments,
+      )
+
+      advance_namespace_stack(name, entry)
+    end
+
+    sig do
+      params(
+        name_or_nesting: T.any(String, T::Array[String]),
+        full_location: Prism::Location,
+        name_location: Prism::Location,
+        parent_class_name: T.nilable(String),
+        comments: T.nilable(String),
+      ).void
+    end
+    def add_class(name_or_nesting, full_location, name_location, parent_class_name: nil, comments: nil)
+      nesting = name_or_nesting.is_a?(Array) ? name_or_nesting : actual_nesting(name_or_nesting)
+      entry = Entry::Class.new(
+        nesting,
+        @file_path,
+        Location.from_prism_location(full_location, @code_units_cache),
+        Location.from_prism_location(name_location, @code_units_cache),
+        comments,
+        parent_class_name,
+      )
+
+      advance_namespace_stack(T.must(nesting.last), entry)
+    end
+
+    sig { params(block: T.proc.params(index: Index, base: Entry::Namespace).void).void }
+    def register_included_hook(&block)
+      owner = @owner_stack.last
+      return unless owner
+
+      @index.register_included_hook(owner.name) do |index, base|
+        block.call(index, base)
+      end
+    end
+
+    sig { void }
+    def pop_namespace_stack
+      @stack.pop
+      @owner_stack.pop
+      @visibility_stack.pop
+    end
+
+    sig { returns(T.nilable(Entry::Namespace)) }
+    def current_owner
+      @owner_stack.last
+    end
+
     private
 
     sig do
@@ -921,5 +980,13 @@ def actual_nesting(name)
 
       corrected_nesting
     end
+
+    sig { params(short_name: String, entry: Entry::Namespace).void }
+    def advance_namespace_stack(short_name, entry)
+      @visibility_stack.push(Entry::Visibility::PUBLIC)
+      @owner_stack << entry
+      @index.add(entry)
+      @stack << short_name
+    end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/enhancement.rb

@@ -8,38 +8,41 @@ class Enhancement
 
     abstract!
 
-    sig { params(index: Index).void }
-    def initialize(index)
-      @index = index
+    @enhancements = T.let([], T::Array[T::Class[Enhancement]])
+
+    class << self
+      extend T::Sig
+
+      sig { params(child: T::Class[Enhancement]).void }
+      def inherited(child)
+        @enhancements << child
+        super
+      end
+
+      sig { params(listener: DeclarationListener).returns(T::Array[Enhancement]) }
+      def all(listener)
+        @enhancements.map { |enhancement| enhancement.new(listener) }
+      end
+
+      # Only available for testing purposes
+      sig { void }
+      def clear
+        @enhancements.clear
+      end
+    end
+
+    sig { params(listener: DeclarationListener).void }
+    def initialize(listener)
+      @listener = listener
     end
 
     # The `on_extend` indexing enhancement is invoked whenever an extend is encountered in the code. It can be used to
     # register for an included callback, similar to what `ActiveSupport::Concern` does in order to auto-extend the
     # `ClassMethods` modules
-    sig do
-      overridable.params(
-        owner: T.nilable(Entry::Namespace),
-        node: Prism::CallNode,
-        file_path: String,
-        code_units_cache: T.any(
-          T.proc.params(arg0: Integer).returns(Integer),
-          Prism::CodeUnitsCache,
-        ),
-      ).void
-    end
-    def on_call_node_enter(owner, node, file_path, code_units_cache); end
-
-    sig do
-      overridable.params(
-        owner: T.nilable(Entry::Namespace),
-        node: Prism::CallNode,
-        file_path: String,
-        code_units_cache: T.any(
-          T.proc.params(arg0: Integer).returns(Integer),
-          Prism::CodeUnitsCache,
-        ),
-      ).void
-    end
-    def on_call_node_leave(owner, node, file_path, code_units_cache); end
+    sig { overridable.params(node: Prism::CallNode).void }
+    def on_call_node_enter(node); end # rubocop:disable RubyLsp/UseRegisterWithHandlerMethod
+
+    sig { overridable.params(node: Prism::CallNode).void }
+    def on_call_node_leave(node); end # rubocop:disable RubyLsp/UseRegisterWithHandlerMethod
   end
 end

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -40,9 +40,6 @@ def initialize
       # Holds the linearized ancestors list for every namespace
       @ancestors = T.let({}, T::Hash[String, T::Array[String]])
 
-      # List of classes that are enhancing the index
-      @enhancements = T.let([], T::Array[Enhancement])
-
       # Map of module name to included hooks that have to be executed when we include the given module
       @included_hooks = T.let(
         {},
@@ -52,12 +49,6 @@ def initialize
       @configuration = T.let(RubyIndexer::Configuration.new, Configuration)
     end
 
-    # Register an enhancement to the index. Enhancements must conform to the `Enhancement` interface
-    sig { params(enhancement: Enhancement).void }
-    def register_enhancement(enhancement)
-      @enhancements << enhancement
-    end
-
     # Register an included `hook` that will be executed when `module_name` is included into any namespace
     sig { params(module_name: String, hook: T.proc.params(index: Index, base: Entry::Namespace).void).void }
     def register_included_hook(module_name, &hook)
@@ -396,7 +387,6 @@ def index_single(indexable_path, source = nil, collect_comments: true)
         result,
         indexable_path.full_path,
         collect_comments: collect_comments,
-        enhancements: @enhancements,
       )
       dispatcher.dispatch(result.value)