Remember singleton nesting and method nesting when locating nodes (#2186

)

lib/ruby_lsp/document.rb

@@ -127,7 +127,10 @@ def locate(node, char_position, node_types: [])
       queue = T.let(node.child_nodes.compact, T::Array[T.nilable(Prism::Node)])
       closest = node
       parent = T.let(nil, T.nilable(Prism::Node))
-      nesting = T.let([], T::Array[T.any(Prism::ClassNode, Prism::ModuleNode)])
+      nesting_nodes = T.let(
+        [],
+        T::Array[T.any(Prism::ClassNode, Prism::ModuleNode, Prism::SingletonClassNode, Prism::DefNode)],
+      )
       call_node = T.let(nil, T.nilable(Prism::CallNode))
 
       until queue.empty?
@@ -151,13 +154,18 @@ def locate(node, char_position, node_types: [])
 
         # If the candidate starts after the end of the previous nesting level, then we've exited that nesting level and
         # need to pop the stack
-        previous_level = nesting.last
-        nesting.pop if previous_level && loc.start_offset > previous_level.location.end_offset
+        previous_level = nesting_nodes.last
+        nesting_nodes.pop if previous_level && loc.start_offset > previous_level.location.end_offset
 
         # Keep track of the nesting where we found the target. This is used to determine the fully qualified name of the
         # target when it is a constant
-        if candidate.is_a?(Prism::ClassNode) || candidate.is_a?(Prism::ModuleNode)
-          nesting << candidate
+        case candidate
+        when Prism::ClassNode, Prism::ModuleNode
+          nesting_nodes << candidate
+        when Prism::SingletonClassNode
+          nesting_nodes << candidate
+        when Prism::DefNode
+          nesting_nodes << candidate
         end
 
         if candidate.is_a?(Prism::CallNode)
@@ -189,11 +197,32 @@ def locate(node, char_position, node_types: [])
       # The correct target is `Foo::Bar` with an empty nesting. `Foo::Bar` should not appear in the nesting stack, even
       # though the class/module node does indeed enclose the target, because it would lead to incorrect behavior
       if closest.is_a?(Prism::ConstantReadNode) || closest.is_a?(Prism::ConstantPathNode)
-        last_level = nesting.last
-        nesting.pop if last_level && last_level.constant_path == closest
+        last_level = nesting_nodes.last
+
+        if (last_level.is_a?(Prism::ModuleNode) || last_level.is_a?(Prism::ClassNode)) &&
+            last_level.constant_path == closest
+          nesting_nodes.pop
+        end
+      end
+
+      nesting = []
+      surrounding_method = T.let(nil, T.nilable(String))
+
+      nesting_nodes.each do |node|
+        case node
+        when Prism::ClassNode, Prism::ModuleNode
+          nesting << node.constant_path.slice
+        when Prism::SingletonClassNode
+          nesting << "<Class:#{nesting.last}>"
+        when Prism::DefNode
+          surrounding_method = node.name.to_s
+          next unless node.receiver.is_a?(Prism::SelfNode)
+
+          nesting << "<Class:#{nesting.last}>"
+        end
       end
 
-      NodeContext.new(closest, parent, nesting.map { |n| n.constant_path.location.slice }, call_node)
+      NodeContext.new(closest, parent, nesting, call_node, surrounding_method)
     end
 
     sig { returns(T::Boolean) }

lib/ruby_lsp/node_context.rb

@@ -16,19 +16,24 @@ class NodeContext
     sig { returns(T.nilable(Prism::CallNode)) }
     attr_reader :call_node
 
+    sig { returns(T.nilable(String)) }
+    attr_reader :surrounding_method
+
     sig do
       params(
         node: T.nilable(Prism::Node),
         parent: T.nilable(Prism::Node),
         nesting: T::Array[String],
         call_node: T.nilable(Prism::CallNode),
+        surrounding_method: T.nilable(String),
       ).void
     end
-    def initialize(node, parent, nesting, call_node)
+    def initialize(node, parent, nesting, call_node, surrounding_method)
       @node = node
       @parent = parent
       @nesting = nesting
       @call_node = call_node
+      @surrounding_method = surrounding_method
     end
 
     sig { returns(String) }

test/ruby_document_test.rb

@@ -672,6 +672,50 @@ class Baz
     assert_equal("Baz", T.must(node_context.node).slice)
   end
 
+  def test_locating_singleton_contexts
+    document = RubyLsp::RubyDocument.new(source: +<<~RUBY, version: 1, uri: URI("file:///foo/bar.rb"))
+      class Foo
+        hello1
+
+        def self.bar
+          hello2
+        end
+
+        class << self
+          hello3
+
+          def baz
+            hello4
+          end
+        end
+
+        def qux
+          hello5
+        end
+      end
+    RUBY
+
+    node_context = document.locate_node({ line: 1, character: 2 })
+    assert_equal(["Foo"], node_context.nesting)
+    assert_nil(node_context.surrounding_method)
+
+    node_context = document.locate_node({ line: 4, character: 4 })
+    assert_equal(["Foo", "<Class:Foo>"], node_context.nesting)
+    assert_equal("bar", node_context.surrounding_method)
+
+    node_context = document.locate_node({ line: 8, character: 4 })
+    assert_equal(["Foo", "<Class:Foo>"], node_context.nesting)
+    assert_nil(node_context.surrounding_method)
+
+    node_context = document.locate_node({ line: 11, character: 6 })
+    assert_equal(["Foo", "<Class:Foo>"], node_context.nesting)
+    assert_equal("baz", node_context.surrounding_method)
+
+    node_context = document.locate_node({ line: 16, character: 6 })
+    assert_equal(["Foo"], node_context.nesting)
+    assert_equal("qux", node_context.surrounding_method)
+  end
+
   private
 
   def assert_error_edit(actual, error_range)