From 4cc387fc412cba9aa173d788cd799e27400e613c Mon Sep 17 00:00:00 2001
From: Weilong <fan01856472@gmail.com>
Date: Thu, 17 Oct 2024 09:52:35 -0400
Subject: [PATCH] add module_function to on_call_node_enter

---
 .../lib/ruby_indexer/declaration_listener.rb  | 23 +++++++++++++++
 lib/ruby_indexer/test/method_test.rb          | 28 +++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
index d96adc04e..0f2fb04c1 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
@@ -312,6 +312,29 @@ def on_call_node_enter(node)
         @visibility_stack.push(Entry::Visibility::PROTECTED)
       when :private
         @visibility_stack.push(Entry::Visibility::PRIVATE)
+      when :module_function
+        name_argument = node.arguments.arguments.first
+        return unless name_argument
+
+        entries = @index.resolve_method(name_argument.value, @owner_stack.last.name)
+        return unless entries
+
+        entries.each do |e|
+          e.visibility = Entry::Visibility::PRIVATE
+
+          singleton = @index.existing_or_new_singleton_class(e.owner.name)
+
+          @index.add(Entry::Method.new(
+            name_argument.value,
+            @file_path,
+            Location.from_prism_location(node.location, @code_units_cache),
+            Location.from_prism_location(node.message_loc, @code_units_cache),
+            collect_comments(node),
+            [],
+            e.visibility,
+            singleton,
+          ))
+        end
       end
 
       @enhancements.each do |enhancement|
diff --git a/lib/ruby_indexer/test/method_test.rb b/lib/ruby_indexer/test/method_test.rb
index ca9bb1282..7a827fd5b 100644
--- a/lib/ruby_indexer/test/method_test.rb
+++ b/lib/ruby_indexer/test/method_test.rb
@@ -123,6 +123,34 @@ def baz; end
       assert_entry("baz", Entry::Method, "/fake/path/foo.rb:9-2:9-14", visibility: Entry::Visibility::PRIVATE)
     end
 
+    def test_visibility_tracking_with_module_function
+      index(<<~RUBY)
+        module Foo
+          def foo; end
+          module_function :foo
+        end
+
+        class Bar
+          include Foo
+        end
+      RUBY
+
+      entries = T.must(@index["foo"])
+      # receive two entries because module_function creates a singleton method
+      # for the Foo module and a private method for classes include Foo module
+      assert_equal(entries.size, 2)
+      # The first entry points to the location of the module_function call
+      first_entry = entries.first
+      assert_equal("Foo", first_entry.owner.name)
+      assert_instance_of(Entry::Module, first_entry.owner)
+      assert_equal(Entry::Visibility::PRIVATE, first_entry.visibility)
+      # The second entry points to the public singleton method
+      second_entry = entries.last
+      assert_equal("Foo::<Class:Foo>", second_entry.owner.name)
+      assert_instance_of(Entry::SingletonClass, second_entry.owner)
+      assert_equal(Entry::Visibility::PRIVATE, second_entry.visibility)
+    end
+
     def test_method_with_parameters
       index(<<~RUBY)
         class Foo