Repository: Shopify/ruby-lsp (Ruby)

Issue:
Support `module_function` in the indexer
When `module_function` is invoked, the given method becomes available in the singleton class and it is turned into a private method on including namespaces. For example:

```ruby
module A
  def foo; end
  module_function :foo
end

A.foo # okay!

class Something
  include A
end

Something.new.foo # not okay, instance method foo is private
```

We need to add processing for `module_function` [here](https://github.com/Shopify/ruby-lsp/blob/ae70ebc6348febe6064789890c221944463597ba/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb#L291). Essentially, when we find an invocation to `module_function`, we want to:

1. Get the method entry that was already inserted in the index. Update the visibility to private
2. Add a new method entry for the same method, but in this case the owner is the singleton class


Discussion:
Hi team,

I am taking a look at this issue and will push something a few days later.
Hello! Sounds good, I have assigned the issue to you. If you hit any blockers or have any doubts, please do not hesitate to ask for assistance here and the team can discuss the implementation.
Hi @vinistock,  

Thank you for pointing out where the method entry should proceed. It's really helpful for me to start looking around the issue 😊.

For the second requirement, I am confused about what you mean by "the same method." Do you mean `on_call_node_enter` in `declaration_listener.rb`?

Also, I would like to double-check with you for the first requirement. I think the fix is to add a new entry point, module_function, to the `on_call_node_enter` and set the visibility as private (please see the following image). Please correct me if I am wrong! I am still attempting to understand what the visibility_stack is and how it works.

![Image](https://github.com/user-attachments/assets/7843e17a-5ecf-4b7a-9aa5-27c506ee12b1)

You're on the right track, but it's slightly different. When the indexing process finds the `module_function` invocation, we already finished processing the method declaration. So we don't want to change the visibility stack, we want to mutate the entry that was already inserted in the index.

For example

```ruby
module Foo
  # We process this method declaration and add an `Entry::Method` for it
  def something; end

  # After adding it to the index, then we find the method call to `module_function`
  module_function :something
end
```

So we need to operate on the entry that we just inserted. Here's the pseudo code of how this should happen

```ruby
case message
# ...
when :module_function
  # get the argument for module function from the node, which is the name of the method we are
  # turning into a module function

  # then find the existing entry for the method. It will be something like
  entries = @index.resolve_method(name_of_method, owner_name)

  entries.each do |e|
    # then turn the existing entries into private (this API may not exist yet)
    e.visibility = Entry::Visibility::PRIVATE

    # and add the singleton version of the method to the index (which is what module_function does)
    # here the important part is that this method's owner will not be the same owner, but the 
    # singleton version of it
    singleton_owner = @index.existing_or_new_singleton_class(name_of_owner)
    new_entry = Entry::Method.new(..., singleton_owner)
  end
end
```
@vinistock, thanks for the pseudo code. I am now writing a test, but I need to figure out how to test the singleton entry.

Right now, the following code is a piece of tests placed in method_test.rb.
```ruby
def test_visibility_tracking_with_module_function
  index(<<~RUBY)
    module Foo
      def foo; end
      module_function :foo
    end
    class Bar
      include Foo
    end
    Bar.new.foo
  RUBY

  assert_entry("foo", Entry::Method, "/fake/path/foo.rb:1-2:1-14", visibility: Entry::Visibility::PRIVATE)
  # How can I make another assertion to test Bar.new.foo?
end
```
You need to get a handle for the entry so that you can assert that the owner is correct ([example](https://github.com/Shopify/ruby-lsp/blob/c2b0231128891a41ac614e5600c7e423bf245292/lib/ruby_indexer/test/method_test.rb#L52)).

When you get the entries for `@index["foo"]`, there should be two entries:

1. One for the original made, which was made private after invoking `module_function`. That one should be owned by the attached class (`Foo`)
2. And the second for the public singleton method, which should be owned by `Foo::<Class:Foo>`
Hi @vinistock, thank you! I think I am ready to push code for reviews, but I got permission denied 😅
You need to fork the repository in order to propose PRs in open source repositories you don't have permissions. This guide explains it in more detail https://docs.github.com/en/get-started/exploring-projects-on-github/contributing-to-a-project.
Hi all, and thanks for your work. I found this issue because the following example is not working for me. Not sure how easy it is to implement, because it requires some form of state - all method definitions after a bare `module_function` become module functions.

```ruby
module Foo
  def a; end # not a module function

  module_function

  def foo; end # This is now a module function

  def bar; end # This is now a module function
end
```
Hey! Indeed we missed that use case. We would essentially need to remember that that `module_function` has been invoked and from that point on we treat method declarations appropriately. I'm going to re-open the issue since we're missing that bit of behaviour.



Patch:
diff --git a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
index d96adc04e5..9d150022dd 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
+++ b/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)


Test patch:
diff --git a/lib/ruby_indexer/test/method_test.rb b/lib/ruby_indexer/test/method_test.rb
index ca9bb1282a..809980e18a 100644
--- a/lib/ruby_indexer/test/method_test.rb
+++ b/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