Track method's visibility during indexing

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -11,6 +11,7 @@ class DeclarationListener
     def initialize(index, dispatcher, parse_result, file_path)
       @index = index
       @file_path = file_path
+      @visibility_stack = T.let([Entry::Visibility::PUBLIC], T::Array[Entry::Visibility])
       @comments_by_line = T.let(
         parse_result.comments.to_h do |c|
           [c.location.start_line, c]
@@ -35,6 +36,7 @@ def initialize(index, dispatcher, parse_result, file_path)
         :on_def_node_enter,
         :on_def_node_leave,
         :on_call_node_enter,
+        :on_call_node_leave,
         :on_multi_write_node_enter,
         :on_constant_path_write_node_enter,
         :on_constant_path_or_write_node_enter,
@@ -50,6 +52,7 @@ def initialize(index, dispatcher, parse_result, file_path)
 
     sig { params(node: Prism::ClassNode).void }
     def on_class_node_enter(node)
+      @visibility_stack.push(Entry::Visibility::PUBLIC)
       name = node.constant_path.location.slice
 
       comments = collect_comments(node)
@@ -77,10 +80,12 @@ def on_class_node_enter(node)
     def on_class_node_leave(node)
       @stack.pop
       @owner_stack.pop
+      @visibility_stack.pop
     end
 
     sig { params(node: Prism::ModuleNode).void }
     def on_module_node_enter(node)
+      @visibility_stack.push(Entry::Visibility::PUBLIC)
       name = node.constant_path.location.slice
 
       comments = collect_comments(node)
@@ -95,6 +100,7 @@ def on_module_node_enter(node)
     def on_module_node_leave(node)
       @stack.pop
       @owner_stack.pop
+      @visibility_stack.pop
     end
 
     sig { params(node: Prism::MultiWriteNode).void }
@@ -198,6 +204,25 @@ def on_call_node_enter(node)
         handle_module_operation(node, :included_modules)
       when :prepend
         handle_module_operation(node, :prepended_modules)
+      when :public
+        @visibility_stack.push(Entry::Visibility::PUBLIC)
+      when :protected
+        @visibility_stack.push(Entry::Visibility::PROTECTED)
+      when :private
+        @visibility_stack.push(Entry::Visibility::PRIVATE)
+      end
+    end
+
+    sig { params(node: Prism::CallNode).void }
+    def on_call_node_leave(node)
+      message = node.name
+      case message
+      when :public, :protected, :private
+        # We want to restore the visibility stack when we leave a method definition with a visibility modifier
+        # e.g. `private def foo; end`
+        if node.arguments&.arguments&.first&.is_a?(Prism::DefNode)
+          @visibility_stack.pop
+        end
       end
     end
 
@@ -206,6 +231,7 @@ def on_def_node_enter(node)
       @inside_def = true
       method_name = node.name.to_s
       comments = collect_comments(node)
+
       case node.receiver
       when nil
         @index << Entry::InstanceMethod.new(
@@ -214,6 +240,7 @@ def on_def_node_enter(node)
           node.location,
           comments,
           node.parameters,
+          current_visibility,
           @owner_stack.last,
         )
       when Prism::SelfNode
@@ -223,6 +250,7 @@ def on_def_node_enter(node)
           node.location,
           comments,
           node.parameters,
+          current_visibility,
           @owner_stack.last,
         )
       end
@@ -257,7 +285,7 @@ def handle_private_constant(node)
       # The private_constant method does not resolve the constant name. It always points to a constant that needs to
       # exist in the current namespace
       entries = @index[fully_qualify_name(name)]
-      entries&.each { |entry| entry.visibility = :private }
+      entries&.each { |entry| entry.visibility = Entry::Visibility::PRIVATE }
     end
 
     sig do
@@ -354,8 +382,15 @@ def handle_attribute(node, reader:, writer:)
 
         next unless name && loc
 
-        @index << Entry::Accessor.new(name, @file_path, loc, comments, @owner_stack.last) if reader
-        @index << Entry::Accessor.new("#{name}=", @file_path, loc, comments, @owner_stack.last) if writer
+        @index << Entry::Accessor.new(name, @file_path, loc, comments, current_visibility, @owner_stack.last) if reader
+        @index << Entry::Accessor.new(
+          "#{name}=",
+          @file_path,
+          loc,
+          comments,
+          current_visibility,
+          @owner_stack.last,
+        ) if writer
       end
     end
 
@@ -380,5 +415,10 @@ def handle_module_operation(node, operation)
       collection = operation == :included_modules ? owner.included_modules : owner.prepended_modules
       collection.concat(names)
     end
+
+    sig { returns(Entry::Visibility) }
+    def current_visibility
+      T.must(@visibility_stack.last)
+    end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/entry.rb

@@ -3,6 +3,14 @@
 
 module RubyIndexer
   class Entry
+    class Visibility < T::Enum
+      enums do
+        PUBLIC = new(:public)
+        PROTECTED = new(:protected)
+        PRIVATE = new(:private)
+      end
+    end
+
     extend T::Sig
 
     sig { returns(String) }
@@ -17,7 +25,7 @@ class Entry
     sig { returns(T::Array[String]) }
     attr_reader :comments
 
-    sig { returns(Symbol) }
+    sig { returns(Visibility) }
     attr_accessor :visibility
 
     sig do
@@ -32,7 +40,7 @@ def initialize(name, file_path, location, comments)
       @name = name
       @file_path = file_path
       @comments = comments
-      @visibility = T.let(:public, Symbol)
+      @visibility = T.let(Visibility::PUBLIC, Visibility)
 
       @location = T.let(
         if location.is_a?(Prism::Location)
@@ -188,11 +196,13 @@ class Member < Entry
           file_path: String,
           location: T.any(Prism::Location, RubyIndexer::Location),
           comments: T::Array[String],
+          visibility: Visibility,
           owner: T.nilable(Entry::Namespace),
         ).void
       end
-      def initialize(name, file_path, location, comments, owner)
+      def initialize(name, file_path, location, comments, visibility, owner) # rubocop:disable Metrics/ParameterLists
         super(name, file_path, location, comments)
+        @visibility = visibility
         @owner = owner
       end
 
@@ -227,11 +237,12 @@ class Method < Member
           location: T.any(Prism::Location, RubyIndexer::Location),
           comments: T::Array[String],
           parameters_node: T.nilable(Prism::ParametersNode),
+          visibility: Visibility,
           owner: T.nilable(Entry::Namespace),
         ).void
       end
-      def initialize(name, file_path, location, comments, parameters_node, owner) # rubocop:disable Metrics/ParameterLists
-        super(name, file_path, location, comments, owner)
+      def initialize(name, file_path, location, comments, parameters_node, visibility, owner) # rubocop:disable Metrics/ParameterLists
+        super(name, file_path, location, comments, visibility, owner)
 
         @parameters = T.let(list_params(parameters_node), T::Array[Parameter])
       end
@@ -377,6 +388,7 @@ class Alias < Entry
       def initialize(target, unresolved_alias)
         super(unresolved_alias.name, unresolved_alias.file_path, unresolved_alias.location, unresolved_alias.comments)
 
+        @visibility = unresolved_alias.visibility
         @target = target
       end
     end

lib/ruby_indexer/test/classes_and_modules_test.rb

@@ -290,13 +290,13 @@ class D; end
       RUBY
 
       b_const = @index["A::B"].first
-      assert_equal(:private, b_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, b_const.visibility)
 
       c_const = @index["A::C"].first
-      assert_equal(:private, c_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, c_const.visibility)
 
       d_const = @index["A::D"].first
-      assert_equal(:public, d_const.visibility)
+      assert_equal(Entry::Visibility::PUBLIC, d_const.visibility)
     end
 
     def test_keeping_track_of_super_classes

lib/ruby_indexer/test/constant_test.rb

@@ -122,13 +122,13 @@ class A
       RUBY
 
       b_const = @index["A::B"].first
-      assert_equal(:private, b_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, b_const.visibility)
 
       c_const = @index["A::C"].first
-      assert_equal(:private, c_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, c_const.visibility)
 
       d_const = @index["A::D"].first
-      assert_equal(:public, d_const.visibility)
+      assert_equal(Entry::Visibility::PUBLIC, d_const.visibility)
     end
 
     def test_marking_constants_as_private_reopening_namespaces
@@ -155,13 +155,13 @@ module B
       RUBY
 
       a_const = @index["A::B::CONST_A"].first
-      assert_equal(:private, a_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, a_const.visibility)
 
       b_const = @index["A::B::CONST_B"].first
-      assert_equal(:private, b_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, b_const.visibility)
 
       c_const = @index["A::B::CONST_C"].first
-      assert_equal(:private, c_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, c_const.visibility)
     end
 
     def test_marking_constants_as_private_with_receiver
@@ -179,10 +179,10 @@ module B
       RUBY
 
       a_const = @index["A::B::CONST_A"].first
-      assert_equal(:private, a_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, a_const.visibility)
 
       b_const = @index["A::B::CONST_B"].first
-      assert_equal(:private, b_const.visibility)
+      assert_equal(Entry::Visibility::PRIVATE, b_const.visibility)
     end
 
     def test_indexing_constant_aliases

lib/ruby_indexer/test/method_test.rb

@@ -71,6 +71,43 @@ def bar
       assert_equal("Bar", second_entry.owner.name)
     end
 
+    def test_visibility_tracking
+      index(<<~RUBY)
+        private def foo
+        end
+
+        def bar; end
+
+        protected
+
+        def baz; end
+      RUBY
+
+      assert_entry("foo", Entry::InstanceMethod, "/fake/path/foo.rb:0-8:1-3", visibility: Entry::Visibility::PRIVATE)
+      assert_entry("bar", Entry::InstanceMethod, "/fake/path/foo.rb:3-0:3-12", visibility: Entry::Visibility::PUBLIC)
+      assert_entry("baz", Entry::InstanceMethod, "/fake/path/foo.rb:7-0:7-12", visibility: Entry::Visibility::PROTECTED)
+    end
+
+    def test_visibility_tracking_with_nested_class_or_modules
+      index(<<~RUBY)
+        class Foo
+          private
+
+          def foo; end
+
+          class Bar
+            def bar; end
+          end
+
+          def baz; end
+        end
+      RUBY
+
+      assert_entry("foo", Entry::InstanceMethod, "/fake/path/foo.rb:3-2:3-14", visibility: Entry::Visibility::PRIVATE)
+      assert_entry("bar", Entry::InstanceMethod, "/fake/path/foo.rb:6-4:6-16", visibility: Entry::Visibility::PUBLIC)
+      assert_entry("baz", Entry::InstanceMethod, "/fake/path/foo.rb:9-2:9-14", visibility: Entry::Visibility::PRIVATE)
+    end
+
     def test_method_with_parameters
       index(<<~RUBY)
         class Foo

lib/ruby_indexer/test/test_case.rb

@@ -15,7 +15,7 @@ def index(source)
       @index.index_single(IndexablePath.new(nil, "/fake/path/foo.rb"), source)
     end
 
-    def assert_entry(expected_name, type, expected_location)
+    def assert_entry(expected_name, type, expected_location, visibility: nil)
       entries = @index[expected_name]
       refute_empty(entries, "Expected #{expected_name} to be indexed")
 
@@ -28,6 +28,8 @@ def assert_entry(expected_name, type, expected_location)
           ":#{location.end_line - 1}-#{location.end_column}"
 
       assert_equal(expected_location, location_string)
+
+      assert_equal(visibility, entry.visibility) if visibility
     end
 
     def refute_entry(expected_name)

lib/ruby_lsp/listeners/completion.rb

@@ -138,7 +138,8 @@ def constant_path_completion(name, range)
           # The only time we may have a private constant reference from outside of the namespace is if we're dealing
           # with ConstantPath and the entry name doesn't start with the current nesting
           first_entry = T.must(entries.first)
-          next if first_entry.visibility == :private && !first_entry.name.start_with?("#{@nesting}::")
+          next if first_entry.visibility == RubyIndexer::Entry::Visibility::PRIVATE &&
+            !first_entry.name.start_with?("#{@nesting}::")
 
           constant_name = first_entry.name.delete_prefix("#{real_namespace}::")
           full_name = aliased_namespace.empty? ? constant_name : "#{aliased_namespace}::#{constant_name}"

lib/ruby_lsp/listeners/definition.rb

@@ -153,7 +153,8 @@ def find_in_index(value)
         # We should only allow jumping to the definition of private constants if the constant is defined in the same
         # namespace as the reference
         first_entry = T.must(entries.first)
-        return if first_entry.visibility == :private && first_entry.name != "#{@nesting.join("::")}::#{value}"
+        return if first_entry.visibility == RubyIndexer::Entry::Visibility::PRIVATE &&
+          first_entry.name != "#{@nesting.join("::")}::#{value}"
 
         entries.each do |entry|
           location = entry.location

lib/ruby_lsp/listeners/hover.rb

@@ -111,7 +111,8 @@ def generate_hover(name, location)
         # We should only show hover for private constants if the constant is defined in the same namespace as the
         # reference
         first_entry = T.must(entries.first)
-        return if first_entry.visibility == :private && first_entry.name != "#{@nesting.join("::")}::#{name}"
+        return if first_entry.visibility == RubyIndexer::Entry::Visibility::PRIVATE &&
+          first_entry.name != "#{@nesting.join("::")}::#{name}"
 
         categorized_markdown_from_index_entries(name, entries).each do |category, content|
           @response_builder.push(content, category: category)

lib/ruby_lsp/requests/workspace_symbol.rb

@@ -39,7 +39,7 @@ def perform
           next if @global_state.typechecker && not_in_dependencies?(file_path)
 
           # We should never show private symbols when searching the entire workspace
-          next if entry.visibility == :private
+          next if entry.visibility == RubyIndexer::Entry::Visibility::PRIVATE
 
           kind = kind_for_entry(entry)
           loc = entry.location