Fix type inference for parent scopes and method definitions with receivers

lib/ruby_lsp/listeners/definition.rb

@@ -249,7 +249,7 @@ def handle_super_node_definition
         return unless surrounding_method
 
         handle_method_definition(
-          surrounding_method,
+          surrounding_method.name,
           @type_inferrer.infer_receiver_type(@node_context),
           inherited_only: true,
         )

lib/ruby_lsp/listeners/hover.rb

@@ -299,7 +299,7 @@ def handle_super_node_hover
         surrounding_method = @node_context.surrounding_method
         return unless surrounding_method
 
-        handle_method_hover(surrounding_method, inherited_only: true)
+        handle_method_hover(surrounding_method.name, inherited_only: true)
       end
 
       #: (String message, ?inherited_only: bool) -> void

lib/ruby_lsp/node_context.rb

@@ -5,6 +5,21 @@ module RubyLsp
   # This class allows listeners to access contextual information about a node in the AST, such as its parent,
   # its namespace nesting, and the surrounding CallNode (e.g. a method call).
   class NodeContext
+    # Represents the surrounding method definition context, tracking both the method name and its receiver
+    class MethodDef
+      #: String
+      attr_reader :name
+
+      #: String?
+      attr_reader :receiver
+
+      #: (String name, String? receiver) -> void
+      def initialize(name, receiver)
+        @name = name
+        @receiver = receiver
+      end
+    end
+
     #: Prism::Node?
     attr_reader :node, :parent
 
@@ -14,7 +29,7 @@ class NodeContext
     #: Prism::CallNode?
     attr_reader :call_node
 
-    #: String?
+    #: MethodDef?
     attr_reader :surrounding_method
 
     #: (Prism::Node? node, Prism::Node? parent, Array[(Prism::ClassNode | Prism::ModuleNode | Prism::SingletonClassNode | Prism::DefNode | Prism::BlockNode | Prism::LambdaNode | Prism::ProgramNode)] nesting_nodes, Prism::CallNode? call_node) -> void
@@ -26,7 +41,7 @@ def initialize(node, parent, nesting_nodes, call_node)
 
       nesting, surrounding_method = handle_nesting_nodes(nesting_nodes)
       @nesting = nesting #: Array[String]
-      @surrounding_method = surrounding_method #: String?
+      @surrounding_method = surrounding_method #: MethodDef?
     end
 
     #: -> String
@@ -52,10 +67,10 @@ def locals_for_scope
 
     private
 
-    #: (Array[(Prism::ClassNode | Prism::ModuleNode | Prism::SingletonClassNode | Prism::DefNode | Prism::BlockNode | Prism::LambdaNode | Prism::ProgramNode)] nodes) -> [Array[String], String?]
+    #: (Array[(Prism::ClassNode | Prism::ModuleNode | Prism::SingletonClassNode | Prism::DefNode | Prism::BlockNode | Prism::LambdaNode | Prism::ProgramNode)] nodes) -> [Array[String], MethodDef?]
     def handle_nesting_nodes(nodes)
       nesting = []
-      surrounding_method = nil #: String?
+      surrounding_method = nil #: MethodDef?
 
       @nesting_nodes.each do |node|
         case node
@@ -64,10 +79,16 @@ def handle_nesting_nodes(nodes)
         when Prism::SingletonClassNode
           nesting << "<#{nesting.flat_map { |n| n.split("::") }.last}>"
         when Prism::DefNode
-          surrounding_method = node.name.to_s
-          next unless node.receiver.is_a?(Prism::SelfNode)
-
-          nesting << "<#{nesting.flat_map { |n| n.split("::") }.last}>"
+          receiver = node.receiver
+
+          surrounding_method = case receiver
+          when Prism::SelfNode
+            MethodDef.new(node.name.to_s, "self")
+          when Prism::ConstantReadNode, Prism::ConstantPathNode
+            MethodDef.new(node.name.to_s, receiver.slice)
+          else
+            MethodDef.new(node.name.to_s, nil)
+          end
         end
       end
 

lib/ruby_lsp/type_inferrer.rb

@@ -135,15 +135,90 @@ def self_receiver_handling(node_context)
       # If we're at the top level, then the invocation is happening on `<main>`, which is a special singleton that
       # inherits from Object
       return Type.new("Object") if nesting.empty?
-      return Type.new(node_context.fully_qualified_name) if node_context.surrounding_method
+
+      surrounding_method = node_context.surrounding_method
+
+      if surrounding_method
+        receiver_name = surrounding_method.receiver
+
+        case receiver_name
+        when "self"
+          # `def self.foo` — self is the singleton of the enclosing class/module
+          return resolve_singleton_type_from_nesting(nesting)
+        when nil
+          # Instance method — self is an instance of the enclosing class/module
+          return resolve_type_from_nesting(nesting)
+        else
+          # Explicit constant receiver (e.g. `def Bar.baz`) — self is that constant's singleton class
+          resolved = resolve_receiver_singleton_type(receiver_name, nesting)
+          return resolved if resolved
+
+          return resolve_type_from_nesting(nesting)
+        end
+      end
 
       # If we're not inside a method, then we're inside the body of a class or module, which is a singleton
-      # context.
-      #
-      # If the class/module definition is using compact style (e.g.: `class Foo::Bar`), then we need to split the name
-      # into its individual parts to build the correct singleton name
-      parts = nesting.flat_map { |part| part.split("::") }
-      Type.new("#{parts.join("::")}::<#{parts.last}>")
+      # context. Resolve through the graph to get the correct fully qualified name
+      resolve_singleton_type_from_nesting(nesting)
+    end
+
+    # Resolves the fully qualified name of the innermost constant from the nesting and returns it as a type.
+    # For instance methods, the nesting won't have singleton markers, so the result is an instance type.
+    # For `def self.` methods, the nesting includes a singleton marker, which is preserved in the result.
+    #: (Array[String] nesting) -> Type
+    def resolve_type_from_nesting(nesting)
+      resolved_name = resolve_nesting_fully_qualified_name(nesting)
+      Type.new(resolved_name)
+    end
+
+    # Resolves the nesting and returns a singleton type (appends `::<Last>`)
+    #: (Array[String] nesting) -> Type
+    def resolve_singleton_type_from_nesting(nesting)
+      resolved_name = resolve_nesting_fully_qualified_name(nesting)
+      last_part = resolved_name.split("::").last #: as !nil
+      Type.new("#{resolved_name}::<#{last_part}>")
+    end
+
+    # Resolves the innermost constant in the nesting through the graph, handling compact-path definitions
+    # like `class Bar::Baz` inside a different module where the lexical nesting doesn't reflect the true
+    # constant hierarchy. Falls back to lexical joining if resolution fails.
+    #: (Array[String] nesting) -> String
+    def resolve_nesting_fully_qualified_name(nesting)
+      nesting_parts = nesting.dup
+      trailing_singletons = [] #: Array[String]
+
+      nesting_parts.reverse_each do |part|
+        break unless part.start_with?("<")
+
+        popped = nesting_parts.pop #: as !nil
+        trailing_singletons.unshift(popped)
+      end
+
+      if nesting_parts.any?
+        resolved = @graph.resolve_constant(
+          nesting_parts.last, #: as !nil
+          nesting_parts[0...-1], #: as !nil
+        )
+
+        if resolved
+          parts = resolved.name.split("::") + trailing_singletons
+          return parts.join("::")
+        end
+      end
+
+      # Fallback to lexical joining if resolution fails
+      nesting.flat_map { |part| part.split("::") }.join("::")
+    end
+
+    #: (String, Array[String]) -> Type?
+    def resolve_receiver_singleton_type(receiver_name, nesting)
+      receiver_declaration = @graph.resolve_constant(receiver_name, nesting)
+      return unless receiver_declaration.is_a?(Rubydex::Namespace)
+
+      singleton = receiver_declaration.singleton_class
+      return unless singleton
+
+      Type.new(singleton.name)
     end
 
     #: (NodeContext node_context) -> Type?