Support module_function in the indexer

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -312,6 +312,8 @@ def on_call_node_enter(node)
         @visibility_stack.push(Entry::Visibility::PROTECTED)
       when :private
         @visibility_stack.push(Entry::Visibility::PRIVATE)
+      when :module_function
+        handle_module_function(node)
       end
 
       @enhancements.each do |enhancement|
@@ -751,6 +753,48 @@ def handle_module_operation(node, operation)
       end
     end
 
+    sig { params(node: Prism::CallNode).void }
+    def handle_module_function(node)
+      arguments_node = node.arguments
+      return unless arguments_node
+
+      owner_name = @owner_stack.last&.name
+      return unless owner_name
+
+      arguments_node.arguments.each do |argument|
+        method_name = case argument
+        when Prism::StringNode
+          argument.content
+        when Prism::SymbolNode
+          argument.value
+        end
+        next unless method_name
+
+        entries = @index.resolve_method(method_name, owner_name)
+        next unless entries
+
+        entries.each do |entry|
+          entry_owner_name = entry.owner&.name
+          next unless entry_owner_name
+
+          entry.visibility = Entry::Visibility::PRIVATE
+
+          singleton = @index.existing_or_new_singleton_class(entry_owner_name)
+          location = Location.from_prism_location(argument.location, @code_units_cache)
+          @index.add(Entry::Method.new(
+            method_name,
+            @file_path,
+            location,
+            location,
+            collect_comments(node)&.concat(entry.comments),
+            entry.signatures,
+            Entry::Visibility::PUBLIC,
+            singleton,
+          ))
+        end
+      end
+    end
+
     sig { returns(Entry::Visibility) }
     def current_visibility
       T.must(@visibility_stack.last)

lib/ruby_indexer/test/method_test.rb

@@ -123,6 +123,32 @@ 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 Test
+          def foo; end
+          def bar; end
+          module_function :foo, "bar"
+        end
+      RUBY
+
+      ["foo", "bar"].each do |keyword|
+        entries = T.must(@index[keyword])
+        # should receive two entries because module_function creates a singleton method
+        # for the Test module and a private method for classes include the Test module
+        assert_equal(entries.size, 2)
+        first_entry, second_entry = *entries
+        # The first entry points to the location of the module_function call
+        assert_equal("Test", 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
+        assert_equal("Test::<Class:Test>", second_entry.owner.name)
+        assert_instance_of(Entry::SingletonClass, second_entry.owner)
+        assert_equal(Entry::Visibility::PUBLIC, second_entry.visibility)
+      end
+    end
+
     def test_method_with_parameters
       index(<<~RUBY)
         class Foo