Allow indexing enhancements to create namespaces

Co-authored-by: Alex Rocha <[email protected]>

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)