From fac03ebc7685047dc61dbb4f05cf8132ac4e86d3 Mon Sep 17 00:00:00 2001
From: Mike Dalton <michaelcdalton@gmail.com>
Date: Thu, 14 Nov 2024 16:29:27 -0600
Subject: [PATCH] Support `private_class_method` in the indexer

Fixes https://github.com/Shopify/ruby-lsp/issues/2781
---
 .../lib/ruby_indexer/declaration_listener.rb  | 39 +++++++++++++++++
 lib/ruby_indexer/test/method_test.rb          | 42 +++++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
index b29bdd7db..a0caaf5e9 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
@@ -314,6 +314,8 @@ def on_call_node_enter(node)
         @visibility_stack.push(Entry::Visibility::PRIVATE)
       when :module_function
         handle_module_function(node)
+      when :private_class_method
+        handle_private_class_method(node)
       end
 
       @enhancements.each do |enhancement|
@@ -805,6 +807,43 @@ def handle_module_function(node)
       end
     end
 
+    sig { params(node: Prism::CallNode).void }
+    def handle_private_class_method(node)
+      node.arguments&.arguments&.each do |argument|
+        string_or_symbol_nodes = case argument
+        when Prism::StringNode, Prism::SymbolNode
+          [argument]
+        when Prism::ArrayNode
+          argument.elements
+        else
+          []
+        end
+
+        string_or_symbol_nodes.each do |string_or_symbol_node|
+          method_name = case string_or_symbol_node
+          when Prism::StringNode
+            string_or_symbol_node.content
+          when Prism::SymbolNode
+            string_or_symbol_node.value
+          end
+          next unless method_name
+
+          owner_name = @owner_stack.last&.name
+          next unless owner_name
+
+          entries = @index.resolve_method(method_name, @index.existing_or_new_singleton_class(owner_name).name)
+          next unless entries
+
+          entries.each do |entry|
+            entry_owner_name = entry.owner&.name
+            next unless entry_owner_name
+
+            entry.visibility = Entry::Visibility::PRIVATE
+          end
+        end
+      end
+    end
+
     sig { returns(Entry::Visibility) }
     def current_visibility
       T.must(@visibility_stack.last)
diff --git a/lib/ruby_indexer/test/method_test.rb b/lib/ruby_indexer/test/method_test.rb
index 809980e18..b8da0a713 100644
--- a/lib/ruby_indexer/test/method_test.rb
+++ b/lib/ruby_indexer/test/method_test.rb
@@ -149,6 +149,48 @@ def bar; end
       end
     end
 
+    def test_private_class_method_visibility_tracking_string_symbol_arguments
+      index(<<~RUBY)
+        class Test
+          def self.foo
+          end
+
+          def self.bar
+          end
+
+          private_class_method("foo", :bar)
+        end
+      RUBY
+
+      ["foo", "bar"].each do |keyword|
+        entries = T.must(@index[keyword])
+        assert_equal(1, entries.size)
+        entry = entries.first
+        assert_equal(Entry::Visibility::PRIVATE, entry.visibility)
+      end
+    end
+
+    def test_private_class_method_visibility_tracking_array_argument
+      index(<<~RUBY)
+        class Test
+          def self.foo
+          end
+
+          def self.bar
+          end
+
+          private_class_method(["foo", :bar])
+        end
+      RUBY
+
+      ["foo", "bar"].each do |keyword|
+        entries = T.must(@index[keyword])
+        assert_equal(1, entries.size)
+        entry = entries.first
+        assert_equal(Entry::Visibility::PRIVATE, entry.visibility)
+      end
+    end
+
     def test_method_with_parameters
       index(<<~RUBY)
         class Foo