type-guessr 0.0.1 → 0.0.3

data/lib/type_guessr/core/ir/nodes.rb

@@ -1,12 +1,44 @@
 # frozen_string_literal: true
 
+require_relative "../node_key_generator"
+
 module TypeGuessr
   module Core
+    # Intermediate Representation (IR) nodes for type inference.
+    # Each node represents a construct in the source code and forms a
+    # reverse dependency graph where nodes point to their dependencies.
+    #
+    # Performance note: All node types use plain Class with positional arguments
+    # and attr_reader/attr_accessor. This avoids the overhead of keyword argument
+    # processing in CRuby's VM (setup_parameters_complex path), yielding the
+    # fastest possible object allocation. Loc (formerly Data.define(:offset)) is
+    # inlined as a plain Integer to eliminate ~2M wrapper object allocations
+    # during indexing.
     module IR
-      # Location information for IR nodes
-      # @param line [Integer] Line number (1-indexed)
-      # @param col_range [Range] Column range
-      Loc = Data.define(:line, :col_range)
+      # Shortcut to NodeKeyGenerator for generating node hash keys
+      NodeKeyGenerator = Core::NodeKeyGenerator
+
+      # Extract the last segment of a class/module path without array allocation.
+      # Uses String#rindex instead of String#split to avoid creating intermediate arrays.
+      # @param path [String] Class path (e.g., "Admin::User", "TypeGuessr::Core::IR::LiteralNode")
+      # @return [String, nil] Last segment (e.g., "User", "LiteralNode"), the path itself if no "::", or nil if empty
+      def self.extract_last_name(path)
+        return nil if path.nil? || path.empty?
+
+        last_sep = path.rindex("::")
+        last_sep ? path[(last_sep + 2)..] : path
+      end
+
+      # Method call signature for duck typing inference
+      # @param name [Symbol] Method name
+      # @param positional_count [Integer, nil] Number of positional arguments (nil if splat used)
+      # @param keywords [Array<Symbol>] Keyword argument names
+      CalledMethod = Data.define(:name, :positional_count, :keywords) do
+        # String representation returns method name for logging/display
+        def to_s
+          name.to_s
+        end
+      end
 
       # Pretty print helper for IR nodes (Prism-style tree output)
       module TreeInspect
@@ -17,23 +49,21 @@ module TypeGuessr
 
         def tree_inspect(indent: "", last: true, root: false)
           lines = if root
-                    ["@ #{self.class.name.split("::").last} (location: #{format_loc})"]
+                    ["@ #{IR.extract_last_name(self.class.name)} (location: #{format_loc})"]
                   else
                     prefix = last ? LAST_BRANCH : BRANCH
                     indent += (last ? SPACE : PIPE)
-                    ["#{indent.delete_suffix(last ? SPACE : PIPE)}#{prefix}@ #{self.class.name.split("::").last} (location: #{format_loc})"]
+                    ["#{indent.delete_suffix(last ? SPACE : PIPE)}#{prefix}@ #{IR.extract_last_name(self.class.name)} (location: #{format_loc})"]
                   end
           lines.concat(tree_inspect_fields(indent))
           lines.join("\n")
         end
 
-        private
-
-        def format_loc
-          loc ? "(#{loc.line},#{loc.col_range.begin})-(#{loc.line},#{loc.col_range.end})" : "∅"
+        private def format_loc
+          loc ? "@#{loc}" : "∅"
         end
 
-        def tree_field(name, value, indent, last: false)
+        private def tree_field(name, value, indent, last: false)
           prefix = last ? LAST_BRANCH : BRANCH
           case value
           when nil
@@ -69,59 +99,37 @@ module TypeGuessr
           end
         end
 
-        def tree_inspect_fields(_indent)
+        private def tree_inspect_fields(_indent)
           []
         end
       end
 
-      # Base class for all IR nodes
-      # IR represents a reverse dependency graph where each node points to nodes it depends on
-      class Node
-        attr_reader :loc
-
-        def initialize(loc:)
-          @loc = loc
-        end
-
-        # Returns all nodes that this node directly depends on
-        # @return [Array<Node>]
-        def dependencies
-          []
-        end
-
-        # Generate a unique hash for this node (type + identifier + line)
-        # @return [String]
-        def node_hash
-          raise NotImplementedError, "#{self.class} must implement node_hash"
-        end
-
-        # Generate a unique key for this node (scope_id + node_hash)
-        # @param scope_id [String] The scope identifier (e.g., "User#save")
-        # @return [String]
-        def node_key(scope_id)
-          "#{scope_id}:#{node_hash}"
-        end
-      end
-
       # Literal value node
       # @param type [TypeGuessr::Core::Types::Type] The type of the literal
       # @param literal_value [Object, nil] The actual literal value (for Symbol, Integer, String)
       # @param values [Array<Node>, nil] Internal value nodes for compound literals (Hash/Array)
-      # @param loc [Loc] Location information
-      #
-      # Examples: "hello" → String, 123 → Integer, [] → Array, {} → Hash
-      # For compound literals, values contains the internal expression nodes
-      # For simple literals, literal_value stores the actual value (e.g., :a for symbols)
-      LiteralNode = Data.define(:type, :literal_value, :values, :loc) do
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class LiteralNode
         include TreeInspect
 
+        attr_reader :type, :literal_value, :values, :called_methods, :loc
+
+        def initialize(type, literal_value, values, called_methods, loc)
+          @type = type
+          @literal_value = literal_value
+          @values = values
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           values || []
         end
 
         def node_hash
-          type_name = type.is_a?(Class) ? type.name.split("::").last : type.class.name.split("::").last
-          "lit:#{type_name}:#{loc&.line}"
+          type_name = type.is_a?(Class) ? IR.extract_last_name(type.name) : IR.extract_last_name(type.class.name)
+          NodeKeyGenerator.literal(type_name, loc)
         end
 
         def node_key(scope_id)
@@ -129,11 +137,12 @@ module TypeGuessr
         end
 
         def tree_inspect_fields(indent)
-          type_str = type.respond_to?(:name) ? type.name : type.class.name.split("::").last
+          type_str = type.respond_to?(:name) ? type.name : IR.extract_last_name(type.class.name)
           [
             tree_field(:type, type_str, indent),
             tree_field(:literal_value, literal_value, indent),
-            tree_field(:values, values, indent, last: true),
+            tree_field(:values, values, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
           ]
         end
       end
@@ -142,18 +151,27 @@ module TypeGuessr
       # @param name [Symbol] Variable name
       # @param value [Node] The node of the assigned value
       # @param called_methods [Array<Symbol>] Methods called on this variable (for method-based inference)
-      # @param loc [Loc] Location information
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
       #
       # Note: called_methods is a shared array object that can be mutated during parsing
-      LocalWriteNode = Data.define(:name, :value, :called_methods, :loc) do
+      class LocalWriteNode
         include TreeInspect
 
+        attr_reader :name, :value, :called_methods, :loc
+
+        def initialize(name, value, called_methods, loc)
+          @name = name
+          @value = value
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           value ? [value] : []
         end
 
         def node_hash
-          "local_write:#{name}:#{loc&.line}"
+          NodeKeyGenerator.local_write(name, loc)
         end
 
         def node_key(scope_id)
@@ -173,18 +191,27 @@ module TypeGuessr
       # @param name [Symbol] Variable name
       # @param write_node [LocalWriteNode, nil] The LocalWriteNode this read references
       # @param called_methods [Array<Symbol>] Methods called on this variable (for method-based inference)
-      # @param loc [Loc] Location information
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
       #
       # Note: called_methods is shared with LocalWriteNode for method-based inference propagation
-      LocalReadNode = Data.define(:name, :write_node, :called_methods, :loc) do
+      class LocalReadNode
         include TreeInspect
 
+        attr_reader :name, :write_node, :called_methods, :loc
+
+        def initialize(name, write_node, called_methods, loc)
+          @name = name
+          @write_node = write_node
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           write_node ? [write_node] : []
         end
 
         def node_hash
-          "local_read:#{name}:#{loc&.line}"
+          NodeKeyGenerator.local_read(name, loc)
         end
 
         def node_key(scope_id)
@@ -205,16 +232,26 @@ module TypeGuessr
       # @param class_name [String, nil] Enclosing class name for deferred resolution
       # @param value [Node] The node of the assigned value
       # @param called_methods [Array<Symbol>] Methods called on this variable (for method-based inference)
-      # @param loc [Loc] Location information
-      InstanceVariableWriteNode = Data.define(:name, :class_name, :value, :called_methods, :loc) do
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class InstanceVariableWriteNode
         include TreeInspect
 
+        attr_reader :name, :class_name, :value, :called_methods, :loc
+
+        def initialize(name, class_name, value, called_methods, loc)
+          @name = name
+          @class_name = class_name
+          @value = value
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           value ? [value] : []
         end
 
         def node_hash
-          "ivar_write:#{name}:#{loc&.line}"
+          NodeKeyGenerator.ivar_write(name, loc)
         end
 
         def node_key(scope_id)
@@ -236,19 +273,29 @@ module TypeGuessr
       # @param class_name [String, nil] Enclosing class name for deferred resolution
       # @param write_node [InstanceVariableWriteNode, nil] The write node this read references
       # @param called_methods [Array<Symbol>] Methods called on this variable (for method-based inference)
-      # @param loc [Loc] Location information
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
       #
       # Note: write_node may be nil at conversion time if assignment appears later.
       # Resolver performs deferred lookup using class_name.
-      InstanceVariableReadNode = Data.define(:name, :class_name, :write_node, :called_methods, :loc) do
+      class InstanceVariableReadNode
         include TreeInspect
 
+        attr_reader :name, :class_name, :write_node, :called_methods, :loc
+
+        def initialize(name, class_name, write_node, called_methods, loc)
+          @name = name
+          @class_name = class_name
+          @write_node = write_node
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           write_node ? [write_node] : []
         end
 
         def node_hash
-          "ivar_read:#{name}:#{loc&.line}"
+          NodeKeyGenerator.ivar_read(name, loc)
         end
 
         def node_key(scope_id)
@@ -270,16 +317,26 @@ module TypeGuessr
       # @param class_name [String, nil] Enclosing class name for deferred resolution
       # @param value [Node] The node of the assigned value
       # @param called_methods [Array<Symbol>] Methods called on this variable (for method-based inference)
-      # @param loc [Loc] Location information
-      ClassVariableWriteNode = Data.define(:name, :class_name, :value, :called_methods, :loc) do
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class ClassVariableWriteNode
         include TreeInspect
 
+        attr_reader :name, :class_name, :value, :called_methods, :loc
+
+        def initialize(name, class_name, value, called_methods, loc)
+          @name = name
+          @class_name = class_name
+          @value = value
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           value ? [value] : []
         end
 
         def node_hash
-          "cvar_write:#{name}:#{loc&.line}"
+          NodeKeyGenerator.cvar_write(name, loc)
         end
 
         def node_key(scope_id)
@@ -301,16 +358,26 @@ module TypeGuessr
       # @param class_name [String, nil] Enclosing class name for deferred resolution
       # @param write_node [ClassVariableWriteNode, nil] The write node this read references
       # @param called_methods [Array<Symbol>] Methods called on this variable (for method-based inference)
-      # @param loc [Loc] Location information
-      ClassVariableReadNode = Data.define(:name, :class_name, :write_node, :called_methods, :loc) do
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class ClassVariableReadNode
         include TreeInspect
 
+        attr_reader :name, :class_name, :write_node, :called_methods, :loc
+
+        def initialize(name, class_name, write_node, called_methods, loc)
+          @name = name
+          @class_name = class_name
+          @write_node = write_node
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           write_node ? [write_node] : []
         end
 
         def node_hash
-          "cvar_read:#{name}:#{loc&.line}"
+          NodeKeyGenerator.cvar_read(name, loc)
         end
 
         def node_key(scope_id)
@@ -333,18 +400,28 @@ module TypeGuessr
       #                      :keyword_optional, :keyword_rest, :block, :forwarding)
       # @param default_value [Node, nil] Default value node (nil if no default)
       # @param called_methods [Array<Symbol>] Methods called on this parameter (for method-based inference)
-      # @param loc [Loc] Location information
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
       #
       # Note: called_methods is a shared array object that can be mutated during parsing
-      ParamNode = Data.define(:name, :kind, :default_value, :called_methods, :loc) do
+      class ParamNode
         include TreeInspect
 
+        attr_reader :name, :kind, :default_value, :called_methods, :loc
+
+        def initialize(name, kind, default_value, called_methods, loc)
+          @name = name
+          @kind = kind
+          @default_value = default_value
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           default_value ? [default_value] : []
         end
 
         def node_hash
-          "param:#{name}:#{loc&.line}"
+          NodeKeyGenerator.param(name, loc)
         end
 
         def node_key(scope_id)
@@ -364,16 +441,26 @@ module TypeGuessr
       # Constant reference node
       # @param name [String] Constant name (e.g., "DEFAULT_NAME", "User::ADMIN")
       # @param dependency [Node] The node where this constant is defined
-      # @param loc [Loc] Location information
-      ConstantNode = Data.define(:name, :dependency, :loc) do
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class ConstantNode
         include TreeInspect
 
+        attr_reader :name, :dependency, :called_methods, :loc
+
+        def initialize(name, dependency, called_methods, loc)
+          @name = name
+          @dependency = dependency
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           dependency ? [dependency] : []
         end
 
         def node_hash
-          "const:#{name}:#{loc&.line}"
+          NodeKeyGenerator.constant(name, loc)
         end
 
         def node_key(scope_id)
@@ -383,7 +470,8 @@ module TypeGuessr
         def tree_inspect_fields(indent)
           [
             tree_field(:name, name, indent),
-            tree_field(:dependency, dependency, indent, last: true),
+            tree_field(:dependency, dependency, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
           ]
         end
       end
@@ -395,10 +483,25 @@ module TypeGuessr
       # @param block_params [Array<BlockParamSlot>] Block parameter slots
       # @param block_body [Node, nil] Block body return node (for inferring block return type)
       # @param has_block [Boolean] Whether a block was provided (even if empty)
-      # @param loc [Loc] Location information
-      CallNode = Data.define(:method, :receiver, :args, :block_params, :block_body, :has_block, :loc) do
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class CallNode
         include TreeInspect
 
+        attr_reader :method, :receiver, :args, :called_methods, :loc
+        attr_accessor :block_params, :block_body, :has_block
+
+        def initialize(method, receiver, args, block_params, block_body, has_block, called_methods, loc)
+          @method = method
+          @receiver = receiver
+          @args = args
+          @block_params = block_params
+          @block_body = block_body
+          @has_block = has_block
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           deps = []
           deps << receiver if receiver
@@ -408,7 +511,7 @@ module TypeGuessr
         end
 
         def node_hash
-          "call:#{method}:#{loc&.line}"
+          NodeKeyGenerator.call(method, loc)
         end
 
         def node_key(scope_id)
@@ -422,7 +525,8 @@ module TypeGuessr
             tree_field(:args, args, indent),
             tree_field(:block_params, block_params, indent),
             tree_field(:block_body, block_body, indent),
-            tree_field(:has_block, has_block, indent, last: true),
+            tree_field(:has_block, has_block, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
           ]
         end
       end
@@ -431,16 +535,26 @@ module TypeGuessr
       # Represents a parameter slot in a block (e.g., |user| in users.each { |user| ... })
       # @param index [Integer] Parameter index (0-based)
       # @param call_node [CallNode] The call node this slot belongs to
-      # @param loc [Loc] Location information for the parameter itself
-      BlockParamSlot = Data.define(:index, :call_node, :loc) do
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class BlockParamSlot
         include TreeInspect
 
+        attr_reader :index, :call_node, :called_methods, :loc
+
+        def initialize(index, call_node, called_methods, loc)
+          @index = index
+          @call_node = call_node
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           [call_node]
         end
 
         def node_hash
-          "bparam:#{index}:#{loc&.line}"
+          NodeKeyGenerator.bparam(index, loc)
         end
 
         def node_key(scope_id)
@@ -450,7 +564,8 @@ module TypeGuessr
         def tree_inspect_fields(indent)
           [
             tree_field(:index, index, indent),
-            tree_field(:call_node, "(CallNode ref)", indent, last: true),
+            tree_field(:call_node, "(CallNode ref)", indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
           ]
         end
       end
@@ -459,16 +574,25 @@ module TypeGuessr
       # Represents the convergence point of multiple branches (if/else, case/when, etc.)
       # The type is the union of all branch types
       # @param branches [Array<Node>] Final nodes from each branch
-      # @param loc [Loc] Location information
-      MergeNode = Data.define(:branches, :loc) do
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class MergeNode
         include TreeInspect
 
+        attr_reader :branches, :called_methods, :loc
+
+        def initialize(branches, called_methods, loc)
+          @branches = branches
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           branches
         end
 
         def node_hash
-          "merge:#{loc&.line}"
+          NodeKeyGenerator.merge(loc)
         end
 
         def node_key(scope_id)
@@ -476,7 +600,50 @@ module TypeGuessr
         end
 
         def tree_inspect_fields(indent)
-          [tree_field(:branches, branches, indent, last: true)]
+          [
+            tree_field(:branches, branches, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
+          ]
+        end
+      end
+
+      # Short-circuit or node
+      # Represents || and ||= operations where LHS is evaluated first,
+      # and RHS is only evaluated if LHS is falsy (nil/false)
+      # @param lhs [Node] Left-hand side (evaluated first)
+      # @param rhs [Node] Right-hand side (evaluated only if LHS is falsy)
+      # @param called_methods [Array<CalledMethod>] Methods called on this node
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class OrNode
+        include TreeInspect
+
+        attr_reader :lhs, :rhs, :called_methods, :loc
+
+        def initialize(lhs, rhs, called_methods, loc)
+          @lhs = lhs
+          @rhs = rhs
+          @called_methods = called_methods
+          @loc = loc
+        end
+
+        def dependencies
+          [lhs, rhs]
+        end
+
+        def node_hash
+          NodeKeyGenerator.or_node(loc)
+        end
+
+        def node_key(scope_id)
+          "#{scope_id}:#{node_hash}"
+        end
+
+        def tree_inspect_fields(indent)
+          [
+            tree_field(:lhs, lhs, indent),
+            tree_field(:rhs, rhs, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
+          ]
         end
       end
 
@@ -486,11 +653,28 @@ module TypeGuessr
       # @param params [Array<ParamNode>] Parameter nodes
       # @param return_node [Node] Node representing the return value
       # @param body_nodes [Array<Node>] All nodes in the method body
-      # @param loc [Loc] Location information
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
       # @param singleton [Boolean] true if this is a singleton method (def self.method_name)
-      DefNode = Data.define(:name, :class_name, :params, :return_node, :body_nodes, :loc, :singleton) do
+      class DefNode
         include TreeInspect
 
+        attr_reader :name, :class_name, :params, :return_node, :body_nodes, :called_methods, :loc, :singleton,
+                    :module_function
+
+        def initialize(name, class_name, params, return_node, body_nodes, called_methods, loc, singleton,
+                       module_function: false)
+          @name = name
+          @class_name = class_name
+          @params = params
+          @return_node = return_node
+          @body_nodes = body_nodes
+          @called_methods = called_methods
+          @loc = loc
+          @singleton = singleton
+          @module_function = module_function
+        end
+
         def dependencies
           deps = params.dup
           deps << return_node if return_node
@@ -499,7 +683,7 @@ module TypeGuessr
         end
 
         def node_hash
-          "def:#{name}:#{loc&.line}"
+          NodeKeyGenerator.def_node(name, loc)
         end
 
         def node_key(scope_id)
@@ -513,6 +697,7 @@ module TypeGuessr
             tree_field(:params, params, indent),
             tree_field(:return_node, return_node, indent),
             tree_field(:body_nodes, body_nodes, indent),
+            tree_field(:called_methods, called_methods, indent),
             tree_field(:singleton, singleton, indent, last: true),
           ]
         end
@@ -521,16 +706,26 @@ module TypeGuessr
       # Class/Module node - container for methods and other definitions
       # @param name [String] Class or module name
       # @param methods [Array<DefNode>] Method definitions in this class/module
-      # @param loc [Loc] Location information
-      ClassModuleNode = Data.define(:name, :methods, :loc) do
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class ClassModuleNode
         include TreeInspect
 
+        attr_reader :name, :methods, :called_methods, :loc
+
+        def initialize(name, methods, called_methods, loc)
+          @name = name
+          @methods = methods
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           methods
         end
 
         def node_hash
-          "class:#{name}:#{loc&.line}"
+          NodeKeyGenerator.class_module(name, loc)
         end
 
         def node_key(scope_id)
@@ -540,7 +735,8 @@ module TypeGuessr
         def tree_inspect_fields(indent)
           [
             tree_field(:name, name, indent),
-            tree_field(:methods, methods, indent, last: true),
+            tree_field(:methods, methods, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
           ]
         end
       end
@@ -548,16 +744,26 @@ module TypeGuessr
       # Self reference node
       # @param class_name [String] Name of the enclosing class/module
       # @param singleton [Boolean] Whether this self is in a singleton method context
-      # @param loc [Loc] Location information
-      SelfNode = Data.define(:class_name, :singleton, :loc) do
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class SelfNode
         include TreeInspect
 
+        attr_reader :class_name, :singleton, :called_methods, :loc
+
+        def initialize(class_name, singleton, called_methods, loc)
+          @class_name = class_name
+          @singleton = singleton
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           []
         end
 
         def node_hash
-          "self:#{class_name}:#{loc&.line}"
+          NodeKeyGenerator.self_node(class_name, loc)
         end
 
         def node_key(scope_id)
@@ -567,23 +773,73 @@ module TypeGuessr
         def tree_inspect_fields(indent)
           [
             tree_field(:class_name, class_name, indent),
-            tree_field(:singleton, singleton, indent, last: true),
+            tree_field(:singleton, singleton, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
+          ]
+        end
+      end
+
+      # Type narrowing wrapper node
+      # Wraps a value node and signals that falsy types should be removed
+      # Used after guard clauses: `return unless x` → x is truthy after guard
+      # @param value [Node] The original node to narrow
+      # @param kind [Symbol] Narrowing kind (:truthy removes nil/false)
+      # @param called_methods [Array<CalledMethod>] Methods called on this node
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class NarrowNode
+        include TreeInspect
+
+        attr_reader :value, :kind, :called_methods, :loc
+
+        def initialize(value, kind, called_methods, loc)
+          @value = value
+          @kind = kind
+          @called_methods = called_methods
+          @loc = loc
+        end
+
+        def dependencies
+          value ? [value] : []
+        end
+
+        def node_hash
+          NodeKeyGenerator.narrow(kind, loc)
+        end
+
+        def node_key(scope_id)
+          "#{scope_id}:#{node_hash}"
+        end
+
+        def tree_inspect_fields(indent)
+          [
+            tree_field(:value, value, indent),
+            tree_field(:kind, kind, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
           ]
         end
       end
 
       # Explicit return statement node
       # @param value [Node, nil] Return value node (nil for bare `return`)
-      # @param loc [Loc] Location information
-      ReturnNode = Data.define(:value, :loc) do
+      # @param called_methods [Array<CalledMethod>] Methods called on this node (for method-based inference)
+      # @param loc [Integer, nil] Byte offset from start of file (0-indexed)
+      class ReturnNode
         include TreeInspect
 
+        attr_reader :value, :called_methods, :loc
+
+        def initialize(value, called_methods, loc)
+          @value = value
+          @called_methods = called_methods
+          @loc = loc
+        end
+
         def dependencies
           value ? [value] : []
         end
 
         def node_hash
-          "return:#{loc&.line}"
+          NodeKeyGenerator.return_node(loc)
         end
 
         def node_key(scope_id)
@@ -591,7 +847,10 @@ module TypeGuessr
         end
 
         def tree_inspect_fields(indent)
-          [tree_field(:value, value, indent, last: true)]
+          [
+            tree_field(:value, value, indent),
+            tree_field(:called_methods, called_methods, indent, last: true),
+          ]
         end
       end
     end