rigortype 0.2.9 → 0.3.1

data/lib/rigor/sig_gen/observation_collector.rb

@@ -6,6 +6,7 @@ require_relative "../environment"
 require_relative "../scope"
 require_relative "../type"
 require_relative "../source/literals"
+require_relative "../source/node_children"
 require_relative "../inference/scope_indexer"
 
 module Rigor
@@ -126,7 +127,7 @@ module Rigor
           end
         end
 
-        node.compact_child_nodes.each { |child| walk_class_decls(child, prefix, accumulator) }
+        node.rigor_each_child { |child| walk_class_decls(child, prefix, accumulator) }
       end
 
       def qualified_constant_path(constant_path)
@@ -146,7 +147,7 @@ module Rigor
         return unless node.is_a?(Prism::Node)
 
         record_call(node, scope_index, bindings, observations) if node.is_a?(Prism::CallNode)
-        node.compact_child_nodes.each { |child| walk_calls(child, scope_index, bindings, observations) }
+        node.rigor_each_child { |child| walk_calls(child, scope_index, bindings, observations) }
       end
 
       def record_call(call_node, scope_index, bindings, observations)
@@ -244,7 +245,7 @@ module Rigor
         recognise_describe(node, bindings)
         recognise_subject_or_let(node, bindings, scope_index)
 
-        node.compact_child_nodes.each { |child| walk_rspec_bindings(child, bindings, scope_index) }
+        node.rigor_each_child { |child| walk_rspec_bindings(child, bindings, scope_index) }
       end
 
       def recognise_describe(node, bindings)

data/lib/rigor/sig_gen/rbs_validity.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require "rbs"
+
+module Rigor
+  module SigGen
+    # Parses what sig-gen is about to emit, before it is emitted.
+    #
+    # Two rendering bugs have shipped RBS that `rbs` itself rejects — a non-identifier record key
+    # (`{ :"data-contrast" => T }`) and a `&block` constructor parameter rendered before the parens
+    # (`(**untyped, ?{ (?) -> void })`). Both were found downstream, as a poisoned `sig/` tree: the file is
+    # quarantined, its types vanish, and (for a project adopting `reject-unparseable-signatures`) the build
+    # fails. A generator that can poison its own consumer's build is worth guarding at the source.
+    #
+    # So this is not a check for ONE bug — it is the guard that turns the whole rendering-bug *class* from a
+    # silent bad artifact into a caught error naming the method that could not be rendered. It is cheap: the
+    # text is already in memory, and `RBS::Parser` is the same C parser the consumer will use, so it is exactly
+    # the oracle that matters (not an approximation of it).
+    module RbsValidity
+      # A rendered method line (`def foo: () -> String`) is a fragment, not a compilation unit — it has to be
+      # wrapped to be parsed at all. The wrapper name is irrelevant to the parse and never reaches disk.
+      PROBE_CLASS = "RigorSigGenProbe"
+
+      module_function
+
+      # @param line [String] a rendered RBS method line, e.g. `def self.parse: (String) -> Integer`.
+      # @return [String, nil] the parse error's first line, or nil when the line is valid.
+      def method_line_error(line)
+        return nil if line.nil?
+
+        source_error("class #{PROBE_CLASS}\n  #{line}\nend\n")
+      end
+
+      # @param source [String] a complete `.rbs` file's text.
+      # @return [String, nil] the parse error's first line, or nil when the file is valid.
+      def source_error(source)
+        ::RBS::Parser.parse_signature(::RBS::Buffer.new(name: "(rigor sig-gen)", content: source))
+        nil
+      rescue ::RBS::ParsingError => e
+        e.message.to_s.lines.first.to_s.strip
+      end
+    end
+  end
+end

data/lib/rigor/sig_gen/renderer.rb

@@ -126,6 +126,7 @@ module Rigor
           when :created then render_write_created(result)
           when :updated then render_write_updated(result)
           when :skipped_outside_sig_root then render_write_skipped(result)
+          when :skipped_invalid_rbs then render_write_invalid(result)
           end
         end
       end
@@ -143,6 +144,15 @@ module Rigor
         @out.puts("skipped #{result.source_path} -> #{result.target_path} (outside sig root)")
       end
 
+      # The assembled file does not parse, so it was NOT written. Writing it would poison the sig tree — the
+      # consumer quarantines an unparseable `.rbs` whole, taking every other type in that file down with it,
+      # including the user's own hand-written ones.
+      def render_write_invalid(result)
+        @out.puts("REFUSED #{result.target_path} — the generated RBS does not parse, so it was not written")
+        @out.puts("  #{result.error}")
+        @out.puts("  This is a bug in Rigor's RBS rendering, not in your code — please report it.")
+      end
+
       def render_write_json(results)
         @out.puts(JSON.pretty_generate({ results: results.map(&:to_h) }))
       end

data/lib/rigor/sig_gen/write_result.rb

@@ -10,19 +10,24 @@ module Rigor
     # - `source_path` — original `.rb` file.
     # - `target_path` — `.rbs` file the writer was responsible for (`nil` when the source path falls outside
     #   the project signature tree, in which case `action` is `:skipped_outside_sig_root`).
-    # - `action` — one of `:created` / `:updated` / `:noop` / `:skipped_outside_sig_root`.
+    # - `action` — one of `:created` / `:updated` / `:noop` / `:skipped_outside_sig_root` /
+    #   `:skipped_invalid_rbs`.
     # - `applied` — the {MethodCandidate}s that actually landed on disk.
     # - `skipped` — the {MethodCandidate}s the writer declined (e.g. tighter-return without `--overwrite`). Each
     #   entry pairs the candidate with a skip reason keyword (`:user_authored`).
+    # - `error` — the parse error, when `action` is `:skipped_invalid_rbs`: the file the writer assembled does
+    #   not parse, so it was NOT written (writing it would poison the project's sig tree — the consumer
+    #   quarantines an unparseable `.rbs`, taking every other type in that file down with it).
     class WriteResult
-      attr_reader :source_path, :target_path, :action, :applied, :skipped
+      attr_reader :source_path, :target_path, :action, :applied, :skipped, :error
 
-      def initialize(source_path:, target_path:, action:, applied: [], skipped: [])
+      def initialize(source_path:, target_path:, action:, applied: [], skipped: [], error: nil)
         @source_path = source_path
         @target_path = target_path
         @action = action
         @applied = applied.freeze
         @skipped = skipped.freeze
+        @error = error
         freeze
       end
 
@@ -33,7 +38,7 @@ module Rigor
           action: action.to_s,
           applied: applied.map(&:to_h),
           skipped: skipped.map { |c, reason| c.to_h.merge(write_skip_reason: reason.to_s) }
-        }
+        }.tap { |h| h[:error] = error if error }
       end
     end
   end

data/lib/rigor/sig_gen/writer.rb

@@ -108,8 +108,15 @@ module Rigor
       end
 
       def create_new(source_path, target, candidates)
+        content = render_new_file(candidates)
+        error = RbsValidity.source_error(content)
+        if error
+          return WriteResult.new(source_path: source_path, target_path: target,
+                                 action: :skipped_invalid_rbs, error: error)
+        end
+
         FileUtils.mkdir_p(target.dirname)
-        target.write(render_new_file(candidates))
+        target.write(content)
         WriteResult.new(source_path: source_path, target_path: target,
                         action: :created, applied: candidates)
       end
@@ -272,35 +279,65 @@ module Rigor
         return WriteResult.new(source_path: source_path, target_path: target, action: :noop) if decls.nil?
 
         state = MergeState.new(source: source, decls: decls, applied: [], skipped: [])
-        supers = merged_superclasses(candidates)
-        candidates.group_by(&:class_name).each { |class_name, methods| merge_class(state, class_name, methods, supers) }
-        merge_class_shells(state, collect_class_shells(candidates), merged_namespace_kinds(candidates))
+        merge_candidates(state, candidates)
 
         action = state.applied.empty? ? :noop : :updated
-        target.write(state.source) if action == :updated
+        unless action == :updated
+          return WriteResult.new(source_path: source_path, target_path: target, action: action,
+                                 applied: state.applied, skipped: state.skipped)
+        end
+
+        # The merge splices text into an existing file by byte offset, so a bug here can produce a file neither
+        # the generator nor the target was responsible for. Parse the assembled result and refuse to write it if
+        # it is broken: a `.rbs` the consumer cannot parse is quarantined whole, so a bad splice would delete
+        # every type in the file — including the user's own, which were fine before we touched them.
+        error = RbsValidity.source_error(state.source)
+        if error
+          return WriteResult.new(source_path: source_path, target_path: target,
+                                 action: :skipped_invalid_rbs, error: error)
+        end
+
+        target.write(state.source)
         WriteResult.new(source_path: source_path, target_path: target,
                         action: action, applied: state.applied, skipped: state.skipped)
       end
 
+      # Applies every class group, then every requested shell, then normalises the layout the three steps
+      # leave behind. Each step mutates `state` in place and re-parses, so the next one sees current offsets.
+      def merge_candidates(state, candidates)
+        supers = merged_superclasses(candidates)
+        kinds = merged_namespace_kinds(candidates)
+        shells = collect_class_shells(candidates)
+        groups = candidates.group_by(&:class_name)
+        # Ancestors first, so `Foo` is created (or found) before `Foo::Bar` looks for a parent to nest under.
+        # A plain sort suffices: a name always sorts before every name it is a strict prefix of.
+        groups.keys.sort.each { |name| merge_class(state, name, groups.fetch(name), kinds, supers) }
+        merge_class_shells(state, shells, kinds, supers)
+        collapse_nested_declarations(state, groups.keys + shells.to_a)
+      end
+
       # ADR-14 gap-#3 (e): for every requested class shell that isn't already declared in the target file,
       # insert an empty `class Const\nend` block inside the nearest existing ancestor. Shells already covered by
       # an existing declaration are silently a no-op. The `applied` accumulator does NOT grow — shells are
       # structural declarations, not methods, so the action-count surface (`updated +N`) keeps reflecting method
       # changes only.
-      def merge_class_shells(state, shells, kinds)
+      def merge_class_shells(state, shells, kinds, supers)
         shells.each do |qualified|
           next if find_class_decl(state.decls, qualified)
 
-          insert_class_shell(state, qualified, kinds)
+          insert_namespace_chain(state, qualified, kinds, supers, [])
         end
       end
 
-      def insert_class_shell(state, qualified, kinds)
+      # Splices the missing part of `qualified`'s namespace chain into the nearest declaration the file
+      # already has for one of its ancestors, or at top level when it has none. `methods` land on the leaf
+      # node, so this is the single insertion path for both an empty class shell and a brand-new class.
+      def insert_namespace_chain(state, qualified, kinds, supers, methods)
         segments = qualified.split("::")
         anchor_segs, missing = split_at_existing_ancestor(state.decls, segments)
         anchor_decl = anchor_segs.empty? ? nil : find_class_decl(state.decls, anchor_segs.join("::"))
-        depth = anchor_decl ? anchor_decl_indent_depth(anchor_decl) : 0
-        snippet = build_shell_snippet(missing, anchor_segs, kinds, depth)
+        depth = anchor_decl ? member_indent_depth(anchor_decl) : 0
+        snippet = render_chain_snippet(missing, anchor_segs, kinds, supers, depth, methods)
         state.source = if anchor_decl
                          insert_before_end(state.source, anchor_decl, snippet)
                        else
@@ -309,6 +346,113 @@ module Rigor
         state.decls = parse_signature(state.source) || state.decls
       end
 
+      # ADR-14 gap-#3 follow-up (c), update half: the create path folds a strict-prefix pair into one nested
+      # tree, but a file written before that fix — or by hand — can still carry the flat sibling layout
+      # (`class Foo` next to a top-level `class Foo::Bar`). After merging, relocate each declaration this run
+      # touched underneath its parent's declaration when the same file holds both, so an update converges on
+      # the same canonical layout a fresh generation would produce.
+      #
+      # Scope is deliberately the touched names only: an unrelated flat pair elsewhere in the file is none of
+      # sig-gen's business, and rewriting it would be a layout change the user never asked for. Shallowest
+      # first, so `Foo::Bar` has already moved under `Foo` by the time `Foo::Bar::Baz` looks for its parent.
+      def collapse_nested_declarations(state, names)
+        names.uniq.sort.each { |name| relocate_under_parent(state, name) }
+      end
+
+      def relocate_under_parent(state, qualified)
+        segments = qualified.split("::")
+        return if segments.size < 2
+
+        parent_name = segments[0...-1].join("::")
+        parent = find_class_decl(state.decls, parent_name)
+        decl = parent && find_class_decl(state.decls, qualified)
+        return if decl.nil? || overlapping?(parent, decl)
+
+        region = decl_region(state.source, decl)
+        block = region && relocated_block(state.source, decl, parent, segments.last, region)
+        return if block.nil?
+
+        apply_relocation(state, region, parent_name, block)
+      end
+
+      # Cuts the declaration's region out, re-parses so the parent's byte range reflects the removal, and
+      # splices the re-indented block back in as the parent's last member. Any step that cannot be carried
+      # out cleanly (an unparseable intermediate, a parent that vanished) abandons the move and leaves the
+      # merged source exactly as it was — a flat layout is cosmetic, a mangled `.rbs` is not.
+      def apply_relocation(state, region, parent_name, block)
+        source = splice_out(state.source, region)
+        decls = parse_signature(source)
+        anchor = decls && find_class_decl(decls, parent_name)
+        return if anchor.nil?
+
+        state.source = insert_before_end(source, anchor, block)
+        state.decls = parse_signature(state.source) || state.decls
+      end
+
+      # True when the two declarations' source ranges are not disjoint — either one already nests the other,
+      # or the file's shape is one this pass does not understand. Both are reasons not to move anything.
+      def overlapping?(one, other)
+        one.location.start_pos < other.location.end_pos && other.location.start_pos < one.location.end_pos
+      end
+
+      # The declaration's full source region, extended to cover its leading comment and annotations and to
+      # end just past the newline that closes it. Returns `nil` when either edge shares a line with something
+      # else, because cutting there would move — or strand — text the declaration does not own.
+      def decl_region(source, decl)
+        start_pos = ([decl.location.start_pos] + leading_positions(decl)).min
+        line_start = line_start_index(source, start_pos)
+        return nil unless blank_range?(source, line_start, start_pos)
+
+        finish = decl.location.end_pos
+        line_end = source.index("\n", finish) || source.size
+        return nil unless blank_range?(source, finish, line_end)
+
+        line_start...(line_end < source.size ? line_end + 1 : line_end)
+      end
+
+      def leading_positions(decl)
+        positions = decl.annotations.filter_map { |a| a.location&.start_pos }
+        comment_location = decl.comment&.location
+        positions << comment_location.start_pos if comment_location
+        positions
+      end
+
+      # The moved text, with its compact `Foo::Bar` head shortened to `Bar` and every line pushed in to the
+      # parent's member depth. The name's own sub-location drives the rewrite, so a superclass, type
+      # parameters, and the whole body survive byte-for-byte.
+      def relocated_block(source, decl, parent, short_name, region)
+        name_location = decl.location[:name]
+        return nil if name_location.nil?
+
+        text = source[region].to_s
+        head = name_location.start_pos - region.begin
+        tail = name_location.end_pos - region.begin
+        renamed = text[0...head].to_s + short_name + text[tail..].to_s
+        reindent(renamed, member_indent_depth(parent) - decl_indent_depth(decl))
+      end
+
+      def reindent(text, delta)
+        return text unless delta.positive?
+
+        prefix = INDENT * delta
+        text.lines.map { |line| line.match?(/\A\s*\z/) ? line : prefix + line }.join
+      end
+
+      # Removes `region`, collapsing the newline run left behind at the seam so the cut never shows up as a
+      # widening gap (or a trailing blank line at EOF) in the file it edited. Spacing that was already fine
+      # is left exactly as the user wrote it.
+      def splice_out(source, region)
+        before = source[0...region.begin].to_s
+        after = source[region.end..].to_s
+        return after.sub(/\A\n+/, "") if before.match?(/\A\s*\z/)
+        return before.sub(/\n{2,}\z/, "\n") if after.match?(/\A\s*\z/)
+
+        seam = before[/\n+\z/].to_s.size + after[/\A\n+/].to_s.size
+        return before + after if seam <= 2
+
+        "#{before.sub(/\n+\z/, "\n\n")}#{after.sub(/\A\n+/, '')}"
+      end
+
       def split_at_existing_ancestor(decls, segments)
         (segments.size - 1).downto(0).each do |i|
           ancestor = segments[0...i].join("::")
@@ -317,37 +461,61 @@ module Rigor
         [[], segments]
       end
 
-      # Pulls the indent depth (in `INDENT` units) one level deeper than the anchor decl's own column.
-      # Pre-existing members might be missing (an empty `class Foo; end`) so the keyword column is the robust
-      # signal.
-      def anchor_decl_indent_depth(decl)
-        decl_column = decl.location[:keyword].start_column
-        (decl_column / INDENT.size) + 1
+      # The indent depth (in `INDENT` units) a member of `decl` sits at: one level deeper than the
+      # declaration's own keyword column. Pre-existing members might be missing (an empty `class Foo; end`)
+      # so the keyword column is the robust signal.
+      def member_indent_depth(decl)
+        decl_indent_depth(decl) + 1
+      end
+
+      def decl_indent_depth(decl)
+        decl.location[:keyword].start_column / INDENT.size
       end
 
-      def build_shell_snippet(missing, anchor_segs, kinds, depth)
+      # Renders the `missing` segment chain as one nested block, carrying `methods` on its leaf, by handing a
+      # synthesised tree node to the create path's renderer. Both paths therefore agree on the keyword, the
+      # superclass suffix, and the indentation of every level. A leaf with no methods is a class shell, which
+      # is exactly the create path's `shell:` node (gap-#3 (e)).
+      def render_chain_snippet(missing, anchor_segs, kinds, supers, depth, methods)
         return "" if missing.empty?
 
-        head, *rest = missing
-        qualified = (anchor_segs + [head]).join("::")
-        indent = INDENT * depth
-        if rest.empty?
-          keyword = kinds[qualified] || :class
-          "#{indent}#{keyword} #{head}\n#{indent}end\n"
-        else
-          inner = build_shell_snippet(rest, anchor_segs + [head], kinds, depth + 1)
-          keyword = kinds[qualified] || :module
-          "#{indent}#{keyword} #{head}\n#{inner}#{indent}end\n"
+        node = missing.reverse.each_with_index.inject(nil) do |child, (segment, index)|
+          { name: segment, children: child ? { child[:name] => child } : {},
+            methods: index.zero? ? methods : [], shell: index.zero? && methods.empty? }
         end
+        render_tree_node(node, kinds, supers, depth, anchor_segs)
       end
 
+      # Splices `snippet` in just before the declaration's closing `end`. When that `end` starts its own line
+      # the insertion point moves to the START of that line: the snippet carries its own indentation, and
+      # anchoring at the keyword would otherwise donate the `end`'s indent to the snippet's first line and
+      # strand the `end` in column zero.
       def insert_before_end(source, decl, snippet)
+        return source if snippet.empty?
+
         end_pos = decl.location[:end].start_pos
-        source[0...end_pos] + snippet + source[end_pos..]
+        line_start = line_start_index(source, end_pos)
+        return source[0...line_start] + snippet + source[line_start..] if blank_range?(source, line_start, end_pos)
+
+        "#{source[0...end_pos]}\n#{snippet}#{source[end_pos..]}"
       end
 
+      # Appends a top-level block, separated from whatever precedes it by exactly one blank line.
       def append_top_level(source, snippet)
-        ends_with_newline?(source) ? source + snippet : "#{source}\n#{snippet}"
+        return snippet if source.empty?
+
+        base = ends_with_newline?(source) ? source : "#{source}\n"
+        base.end_with?("\n\n") ? base + snippet : "#{base}\n#{snippet}"
+      end
+
+      def line_start_index(source, pos)
+        return 0 unless pos.positive?
+
+        (source.rindex("\n", pos - 1) || -1) + 1
+      end
+
+      def blank_range?(source, from, to)
+        source[from...to].to_s.match?(/\A[ \t]*\z/)
       end
 
       def parse_signature(source)
@@ -357,14 +525,15 @@ module Rigor
         nil
       end
 
-      def merge_class(state, class_name, methods, supers = {})
+      def merge_class(state, class_name, methods, kinds, supers)
         decl = find_class_decl(state.decls, class_name)
-        state.source = if decl.nil?
-                         append_new_class(state.source, class_name, methods, state.applied, supers[class_name])
-                       else
-                         merge_into_existing_class(state.source, decl, methods, state.applied, state.skipped)
-                       end
-        state.decls = parse_signature(state.source) || state.decls
+        if decl.nil?
+          state.applied.concat(methods)
+          insert_namespace_chain(state, class_name, kinds, supers, methods)
+        else
+          state.source = merge_into_existing_class(state.source, decl, methods, state.applied, state.skipped)
+          state.decls = parse_signature(state.source) || state.decls
+        end
       end
 
       # Walks the parsed decl tree recursively, tracking the enclosing module/class prefix, and returns the
@@ -389,16 +558,6 @@ module Rigor
         nil
       end
 
-      # Appends an entirely new `class Foo … end` block at the end of the file (with a leading blank line as
-      # separator).
-      def append_new_class(source, class_name, methods, applied, superclass = nil)
-        body = methods.map { |c| "#{INDENT}#{c.rbs}" }.join("\n")
-        header = superclass ? "class #{class_name} < #{superclass}" : "class #{class_name}"
-        snippet = "\n#{header}\n#{body}\nend\n"
-        applied.concat(methods)
-        ends_with_newline?(source) ? source + snippet : "#{source}\n#{snippet}"
-      end
-
       def ends_with_newline?(source)
         source.end_with?("\n")
       end
@@ -455,9 +614,8 @@ module Rigor
       def insert_into_class(source, decl, new_methods)
         return source if new_methods.empty?
 
-        end_pos = decl.location[:end].start_pos
-        addition = new_methods.map { |c| "#{INDENT}#{c.rbs}\n" }.join
-        source[0...end_pos] + addition + source[end_pos..]
+        indent = INDENT * member_indent_depth(decl)
+        insert_before_end(source, decl, new_methods.map { |c| "#{indent}#{c.rbs}\n" }.join)
       end
 
       # Walks the class's existing method declarations; for each replaceable candidate that matches a member

data/lib/rigor/source/node_children.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+require "prism"
+
+module Rigor
+  module Source
+    # Allocation-free replacement for the `node.compact_child_nodes.each { … }` idiom.
+    #
+    # `Prism::Node#compact_child_nodes` allocates a fresh Array on every call — for the ~43 leaf node classes
+    # (`IntegerNode`, `LocalVariableReadNode`, `NilNode`, …) it is literally `def compact_child_nodes; []; end`.
+    # Rigor's tree walkers call it unconditionally on every node of every walk, so one full-tree walk allocates one
+    # Array per node visited. On leaf-heavy sources (a Ragel-generated parser has hundreds of thousands of
+    # integer-literal leaves) these throwaway arrays are the single largest allocation source in a run — over half
+    # of all allocations on mail's `lib`.
+    #
+    # Loading this file compiles a `#rigor_each_child` method onto every `Prism::*Node` class (additive reopening of
+    # a dependency's class, the `Cache::RbsEnvironmentMarshalPatch` precedent; the `rigor_` prefix keeps it
+    # collision-free). It yields the same child nodes, in the same order, without materialising the intermediate
+    # array: each child-bearing field is read directly in field declaration order — exactly the order
+    # `compact_child_nodes` emits — and a `NodeListField`'s already-materialised array is iterated in place (the
+    # reader returns the stored array, not a copy). Nil optional children and nil list elements are skipped,
+    # mirroring `compact_child_nodes`'s "compact" semantics. Leaf classes compile to an empty method.
+    #
+    # A method compiled onto the node class is the fastest dispatch available for this shape — one virtual send on
+    # the receiver, the same mechanism `compact_child_nodes` itself uses, with zero allocation. The reflective
+    # alternatives were measured and rejected: a `public_send`-per-field loop is ~45 % *slower* than
+    # `compact_child_nodes` on a full-tree walk (dynamic dispatch per field swamps the allocation win), and a
+    # central generated-method table dispatched via `Module#send` + a per-class Hash lookup still regressed
+    # Rails-corpus wall ~5 % under YJIT. Walkers therefore call `node.rigor_each_child { … }` directly.
+    #
+    # The field map is derived once at load from `Prism::Reflection`, so it tracks whatever `prism` version resolves
+    # at runtime (ADR-79) instead of a hand-maintained table. `spec/rigor/source/node_children_spec.rb` is the
+    # binding contract: over a corpus of real source it asserts — for every node reached — that both
+    # `#rigor_each_child` and {each_child} yield output element-for-element identical (object identity and order) to
+    # `compact_child_nodes`.
+    module NodeChildren
+      module_function
+
+      # The concrete `Prism::*Node` classes (`< Prism::Node`).
+      NODE_CLASSES =
+        Prism.constants.grep(/Node\z/).filter_map do |name|
+          const = Prism.const_get(name)
+          const if const.is_a?(Class) && const < Prism::Node
+        end.freeze
+
+      # Shared frozen empty entry for leaf classes. Identity (`equal?`) distinguishes a leaf entry from a
+      # childless-in-practice non-leaf.
+      EMPTY = [].freeze
+
+      # node class => frozen Array of `[reader_symbol, :node | :node_optional | :list]` pairs, in field declaration
+      # order (the order `compact_child_nodes` emits its children). Leaf classes map to {EMPTY}. This is the source
+      # of truth the per-class methods are compiled from, exposed for introspection and the equivalence spec.
+      CHILD_READERS =
+        NODE_CLASSES.each_with_object({}) do |klass, map|
+          pairs = Prism::Reflection.fields_for(klass).filter_map do |field|
+            case field
+            when Prism::Reflection::OptionalNodeField
+              [field.name, :node_optional].freeze
+            when Prism::Reflection::NodeField
+              [field.name, :node].freeze
+            when Prism::Reflection::NodeListField
+              [field.name, :list].freeze
+            end
+          end
+          map[klass] = pairs.empty? ? EMPTY : pairs.freeze
+        end.freeze
+
+      # The leaf classes — no child-bearing field, so `compact_child_nodes` is always `[]` and `#rigor_each_child`
+      # compiles to an empty method. Exposed for the equivalence spec and introspection.
+      LEAF_CLASSES = NODE_CLASSES.select { |klass| CHILD_READERS[klass].equal?(EMPTY) }.to_set.freeze
+
+      # Codegen (not per-call reflection) so iteration reads fields directly at hand-written speed — see the module
+      # comment for the measurements ruling out the reflective forms. Each statement mirrors `compact_child_nodes`'s
+      # own form for that field kind exactly: a required `NodeField` is yielded unconditionally (`compact <<
+      # predicate`), an `OptionalNodeField` behind a truthiness guard (`compact << receiver if receiver`), and a
+      # `NodeListField`'s stored array is iterated unfiltered (`compact.concat(requireds)`) — no extra nil checks
+      # beyond what Prism itself performs. The explicit `self.` receiver keeps a hypothetical future field named
+      # `child` from parsing as the just-assigned local. For a `CallNode` (three optional-node fields) the emitted
+      # source is:
+      #
+      #   def rigor_each_child
+      #     child = self.receiver
+      #     yield child if child
+      #     child = self.arguments
+      #     yield child if child
+      #     child = self.block
+      #     yield child if child
+      #   end
+      NODE_CLASSES.each do |klass|
+        statements = CHILD_READERS[klass].map do |reader, kind|
+          case kind
+          when :list          then "self.#{reader}.each { |child| yield child }"
+          when :node_optional then "child = self.#{reader}\n  yield child if child"
+          else                     "yield self.#{reader}"
+          end
+        end
+        klass.class_eval(<<~RUBY, __FILE__, __LINE__ + 1) # rubocop:disable Style/DocumentDynamicEvalDefinition
+          def rigor_each_child
+            #{statements.join("\n  ")}
+          end
+        RUBY
+      end
+
+      # Yield each direct child `Prism::Node` of `node` in `compact_child_nodes` order, without allocating an
+      # intermediate array. Nil / non-node input yields nothing (an optional field can be nil, so a caller may pass
+      # e.g. `node.body`). Pure and re-entrant; `break` / `next` / `return` in the block behave exactly as they
+      # would with `compact_child_nodes.each`. Walkers holding a known `Prism::Node` call `#rigor_each_child`
+      # directly; this wrapper is the nil-tolerant entry.
+      #
+      # @yieldparam child [Prism::Node]
+      def each_child(node, &)
+        node.rigor_each_child(&) if node.is_a?(Prism::Node)
+      end
+    end
+  end
+end

data/lib/rigor/source/node_locator.rb

@@ -2,6 +2,8 @@
 
 require "prism"
 
+require_relative "node_children"
+
 module Rigor
   module Source
     # Locates the deepest Prism AST node enclosing a given source position.
@@ -83,7 +85,7 @@ module Rigor
         return nil unless node.is_a?(Prism::Node)
         return nil unless contains?(node, offset)
 
-        node.compact_child_nodes.each do |child|
+        node.rigor_each_child do |child|
           deeper = descend(child, offset)
           return deeper if deeper
         end

data/lib/rigor/source/node_walker.rb

@@ -2,6 +2,8 @@
 
 require "prism"
 
+require_relative "node_children"
+
 module Rigor
   module Source
     # Yields every `Prism::Node` reachable from a root in DFS pre-order.
@@ -28,7 +30,7 @@ module Rigor
         return unless node.is_a?(Prism::Node)
 
         yield node
-        node.compact_child_nodes.each { |child| walk(child, &) }
+        node.rigor_each_child { |child| walk(child, &) }
       end
 
       # Like {.each}, but also yields the node's lexical ancestor chain (outermost first, EXCLUDING the node
@@ -51,7 +53,7 @@ module Rigor
 
         block.call(node, ancestors)
         ancestors.push(node)
-        node.compact_child_nodes.each { |child| walk_with_ancestors(child, ancestors, &block) }
+        node.rigor_each_child { |child| walk_with_ancestors(child, ancestors, &block) }
         ancestors.pop
       end
     end

data/lib/rigor/source.rb

@@ -10,6 +10,7 @@ module Rigor
   end
 end
 
+require_relative "source/node_children"
 require_relative "source/node_locator"
 require_relative "source/node_walker"
 require_relative "source/literals"