Track method visibility in the index

This commit supports visibility tracking for methods in the index by:

- Adding a visibility field to the Member entry, which is a superclass of
  Accessor, Method, and Alias.
- In the DeclarationListener, create a nested stack structure to track
  the visibility of the current scope.

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([[:public]], T::Array[T::Array[Symbol]])
       @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)
+      create_visibility_scope
       name = node.constant_path.location.slice
 
       comments = collect_comments(node)
@@ -66,6 +69,7 @@ def on_class_node_enter(node)
         node.location,
         comments,
         parent_class,
+        :public,
       )
 
       @owner_stack << entry
@@ -77,14 +81,16 @@ def on_class_node_enter(node)
     def on_class_node_leave(node)
       @stack.pop
       @owner_stack.pop
+      drop_visibility_scope
     end
 
     sig { params(node: Prism::ModuleNode).void }
     def on_module_node_enter(node)
+      create_visibility_scope
       name = node.constant_path.location.slice
 
       comments = collect_comments(node)
-      entry = Entry::Module.new(fully_qualify_name(name), @file_path, node.location, comments)
+      entry = Entry::Module.new(fully_qualify_name(name), @file_path, node.location, comments, :public)
 
       @owner_stack << entry
       @index << entry
@@ -95,6 +101,7 @@ def on_module_node_enter(node)
     def on_module_node_leave(node)
       @stack.pop
       @owner_stack.pop
+      drop_visibility_scope
     end
 
     sig { params(node: Prism::MultiWriteNode).void }
@@ -198,6 +205,21 @@ def on_call_node_enter(node)
         handle_module_operation(node, :included_modules)
       when :prepend
         handle_module_operation(node, :prepended_modules)
+      when :public, :protected, :private
+        change_scope_visibility(message)
+      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)
+          restore_scope_visibility
+        end
       end
     end
 
@@ -206,6 +228,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 +237,7 @@ def on_def_node_enter(node)
           node.location,
           comments,
           node.parameters,
+          current_visibility,
           @owner_stack.last,
         )
       when Prism::SelfNode
@@ -223,6 +247,7 @@ def on_def_node_enter(node)
           node.location,
           comments,
           node.parameters,
+          current_visibility,
           @owner_stack.last,
         )
       end
@@ -284,19 +309,43 @@ def add_constant(node, name, value = nil)
 
       @index << case value
       when Prism::ConstantReadNode, Prism::ConstantPathNode
-        Entry::UnresolvedAlias.new(value.slice, @stack.dup, name, @file_path, node.location, comments)
+        Entry::UnresolvedAlias.new(
+          value.slice,
+          @stack.dup,
+          name,
+          @file_path,
+          node.location,
+          comments,
+          current_visibility,
+        )
       when Prism::ConstantWriteNode, Prism::ConstantAndWriteNode, Prism::ConstantOrWriteNode,
         Prism::ConstantOperatorWriteNode
 
         # If the right hand side is another constant assignment, we need to visit it because that constant has to be
         # indexed too
-        Entry::UnresolvedAlias.new(value.name.to_s, @stack.dup, name, @file_path, node.location, comments)
+        Entry::UnresolvedAlias.new(
+          value.name.to_s,
+          @stack.dup,
+          name,
+          @file_path,
+          node.location,
+          comments,
+          current_visibility,
+        )
       when Prism::ConstantPathWriteNode, Prism::ConstantPathOrWriteNode, Prism::ConstantPathOperatorWriteNode,
         Prism::ConstantPathAndWriteNode
 
-        Entry::UnresolvedAlias.new(value.target.slice, @stack.dup, name, @file_path, node.location, comments)
+        Entry::UnresolvedAlias.new(
+          value.target.slice,
+          @stack.dup,
+          name,
+          @file_path,
+          node.location,
+          comments,
+          current_visibility,
+        )
       else
-        Entry::Constant.new(name, @file_path, node.location, comments)
+        Entry::Constant.new(name, @file_path, node.location, comments, current_visibility)
       end
     end
 
@@ -354,8 +403,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 +436,32 @@ def handle_module_operation(node, operation)
       collection = operation == :included_modules ? owner.included_modules : owner.prepended_modules
       collection.concat(names)
     end
+
+    sig { returns(Symbol) }
+    def current_visibility
+      T.must(@visibility_stack.last&.last)
+    end
+
+    sig { params(visibility: Symbol).void }
+    def change_scope_visibility(visibility)
+      scope_visibility_stack = T.must(@visibility_stack.last)
+      scope_visibility_stack.push(visibility)
+    end
+
+    sig { void }
+    def restore_scope_visibility
+      scope_visibility_stack = T.must(@visibility_stack.last)
+      scope_visibility_stack.pop
+    end
+
+    sig { void }
+    def create_visibility_scope
+      @visibility_stack.push([:public])
+    end
+
+    sig { void }
+    def drop_visibility_scope
+      @visibility_stack.pop
+    end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/entry.rb

@@ -26,13 +26,14 @@ class Entry
         file_path: String,
         location: T.any(Prism::Location, RubyIndexer::Location),
         comments: T::Array[String],
+        visibility: Symbol,
       ).void
     end
-    def initialize(name, file_path, location, comments)
+    def initialize(name, file_path, location, comments, visibility)
       @name = name
       @file_path = file_path
       @comments = comments
-      @visibility = T.let(:public, Symbol)
+      @visibility = T.let(visibility, Symbol)
 
       @location = T.let(
         if location.is_a?(Prism::Location)
@@ -89,10 +90,11 @@ class Class < Namespace
           location: T.any(Prism::Location, RubyIndexer::Location),
           comments: T::Array[String],
           parent_class: T.nilable(String),
+          visibility: Symbol,
         ).void
       end
-      def initialize(name, file_path, location, comments, parent_class)
-        super(name, file_path, location, comments)
+      def initialize(name, file_path, location, comments, parent_class, visibility) # rubocop:disable Metrics/ParameterLists
+        super(name, file_path, location, comments, visibility)
         @parent_class = T.let(parent_class, T.nilable(String))
       end
     end
@@ -182,17 +184,22 @@ class Member < Entry
       sig { returns(T.nilable(Entry::Namespace)) }
       attr_reader :owner
 
+      sig { returns(Symbol) }
+      attr_reader :visibility
+
       sig do
         params(
           name: String,
           file_path: String,
           location: T.any(Prism::Location, RubyIndexer::Location),
           comments: T::Array[String],
+          visibility: Symbol,
           owner: T.nilable(Entry::Namespace),
         ).void
       end
-      def initialize(name, file_path, location, comments, owner)
-        super(name, file_path, location, comments)
+      def initialize(name, file_path, location, comments, visibility, owner) # rubocop:disable Metrics/ParameterLists
+        super(name, file_path, location, comments, visibility)
+        @visibility = visibility
         @owner = owner
       end
 
@@ -227,11 +234,12 @@ class Method < Member
           location: T.any(Prism::Location, RubyIndexer::Location),
           comments: T::Array[String],
           parameters_node: T.nilable(Prism::ParametersNode),
+          visibility: Symbol,
           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
@@ -356,10 +364,11 @@ class UnresolvedAlias < Entry
           file_path: String,
           location: T.any(Prism::Location, RubyIndexer::Location),
           comments: T::Array[String],
+          visibility: Symbol,
         ).void
       end
-      def initialize(target, nesting, name, file_path, location, comments) # rubocop:disable Metrics/ParameterLists
-        super(name, file_path, location, comments)
+      def initialize(target, nesting, name, file_path, location, comments, visibility) # rubocop:disable Metrics/ParameterLists
+        super(name, file_path, location, comments, visibility)
 
         @target = target
         @nesting = nesting
@@ -375,7 +384,13 @@ class Alias < Entry
 
       sig { params(target: String, unresolved_alias: UnresolvedAlias).void }
       def initialize(target, unresolved_alias)
-        super(unresolved_alias.name, unresolved_alias.file_path, unresolved_alias.location, unresolved_alias.comments)
+        super(
+          unresolved_alias.name,
+          unresolved_alias.file_path,
+          unresolved_alias.location,
+          unresolved_alias.comments,
+          unresolved_alias.visibility
+        )
 
         @target = target
       end

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: :private)
+      assert_entry("bar", Entry::InstanceMethod, "/fake/path/foo.rb:3-0:3-12", visibility: :public)
+      assert_entry("baz", Entry::InstanceMethod, "/fake/path/foo.rb:7-0:7-12", 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: :private)
+      assert_entry("bar", Entry::InstanceMethod, "/fake/path/foo.rb:6-4:6-16", visibility: :public)
+      assert_entry("baz", Entry::InstanceMethod, "/fake/path/foo.rb:9-2:9-14", 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)