From 290c082b3bf2fb6d5244598ca6d6239937f892c3 Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Tue, 14 Jan 2025 09:31:15 -0500
Subject: [PATCH] Fix singleton indexing under compact namespaces

---
 .../lib/ruby_indexer/declaration_listener.rb  | 14 +++++++++++---
 .../test/classes_and_modules_test.rb          | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
index f46a034bc..52aa5f3c7 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
@@ -143,7 +143,7 @@ def on_singleton_class_node_enter(node)
 
       if current_owner
         expression = node.expression
-        name = (expression.is_a?(Prism::SelfNode) ? "<Class:#{@stack.last}>" : "<Class:#{expression.slice}>")
+        name = (expression.is_a?(Prism::SelfNode) ? "<Class:#{last_name_in_stack}>" : "<Class:#{expression.slice}>")
         real_nesting = actual_nesting(name)
 
         existing_entries = T.cast(@index[real_nesting.join("::")], T.nilable(T::Array[Entry::SingletonClass]))
@@ -376,7 +376,6 @@ def on_def_node_enter(node)
         ))
 
         @owner_stack << singleton
-        @stack << "<Class:#{@stack.last}>"
       end
     end
 
@@ -386,7 +385,6 @@ def on_def_node_leave(node)
 
       if node.receiver.is_a?(Prism::SelfNode)
         @owner_stack.pop
-        @stack.pop
       end
     end
 
@@ -1127,5 +1125,15 @@ def advance_namespace_stack(short_name, entry)
       @index.add(entry)
       @stack << short_name
     end
+
+    # Returns the last name in the stack not as we found it, but in terms of declared constants. For example, if the
+    # last entry in the stack is a compact namespace like `Foo::Bar`, then the last name is `Bar`
+    sig { returns(T.nilable(String)) }
+    def last_name_in_stack
+      name = @stack.last
+      return unless name
+
+      name.split("::").last
+    end
   end
 end
diff --git a/lib/ruby_indexer/test/classes_and_modules_test.rb b/lib/ruby_indexer/test/classes_and_modules_test.rb
index 91227a966..27c332c51 100644
--- a/lib/ruby_indexer/test/classes_and_modules_test.rb
+++ b/lib/ruby_indexer/test/classes_and_modules_test.rb
@@ -647,5 +647,24 @@ class Foo
       entry = @index["Foo"].first
       assert_empty(entry.comments)
     end
+
+    def test_singleton_inside_compact_namespace
+      index(<<~RUBY)
+        module Foo::Bar
+          class << self
+            def baz; end
+          end
+        end
+      RUBY
+
+      # Verify we didn't index the incorrect name
+      assert_nil(@index["Foo::Bar::<Class:Foo::Bar>"])
+
+      # Verify we indexed the correct name
+      assert_entry("Foo::Bar::<Class:Bar>", Entry::SingletonClass, "/fake/path/foo.rb:1-2:3-5")
+
+      method = @index["baz"]&.first
+      assert_equal("Foo::Bar::<Class:Bar>", method.owner.name)
+    end
   end
 end