Handle module_function with no arguments

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -24,7 +24,7 @@ def initialize(index, dispatcher, parse_result, uri, collect_comments: false)
       @index = index
       @uri = uri
       @enhancements = T.let(Enhancement.all(self), T::Array[Enhancement])
-      @visibility_stack = T.let([Entry::Visibility::PUBLIC], T::Array[Entry::Visibility])
+      @visibility_stack = T.let([VisibilityScope.public_scope], T::Array[VisibilityScope])
       @comments_by_line = T.let(
         parse_result.comments.to_h do |c|
           [c.location.start_line, c]
@@ -138,7 +138,7 @@ def on_module_node_leave(node)
 
     sig { params(node: Prism::SingletonClassNode).void }
     def on_singleton_class_node_enter(node)
-      @visibility_stack.push(Entry::Visibility::PUBLIC)
+      @visibility_stack.push(VisibilityScope.public_scope)
 
       current_owner = @owner_stack.last
 
@@ -280,15 +280,15 @@ def on_call_node_enter(node)
       when :include, :prepend, :extend
         handle_module_operation(node, message)
       when :public
-        @visibility_stack.push(Entry::Visibility::PUBLIC)
+        @visibility_stack.push(VisibilityScope.public_scope)
       when :protected
-        @visibility_stack.push(Entry::Visibility::PROTECTED)
+        @visibility_stack.push(VisibilityScope.new(visibility: Entry::Visibility::PROTECTED))
       when :private
-        @visibility_stack.push(Entry::Visibility::PRIVATE)
+        @visibility_stack.push(VisibilityScope.new(visibility: Entry::Visibility::PRIVATE))
       when :module_function
         handle_module_function(node)
       when :private_class_method
-        @visibility_stack.push(Entry::Visibility::PRIVATE)
+        @visibility_stack.push(VisibilityScope.new(visibility: Entry::Visibility::PRIVATE))
         handle_private_class_method(node)
       end
 
@@ -324,42 +324,61 @@ def on_call_node_leave(node)
 
     sig { params(node: Prism::DefNode).void }
     def on_def_node_enter(node)
+      owner = @owner_stack.last
+      return unless owner
+
       @inside_def = true
       method_name = node.name.to_s
       comments = collect_comments(node)
+      scope = current_visibility_scope
 
       case node.receiver
       when nil
+        location = Location.from_prism_location(node.location, @code_units_cache)
+        name_location = Location.from_prism_location(node.name_loc, @code_units_cache)
+        signatures = [Entry::Signature.new(list_params(node.parameters))]
+
         @index.add(Entry::Method.new(
           method_name,
           @uri,
-          Location.from_prism_location(node.location, @code_units_cache),
-          Location.from_prism_location(node.name_loc, @code_units_cache),
+          location,
+          name_location,
           comments,
-          [Entry::Signature.new(list_params(node.parameters))],
-          current_visibility,
-          @owner_stack.last,
+          signatures,
+          scope.visibility,
+          owner,
         ))
-      when Prism::SelfNode
-        owner = @owner_stack.last
 
-        if owner
+        if scope.module_func
           singleton = @index.existing_or_new_singleton_class(owner.name)
 
           @index.add(Entry::Method.new(
             method_name,
             @uri,
-            Location.from_prism_location(node.location, @code_units_cache),
-            Location.from_prism_location(node.name_loc, @code_units_cache),
+            location,
+            name_location,
             comments,
-            [Entry::Signature.new(list_params(node.parameters))],
-            current_visibility,
+            signatures,
+            Entry::Visibility::PUBLIC,
             singleton,
           ))
-
-          @owner_stack << singleton
-          @stack << "<Class:#{@stack.last}>"
         end
+      when Prism::SelfNode
+        singleton = @index.existing_or_new_singleton_class(owner.name)
+
+        @index.add(Entry::Method.new(
+          method_name,
+          @uri,
+          Location.from_prism_location(node.location, @code_units_cache),
+          Location.from_prism_location(node.name_loc, @code_units_cache),
+          comments,
+          [Entry::Signature.new(list_params(node.parameters))],
+          scope.visibility,
+          singleton,
+        ))
+
+        @owner_stack << singleton
+        @stack << "<Class:#{@stack.last}>"
       end
     end
 
@@ -834,6 +853,8 @@ def handle_attribute(node, reader:, writer:)
       return unless receiver.nil? || receiver.is_a?(Prism::SelfNode)
 
       comments = collect_comments(node)
+      scope = current_visibility_scope
+
       arguments.each do |argument|
         name, loc = case argument
         when Prism::SymbolNode
@@ -850,7 +871,7 @@ def handle_attribute(node, reader:, writer:)
             @uri,
             Location.from_prism_location(loc, @code_units_cache),
             comments,
-            current_visibility,
+            scope.visibility,
             @owner_stack.last,
           ))
         end
@@ -862,7 +883,7 @@ def handle_attribute(node, reader:, writer:)
           @uri,
           Location.from_prism_location(loc, @code_units_cache),
           comments,
-          current_visibility,
+          scope.visibility,
           @owner_stack.last,
         ))
       end
@@ -904,11 +925,20 @@ def handle_module_operation(node, operation)
 
     sig { params(node: Prism::CallNode).void }
     def handle_module_function(node)
+      # Invoking `module_function` in a class raises
+      owner = @owner_stack.last
+      return unless owner.is_a?(Entry::Module)
+
       arguments_node = node.arguments
-      return unless arguments_node
 
-      owner_name = @owner_stack.last&.name
-      return unless owner_name
+      # If `module_function` is invoked without arguments, all methods defined after it become singleton methods and the
+      # visibility for instance methods changes to private
+      unless arguments_node
+        @visibility_stack.push(VisibilityScope.module_function_scope)
+        return
+      end
+
+      owner_name = owner.name
 
       arguments_node.arguments.each do |argument|
         method_name = case argument
@@ -983,8 +1013,8 @@ def handle_private_class_method(node)
       end
     end
 
-    sig { returns(Entry::Visibility) }
-    def current_visibility
+    sig { returns(VisibilityScope) }
+    def current_visibility_scope
       T.must(@visibility_stack.last)
     end
 
@@ -1091,7 +1121,7 @@ def actual_nesting(name)
 
     sig { params(short_name: String, entry: Entry::Namespace).void }
     def advance_namespace_stack(short_name, entry)
-      @visibility_stack.push(Entry::Visibility::PUBLIC)
+      @visibility_stack.push(VisibilityScope.public_scope)
       @owner_stack << entry
       @index.add(entry)
       @stack << short_name

lib/ruby_indexer/lib/ruby_indexer/visibility_scope.rb

@@ -0,0 +1,36 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyIndexer
+  # Represents the visibility scope in a Ruby namespace. This keeps track of whether methods are in a public, private or
+  # protected section, and whether they are module functions.
+  class VisibilityScope
+    extend T::Sig
+
+    class << self
+      extend T::Sig
+
+      sig { returns(T.attached_class) }
+      def module_function_scope
+        new(module_func: true, visibility: Entry::Visibility::PRIVATE)
+      end
+
+      sig { returns(T.attached_class) }
+      def public_scope
+        new
+      end
+    end
+
+    sig { returns(Entry::Visibility) }
+    attr_reader :visibility
+
+    sig { returns(T::Boolean) }
+    attr_reader :module_func
+
+    sig { params(visibility: Entry::Visibility, module_func: T::Boolean).void }
+    def initialize(visibility: Entry::Visibility::PUBLIC, module_func: false)
+      @visibility = visibility
+      @module_func = module_func
+    end
+  end
+end

lib/ruby_indexer/ruby_indexer.rb

@@ -5,6 +5,7 @@
 require "did_you_mean"
 
 require "ruby_indexer/lib/ruby_indexer/uri"
+require "ruby_indexer/lib/ruby_indexer/visibility_scope"
 require "ruby_indexer/lib/ruby_indexer/declaration_listener"
 require "ruby_indexer/lib/ruby_indexer/reference_finder"
 require "ruby_indexer/lib/ruby_indexer/enhancement"

lib/ruby_indexer/test/instance_variables_test.rb

@@ -216,5 +216,25 @@ def something.bar
       assert_instance_of(Entry::Class, owner)
       assert_equal("Foo", owner.name)
     end
+
+    def test_module_function_does_not_impact_instance_variables
+      # One possible way of implementing `module_function` would be to push a fake singleton class to the stack, so that
+      # methods are inserted into it. However, that would be incorrect because it would then bind instance variables to
+      # the wrong type. This test is here to prevent that from happening.
+      index(<<~RUBY)
+        module Foo
+          module_function
+
+          def something; end
+
+          @a = 123
+        end
+      RUBY
+
+      entry = T.must(@index["@a"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::SingletonClass, owner)
+      assert_equal("Foo::<Class:Foo>", owner.name)
+    end
   end
 end

lib/ruby_indexer/test/method_test.rb

@@ -88,19 +88,21 @@ def bar
 
     def test_visibility_tracking
       index(<<~RUBY)
-        private def foo
-        end
+        class Foo
+          private def foo
+          end
 
-        def bar; end
+          def bar; end
 
-        protected
+          protected
 
-        def baz; end
+          def baz; end
+        end
       RUBY
 
-      assert_entry("foo", Entry::Method, "/fake/path/foo.rb:0-8:1-3", visibility: Entry::Visibility::PRIVATE)
-      assert_entry("bar", Entry::Method, "/fake/path/foo.rb:3-0:3-12", visibility: Entry::Visibility::PUBLIC)
-      assert_entry("baz", Entry::Method, "/fake/path/foo.rb:7-0:7-12", visibility: Entry::Visibility::PROTECTED)
+      assert_entry("foo", Entry::Method, "/fake/path/foo.rb:1-10:2-5", visibility: Entry::Visibility::PRIVATE)
+      assert_entry("bar", Entry::Method, "/fake/path/foo.rb:4-2:4-14", visibility: Entry::Visibility::PUBLIC)
+      assert_entry("baz", Entry::Method, "/fake/path/foo.rb:8-2:8-14", visibility: Entry::Visibility::PROTECTED)
     end
 
     def test_visibility_tracking_with_nested_class_or_modules
@@ -846,6 +848,67 @@ def baz(a, b)
       assert_signature_matches(entry, "baz(1)")
     end
 
+    def test_module_function_with_no_arguments
+      index(<<~RUBY)
+        module Foo
+          def bar; end
+
+          module_function
+
+          def baz; end
+          attr_reader :attribute
+
+          public
+
+          def qux; end
+        end
+      RUBY
+
+      entry = T.must(@index["bar"].first)
+      assert_predicate(entry, :public?)
+      assert_equal("Foo", T.must(entry.owner).name)
+
+      instance_baz, singleton_baz = T.must(@index["baz"])
+      assert_predicate(instance_baz, :private?)
+      assert_equal("Foo", T.must(instance_baz.owner).name)
+
+      assert_predicate(singleton_baz, :public?)
+      assert_equal("Foo::<Class:Foo>", T.must(singleton_baz.owner).name)
+
+      # After invoking `public`, the state of `module_function` is reset
+      instance_qux, singleton_qux = T.must(@index["qux"])
+      assert_nil(singleton_qux)
+      assert_predicate(instance_qux, :public?)
+      assert_equal("Foo", T.must(instance_baz.owner).name)
+
+      # Attributes are not turned into class methods, they do become private
+      instance_attribute, singleton_attribute = @index["attribute"]
+      assert_nil(singleton_attribute)
+      assert_equal("Foo", T.must(instance_attribute.owner).name)
+      assert_predicate(instance_attribute, :private?)
+    end
+
+    def test_module_function_does_nothing_in_classes
+      # Invoking `module_function` in a class raises an error. We simply ignore it
+      index(<<~RUBY)
+        class Foo
+          def bar; end
+
+          module_function
+
+          def baz; end
+        end
+      RUBY
+
+      entry = T.must(@index["bar"].first)
+      assert_predicate(entry, :public?)
+      assert_equal("Foo", T.must(entry.owner).name)
+
+      entry = T.must(@index["baz"].first)
+      assert_predicate(entry, :public?)
+      assert_equal("Foo", T.must(entry.owner).name)
+    end
+
     private
 
     sig { params(entry: Entry::Method, call_string: String).void }