From ba3e595c8dd76f87b3bc19e96e1cffd5cff9db67 Mon Sep 17 00:00:00 2001
From: Stan Lo <stan.lo@shopify.com>
Date: Tue, 22 Oct 2024 20:37:05 +0100
Subject: [PATCH] Pass index to enhancement's constructor

---
 lib/ruby_indexer/lib/ruby_indexer/enhancement.rb | 11 +++++++----
 lib/ruby_indexer/test/enhancements_test.rb       |  8 ++++----
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb b/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb
index f0f11d390..a654b87e8 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb
@@ -8,12 +8,16 @@ class Enhancement
 
     abstract!
 
+    sig { params(index: Index).void }
+    def initialize(index)
+      @index = index
+    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(
-        index: Index,
         owner: T.nilable(Entry::Namespace),
         node: Prism::CallNode,
         file_path: String,
@@ -23,11 +27,10 @@ class Enhancement
         ),
       ).void
     end
-    def on_call_node_enter(index, owner, node, file_path, code_units_cache); end
+    def on_call_node_enter(owner, node, file_path, code_units_cache); end
 
     sig do
       overridable.params(
-        index: Index,
         owner: T.nilable(Entry::Namespace),
         node: Prism::CallNode,
         file_path: String,
@@ -37,6 +40,6 @@ def on_call_node_enter(index, owner, node, file_path, code_units_cache); end
         ),
       ).void
     end
-    def on_call_node_leave(index, owner, node, file_path, code_units_cache); end
+    def on_call_node_leave(owner, node, file_path, code_units_cache); end
   end
 end
diff --git a/lib/ruby_indexer/test/enhancements_test.rb b/lib/ruby_indexer/test/enhancements_test.rb
index e5dd106bd..aa83f4184 100644
--- a/lib/ruby_indexer/test/enhancements_test.rb
+++ b/lib/ruby_indexer/test/enhancements_test.rb
@@ -48,7 +48,7 @@ def on_call_node_enter(index, owner, node, file_path, code_units_cache)
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
       index(<<~RUBY)
         module ActiveSupport
           module Concern
@@ -124,7 +124,7 @@ def on_call_node_enter(index, owner, node, file_path, code_units_cache)
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
       index(<<~RUBY)
         module ActiveSupport
           module Concern
@@ -169,7 +169,7 @@ def name
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
 
       _stdout, stderr = capture_io do
         index(<<~RUBY)
@@ -204,7 +204,7 @@ def name
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
 
       _stdout, stderr = capture_io do
         index(<<~RUBY)