From 47433a16ec64bf7a4cac2cf024fbe35415809f5c Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 03:23:27 +0200
Subject: [PATCH 01/17] deferred token stream

---
 lib/rdoc/generator/markup.rb             |  5 +-
 lib/rdoc/parser/ruby.rb                  | 17 ++---
 lib/rdoc/parser/ruby_colorizer.rb        | 48 ++++++++++++--
 lib/rdoc/token_stream.rb                 | 11 ++--
 test/rdoc/code_object/any_method_test.rb | 18 ++++++
 test/rdoc/parser/ruby_colorizer_test.rb  | 82 ++++++++++++++++++++++++
 test/rdoc/parser/ruby_test.rb            | 14 ++++
 test/rdoc/rdoc_rubygems_hook_test.rb     |  9 ++-
 test/rdoc/rdoc_token_stream_test.rb      | 78 +++++++++++++++++++++-
 9 files changed, 254 insertions(+), 28 deletions(-)

diff --git a/lib/rdoc/generator/markup.rb b/lib/rdoc/generator/markup.rb
index d6a56adf97..6a2b235b5e 100644
--- a/lib/rdoc/generator/markup.rb
+++ b/lib/rdoc/generator/markup.rb
@@ -138,9 +138,10 @@ def add_location_comment(src)
   # Prepends line numbers if +options.line_numbers+ is true.
 
   def markup_code
-    return '' if !@token_stream
+    tokens = token_stream
+    return '' if !tokens
 
-    src = RDoc::TokenStream.to_html @token_stream
+    src = RDoc::TokenStream.to_html tokens
 
     # dedent the source
     common_indent = src.length
diff --git a/lib/rdoc/parser/ruby.rb b/lib/rdoc/parser/ruby.rb
index 63bec37426..a3e7a10bf8 100644
--- a/lib/rdoc/parser/ruby.rb
+++ b/lib/rdoc/parser/ruby.rb
@@ -251,11 +251,8 @@ def record_location(container) # :nodoc:
 
   def scan
     @lines = @content.lines
-    result = Prism.parse_lex(@content)
-    @program_node, unordered_tokens = result.value
-    # Heredoc tokens are not in start_offset order.
-    # Need to sort them to use bsearch for finding tokens from location.
-    @prism_tokens = unordered_tokens.map(&:first).sort_by { |t| t.location.start_offset }
+    result = Prism.parse(@content)
+    @program_node = result.value
     @line_nodes = {}
     prepare_line_nodes(@program_node)
     prepare_comments(result.comments)
@@ -367,10 +364,9 @@ def parse_comment_tomdoc(container, comment, line_no, start_line)
     meth.call_seq = signature
     return unless meth.name
 
-    meth.start_collecting_tokens(:ruby)
     node = @line_nodes[line_no]
     tokens = node ? syntax_highlighted_tokens(node) : []
-    tokens.each { |token| meth.token_stream << token }
+    meth.start_collecting_tokens(:ruby, tokens)
 
     container.add_method meth
     meth.comment = comment
@@ -556,7 +552,7 @@ def extract_section_comment(comment_text, prefix_line_count) # :nodoc:
   # Returns syntax highlighted tokens of the given node
 
   def syntax_highlighted_tokens(node)
-    RDoc::Parser::RubyColorizer.partial_colorize(@content, node, @prism_tokens)
+    RDoc::Parser::RubyColorizer.deferred_token_stream(@content, node)
   end
 
   # Handles `public :foo, :bar` `private :foo, :bar` and `protected :foo, :bar`
@@ -750,10 +746,7 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
     meth.block_params ||= block_params if block_params
     meth.type_signature_lines = type_signature_lines
     record_location(meth)
-    meth.start_collecting_tokens(:ruby)
-    tokens.each do |token|
-      meth.token_stream << token
-    end
+    meth.start_collecting_tokens(:ruby, tokens)
 
     # Rename after add_method to register duplicated 'new' and 'initialize'
     # defined in c and ruby.
diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index fdfbfe0b23..7d94d98fb3 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -10,6 +10,32 @@ module RDoc::Parser::RubyColorizer
 
   ColoredToken = Struct.new(:kind, :text)
 
+  # Defers Ruby source colorization until a generator requests the tokens.
+  class DeferredTokenStream
+    def initialize(source, start_column)
+      @source = source
+      @start_column = start_column
+      @tokens = nil
+    end
+
+    def materialize
+      return @tokens if @tokens
+
+      source = @source
+      result = Prism.parse_lex(source)
+      program_node, unordered_tokens = result.value
+      prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
+      node = program_node.statements.body.first.breadth_first_search do |candidate|
+        candidate.location.start_offset == 0 && (candidate.is_a?(Prism::DefNode) || candidate.is_a?(Prism::CallNode))
+      end
+      tokens = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
+      tokens.unshift(ColoredToken.new(:plain, ' ' * @start_column)) if @start_column > 0
+      @tokens = tokens
+      @source = nil
+      @tokens
+    end
+  end
+
   # Prism operator token types except assignment '='
   OP_TOKENS = %i[
     AMPERSAND AMPERSAND_AMPERSAND
@@ -57,6 +83,18 @@ def colorize(code)
       partial_colorize(code, program_node, prism_tokens, 0, code.bytesize)
     end
 
+    # Returns a token stream that retains only the source needed to colorize
+    # +node+ when first accessed.
+    def deferred_token_stream(whole_code, node)
+      visitor = NodeColorizeVisitor.new(false)
+      node.accept(visitor)
+      start_offset = node.location.start_offset
+      end_offset = [node.location.end_offset, visitor.effective_end_offset].max
+      source = String.new(capacity: end_offset - start_offset, encoding: whole_code.encoding)
+      source << whole_code.byteslice(start_offset...end_offset)
+      DeferredTokenStream.new(source, node.location.start_column)
+    end
+
     # Colorize partial +node+ in +whole_code+ and returns colored token stream.
     def partial_colorize(whole_code, node, prism_tokens, start_offset = nil, end_offset = nil)
       start_offset ||= node.location.start_offset
@@ -133,10 +171,11 @@ def normal_tokens(tokens)
   # Visitor to determine node colorizing which can't be determined by tokens.
   # STRING_CONTENT/EMBEXPR_BEGIN/EMBEXPR_END in string/regexp/symbol have different colorizing
   class NodeColorizeVisitor < Prism::Visitor # :nodoc:
-    attr_reader :tokens
+    attr_reader :effective_end_offset, :tokens
 
-    def initialize
-      @tokens = []
+    def initialize(collect_tokens = true)
+      @effective_end_offset = 0
+      @tokens = [] if collect_tokens
     end
 
     def visit_symbol_node(node)
@@ -221,7 +260,8 @@ def visit_def_node(node)
     def push_location(kind, location)
       # Only push tokens that have a non-zero length
       if location && location.start_offset < location.end_offset
-        @tokens << [kind, location.start_offset, location.end_offset]
+        @effective_end_offset = location.end_offset if location.end_offset > @effective_end_offset
+        @tokens << [kind, location.start_offset, location.end_offset] if @tokens
       end
     end
 
diff --git a/lib/rdoc/token_stream.rb b/lib/rdoc/token_stream.rb
index 71bd4a7078..bd0fabf172 100644
--- a/lib/rdoc/token_stream.rb
+++ b/lib/rdoc/token_stream.rb
@@ -50,22 +50,22 @@ def self.to_html(token_stream)
   # Adds +tokens+ to the collected tokens
 
   def add_tokens(tokens)
-    @token_stream.concat(tokens)
+    token_stream.concat(tokens)
   end
 
   ##
   # Adds one +token+ to the collected tokens
 
   def add_token(token)
-    @token_stream.push(token)
+    token_stream.push(token)
   end
 
   ##
   # Starts collecting tokens
   #
 
-  def collect_tokens(language)
-    @token_stream = []
+  def collect_tokens(language, initial_stream = [])
+    @token_stream = initial_stream
     @token_stream_language = language
   end
 
@@ -75,13 +75,14 @@ def collect_tokens(language)
   # Remove the last token from the collected tokens
 
   def pop_token
-    @token_stream.pop
+    token_stream.pop
   end
 
   ##
   # Current token stream
 
   def token_stream
+    @token_stream = @token_stream.materialize if @token_stream.respond_to?(:materialize)
     @token_stream
   end
 
diff --git a/test/rdoc/code_object/any_method_test.rb b/test/rdoc/code_object/any_method_test.rb
index 3f3b681407..318e4fce61 100644
--- a/test/rdoc/code_object/any_method_test.rb
+++ b/test/rdoc/code_object/any_method_test.rb
@@ -202,6 +202,24 @@ def test_marshal_dump
     assert_equal section,        loaded.section
   end
 
+  def test_marshal_dump_does_not_materialize_deferred_tokens
+    source = 'def method; :unique_deferred_body; end'
+    node = Prism.parse(source).value.statements.body.first
+    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(source, node)
+    method = RDoc::AnyMethod.new 'method'
+    top_level = @store.add_file('deferred.rb')
+    method.record_location top_level
+    top_level.add_class(RDoc::ClassModule, 'Deferred').add_method(method)
+    method.start_collecting_tokens(:ruby, deferred)
+
+    dump = Marshal.dump(method)
+
+    assert_nil deferred.instance_variable_get(:@tokens)
+    assert_equal source, deferred.instance_variable_get(:@source)
+    refute_includes dump, 'unique_deferred_body'
+    assert_nil Marshal.load(dump).token_stream
+  end
+
   def test_marshal_dump_with_type_signature
     @store.path = Dir.tmpdir
     top_level = @store.add_file 'file.rb'
diff --git a/test/rdoc/parser/ruby_colorizer_test.rb b/test/rdoc/parser/ruby_colorizer_test.rb
index 9b7ce92211..79647b13e2 100644
--- a/test/rdoc/parser/ruby_colorizer_test.rb
+++ b/test/rdoc/parser/ruby_colorizer_test.rb
@@ -7,6 +7,88 @@ def token(kind, text)
     RDoc::Parser::RubyColorizer::ColoredToken.new(kind, text)
   end
 
+  def assert_deferred_matches_eager(code)
+    parse_result = Prism.parse_lex(code)
+    program_node, unordered_tokens = parse_result.value
+    prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
+    node = block_given? ? yield(program_node) : program_node.statements.body.first
+    eager = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
+    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, node)
+
+    assert_equal eager.map { |token| [token.kind, token.text] }, deferred.materialize.map { |token| [token.kind, token.text] }
+    assert_same deferred.materialize, deferred.materialize
+    assert_nil deferred.instance_variable_get(:@source)
+    deferred
+  end
+
+  def test_deferred_token_stream
+    assert_deferred_matches_eager <<~'RUBY'
+        def ordinary(value = '💎')
+          "string#{value}"
+          /regexp#{value}/
+          :"symbol#{value}"
+        end
+    RUBY
+
+    assert_deferred_matches_eager <<~RUBY
+      def endless = 42
+    RUBY
+  end
+
+  def test_deferred_token_stream_with_heredocs
+    assert_deferred_matches_eager <<~'RUBY'
+      def heredocs
+        string = <<~TEXT
+          string #{value}
+        TEXT
+        command = <<~`COMMAND`
+          echo value
+        COMMAND
+        [string, command]
+      end
+    RUBY
+
+    assert_deferred_matches_eager <<~'RUBY'
+      def endless = <<~TEXT
+        value
+      TEXT
+    RUBY
+
+    assert_deferred_matches_eager <<~'RUBY'
+      method <<~ONE, <<~TWO
+        one
+      ONE
+        two
+      TWO
+    RUBY
+  end
+
+  def test_deferred_token_stream_with_heredoc_and_sibling_statement
+    code = <<~'RUBY'
+      method <<~TEXT; sibling
+        value
+      TEXT
+    RUBY
+
+    assert_deferred_matches_eager code
+
+    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, Prism.parse(code).value.statements.body.first)
+    sibling = deferred.materialize.find { |token| token.text == 'sibling' }
+    assert_equal :plain, sibling.kind
+  end
+
+  def test_deferred_token_stream_with_heredoc_in_parent_expression
+    code = <<~'RUBY'
+      method(<<~TEXT) && sibling
+        value
+      TEXT
+    RUBY
+
+    deferred = assert_deferred_matches_eager(code) { |program| program.statements.body.first.left }
+    sibling = deferred.materialize.find { |token| token.text == 'sibling' }
+    assert_equal :plain, sibling.kind
+  end
+
   def test_partial_colorize
     code = <<~RUBY
       class A
diff --git a/test/rdoc/parser/ruby_test.rb b/test/rdoc/parser/ruby_test.rb
index 3ebfd5891b..9c3c90954d 100644
--- a/test/rdoc/parser/ruby_test.rb
+++ b/test/rdoc/parser/ruby_test.rb
@@ -782,6 +782,9 @@ class Foo
     assert_equal @top_level, one.file
     assert_equal @top_level, two.file
     assert_equal @top_level, three.file
+    assert_equal [], one.token_stream
+    assert_equal [], two.token_stream
+    assert_equal [], three.token_stream
   end
 
   def test_invalid_meta_method
@@ -953,6 +956,7 @@ class Foo
     assert_equal 'my method', method.comment.text.strip
     assert_equal 4, method.line
     assert_equal @top_level, method.file
+    assert_equal '  add_my_method :method_foo, :arg', method.token_stream.map(&:text).join
   end
 
   def test_first_comment_is_not_a_meta_method
@@ -1358,6 +1362,11 @@ def m3; end
     assert_equal [:private, :public, :private, :private], instance_methods.map(&:visibility)
     assert_equal ['m1', 'm3', 'm4'], singleton_methods.map(&:name)
     assert_equal [:public, :public, :public], singleton_methods.map(&:visibility)
+
+    instance_m4 = instance_methods.last
+    singleton_m4 = singleton_methods.last
+    assert_same instance_m4.instance_variable_get(:@token_stream), singleton_m4.instance_variable_get(:@token_stream)
+    assert_same instance_m4.token_stream, singleton_m4.token_stream
   end
 
   def test_class_method_visibility
@@ -2711,6 +2720,8 @@ def foo
     RUBY
 
     foo, bar = @top_level.classes.first.method_list
+    assert_kind_of RDoc::Parser::RubyColorizer::DeferredTokenStream, foo.instance_variable_get(:@token_stream)
+    assert_kind_of RDoc::Parser::RubyColorizer::DeferredTokenStream, bar.instance_variable_get(:@token_stream)
     assert_equal(['  ', 'def', ' ', 'foo', "\n", '    ', '42', "\n", '  ', 'end'], foo.token_stream.map(&:text))
     assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
   end
@@ -2764,6 +2775,8 @@ class C
         #
         # field - A field name.
 
+        find_by_fields :dynamic
+
       end
     RUBY
 
@@ -2783,6 +2796,7 @@ class C
     expected.file = @top_level
 
     assert_equal expected, m.comment.parse
+    assert_equal '  find_by_fields :dynamic', m.token_stream.map(&:text).join
   end
 
   def test_tomdoc_postprocess
diff --git a/test/rdoc/rdoc_rubygems_hook_test.rb b/test/rdoc/rdoc_rubygems_hook_test.rb
index bb122f0e41..99fc91b52e 100644
--- a/test/rdoc/rdoc_rubygems_hook_test.rb
+++ b/test/rdoc/rdoc_rubygems_hook_test.rb
@@ -11,7 +11,7 @@ def setup
       s.platform    = Gem::Platform::RUBY
       s.name        = "a"
       s.version     = 2
-      s.rdoc_options = %w[--main MyTitle]
+      s.rdoc_options = %w[--main MyTitle --format=ri]
       s.extra_rdoc_files = %w[README]
     end
     @tempdir = File.realpath(Dir.mktmpdir("test_rubygems_hook_"))
@@ -41,7 +41,9 @@ def setup
     File.open(File.join(@tempdir, 'a-2', 'lib', 'a.rb'), 'w') do |f|
       f.puts '# comment'
       f.puts '# :include: include.txt'
-      f.puts 'class A; end'
+      f.puts 'class A'
+      f.puts '  def visible_method; :visible_body; end'
+      f.puts 'end'
     end
     File.open(File.join(@tempdir, 'a-2', 'include.txt'), 'w') do |f|
       f.puts 'included content'
@@ -136,6 +138,8 @@ def test_generate_all
 
     assert @hook.rdoc_installed?
     assert @hook.ri_installed?
+    assert_path_exist File.join(@a.doc_dir('ri'), 'cache.ri')
+    assert_include File.read(File.join(@a.doc_dir('rdoc'), 'A.html')), 'visible_body'
 
     rdoc = @hook.instance_variable_get :@rdoc
 
@@ -242,6 +246,7 @@ def test_generate_no_overwrite
   end
 
   def test_generate_with_ri_opt
+    @a.rdoc_options.delete '--format=ri'
     @a.rdoc_options << '--ri'
     FileUtils.mkdir_p @a.doc_dir
     FileUtils.mkdir_p File.join(@a.gem_dir, 'lib')
diff --git a/test/rdoc/rdoc_token_stream_test.rb b/test/rdoc/rdoc_token_stream_test.rb
index 254811c72a..f02096d543 100644
--- a/test/rdoc/rdoc_token_stream_test.rb
+++ b/test/rdoc/rdoc_token_stream_test.rb
@@ -3,6 +3,26 @@
 
 class RDocTokenStreamTest < RDoc::TestCase
 
+  class DeferredStream
+    attr_reader :materializations
+
+    def initialize(tokens)
+      @tokens = tokens
+      @materializations = 0
+    end
+
+    def materialize
+      @materializations += 1
+      @tokens
+    end
+  end
+
+  def token_collector
+    Class.new do
+      include RDoc::TokenStream
+    end.new
+  end
+
   def test_class_to_html
     tokens = [
       { kind: :constant, text: 'CONSTANT' },
@@ -78,13 +98,56 @@ def test_add_token
   end
 
   def test_collect_tokens
-    foo = Class.new do
-      include RDoc::TokenStream
-    end.new
+    foo = token_collector
     foo.collect_tokens(:ruby)
     assert_equal [], foo.token_stream
   end
 
+  def test_collect_deferred_tokens
+    tokens = [{ kind: :identifier, text: 'foo' }]
+    deferred = DeferredStream.new(tokens)
+    foo = token_collector
+
+    foo.collect_tokens(:ruby, deferred)
+
+    assert_equal 0, deferred.materializations
+    assert_same tokens, foo.token_stream
+    assert_same tokens, foo.token_stream
+    assert_equal 1, deferred.materializations
+  end
+
+  def test_collect_existing_token_array
+    tokens = [{ kind: :identifier, text: 'foo' }]
+    foo = token_collector
+
+    foo.collect_tokens(:c, tokens)
+
+    assert_same tokens, foo.token_stream
+  end
+
+  def test_mutating_deferred_tokens
+    deferred = DeferredStream.new([:first])
+    foo = token_collector
+    foo.collect_tokens(:ruby, deferred)
+    foo.add_token(:second)
+    assert_equal [:first, :second], foo.token_stream
+    assert_equal 1, deferred.materializations
+
+    deferred = DeferredStream.new([:first])
+    foo = token_collector
+    foo.collect_tokens(:ruby, deferred)
+    foo.add_tokens([:second, :third])
+    assert_equal [:first, :second, :third], foo.token_stream
+    assert_equal 1, deferred.materializations
+
+    deferred = DeferredStream.new([:first])
+    foo = token_collector
+    foo.collect_tokens(:ruby, deferred)
+    assert_equal :first, foo.pop_token
+    assert_equal [], foo.token_stream
+    assert_equal 1, deferred.materializations
+  end
+
   def test_pop_token
     foo = Class.new do
       include RDoc::TokenStream
@@ -126,4 +189,13 @@ def initialize
     end.new
     assert_equal "", foo.tokens_to_s
   end
+
+  def test_tokens_to_s_with_deferred_tokens
+    foo = token_collector
+    deferred = DeferredStream.new([{ kind: :identifier, text: 'foo' }])
+    foo.collect_tokens(:ruby, deferred)
+
+    assert_equal 'foo', foo.tokens_to_s
+    assert_equal 1, deferred.materializations
+  end
 end

From b054da9d8f567afca8d770033b15fed982957b18 Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 12:54:29 +0200
Subject: [PATCH 02/17] minimize tests

---
 lib/rdoc/token_stream.rb                 |  6 +-
 test/rdoc/code_object/any_method_test.rb | 18 ------
 test/rdoc/parser/ruby_colorizer_test.rb  | 80 +++---------------------
 test/rdoc/parser/ruby_test.rb            | 13 ----
 test/rdoc/rdoc_rubygems_hook_test.rb     |  9 +--
 test/rdoc/rdoc_token_stream_test.rb      | 78 ++---------------------
 6 files changed, 18 insertions(+), 186 deletions(-)

diff --git a/lib/rdoc/token_stream.rb b/lib/rdoc/token_stream.rb
index bd0fabf172..274f171183 100644
--- a/lib/rdoc/token_stream.rb
+++ b/lib/rdoc/token_stream.rb
@@ -50,14 +50,14 @@ def self.to_html(token_stream)
   # Adds +tokens+ to the collected tokens
 
   def add_tokens(tokens)
-    token_stream.concat(tokens)
+    @token_stream.concat(tokens)
   end
 
   ##
   # Adds one +token+ to the collected tokens
 
   def add_token(token)
-    token_stream.push(token)
+    @token_stream.push(token)
   end
 
   ##
@@ -75,7 +75,7 @@ def collect_tokens(language, initial_stream = [])
   # Remove the last token from the collected tokens
 
   def pop_token
-    token_stream.pop
+    @token_stream.pop
   end
 
   ##
diff --git a/test/rdoc/code_object/any_method_test.rb b/test/rdoc/code_object/any_method_test.rb
index 318e4fce61..3f3b681407 100644
--- a/test/rdoc/code_object/any_method_test.rb
+++ b/test/rdoc/code_object/any_method_test.rb
@@ -202,24 +202,6 @@ def test_marshal_dump
     assert_equal section,        loaded.section
   end
 
-  def test_marshal_dump_does_not_materialize_deferred_tokens
-    source = 'def method; :unique_deferred_body; end'
-    node = Prism.parse(source).value.statements.body.first
-    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(source, node)
-    method = RDoc::AnyMethod.new 'method'
-    top_level = @store.add_file('deferred.rb')
-    method.record_location top_level
-    top_level.add_class(RDoc::ClassModule, 'Deferred').add_method(method)
-    method.start_collecting_tokens(:ruby, deferred)
-
-    dump = Marshal.dump(method)
-
-    assert_nil deferred.instance_variable_get(:@tokens)
-    assert_equal source, deferred.instance_variable_get(:@source)
-    refute_includes dump, 'unique_deferred_body'
-    assert_nil Marshal.load(dump).token_stream
-  end
-
   def test_marshal_dump_with_type_signature
     @store.path = Dir.tmpdir
     top_level = @store.add_file 'file.rb'
diff --git a/test/rdoc/parser/ruby_colorizer_test.rb b/test/rdoc/parser/ruby_colorizer_test.rb
index 79647b13e2..1977a31377 100644
--- a/test/rdoc/parser/ruby_colorizer_test.rb
+++ b/test/rdoc/parser/ruby_colorizer_test.rb
@@ -7,86 +7,20 @@ def token(kind, text)
     RDoc::Parser::RubyColorizer::ColoredToken.new(kind, text)
   end
 
-  def assert_deferred_matches_eager(code)
-    parse_result = Prism.parse_lex(code)
-    program_node, unordered_tokens = parse_result.value
-    prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
-    node = block_given? ? yield(program_node) : program_node.statements.body.first
-    eager = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
-    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, node)
-
-    assert_equal eager.map { |token| [token.kind, token.text] }, deferred.materialize.map { |token| [token.kind, token.text] }
-    assert_same deferred.materialize, deferred.materialize
-    assert_nil deferred.instance_variable_get(:@source)
-    deferred
-  end
-
   def test_deferred_token_stream
-    assert_deferred_matches_eager <<~'RUBY'
-        def ordinary(value = '💎')
-          "string#{value}"
-          /regexp#{value}/
-          :"symbol#{value}"
-        end
-    RUBY
-
-    assert_deferred_matches_eager <<~RUBY
-      def endless = 42
-    RUBY
-  end
-
-  def test_deferred_token_stream_with_heredocs
-    assert_deferred_matches_eager <<~'RUBY'
-      def heredocs
-        string = <<~TEXT
-          string #{value}
-        TEXT
-        command = <<~`COMMAND`
-          echo value
-        COMMAND
-        [string, command]
-      end
-    RUBY
-
-    assert_deferred_matches_eager <<~'RUBY'
-      def endless = <<~TEXT
-        value
-      TEXT
-    RUBY
-
-    assert_deferred_matches_eager <<~'RUBY'
-      method <<~ONE, <<~TWO
-        one
-      ONE
-        two
-      TWO
-    RUBY
-  end
-
-  def test_deferred_token_stream_with_heredoc_and_sibling_statement
-    code = <<~'RUBY'
-      method <<~TEXT; sibling
-        value
-      TEXT
-    RUBY
-
-    assert_deferred_matches_eager code
-
-    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, Prism.parse(code).value.statements.body.first)
-    sibling = deferred.materialize.find { |token| token.text == 'sibling' }
-    assert_equal :plain, sibling.kind
-  end
-
-  def test_deferred_token_stream_with_heredoc_in_parent_expression
     code = <<~'RUBY'
       method(<<~TEXT) && sibling
         value
       TEXT
     RUBY
+    program_node, unordered_tokens = Prism.parse_lex(code).value
+    node = program_node.statements.body.first.left
+    prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
+    expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
+    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, node)
 
-    deferred = assert_deferred_matches_eager(code) { |program| program.statements.body.first.left }
-    sibling = deferred.materialize.find { |token| token.text == 'sibling' }
-    assert_equal :plain, sibling.kind
+    assert_equal expected, deferred.materialize
+    assert_same deferred.materialize, deferred.materialize
   end
 
   def test_partial_colorize
diff --git a/test/rdoc/parser/ruby_test.rb b/test/rdoc/parser/ruby_test.rb
index 9c3c90954d..07ad499651 100644
--- a/test/rdoc/parser/ruby_test.rb
+++ b/test/rdoc/parser/ruby_test.rb
@@ -782,9 +782,6 @@ class Foo
     assert_equal @top_level, one.file
     assert_equal @top_level, two.file
     assert_equal @top_level, three.file
-    assert_equal [], one.token_stream
-    assert_equal [], two.token_stream
-    assert_equal [], three.token_stream
   end
 
   def test_invalid_meta_method
@@ -956,7 +953,6 @@ class Foo
     assert_equal 'my method', method.comment.text.strip
     assert_equal 4, method.line
     assert_equal @top_level, method.file
-    assert_equal '  add_my_method :method_foo, :arg', method.token_stream.map(&:text).join
   end
 
   def test_first_comment_is_not_a_meta_method
@@ -1362,11 +1358,6 @@ def m3; end
     assert_equal [:private, :public, :private, :private], instance_methods.map(&:visibility)
     assert_equal ['m1', 'm3', 'm4'], singleton_methods.map(&:name)
     assert_equal [:public, :public, :public], singleton_methods.map(&:visibility)
-
-    instance_m4 = instance_methods.last
-    singleton_m4 = singleton_methods.last
-    assert_same instance_m4.instance_variable_get(:@token_stream), singleton_m4.instance_variable_get(:@token_stream)
-    assert_same instance_m4.token_stream, singleton_m4.token_stream
   end
 
   def test_class_method_visibility
@@ -2721,7 +2712,6 @@ def foo
 
     foo, bar = @top_level.classes.first.method_list
     assert_kind_of RDoc::Parser::RubyColorizer::DeferredTokenStream, foo.instance_variable_get(:@token_stream)
-    assert_kind_of RDoc::Parser::RubyColorizer::DeferredTokenStream, bar.instance_variable_get(:@token_stream)
     assert_equal(['  ', 'def', ' ', 'foo', "\n", '    ', '42', "\n", '  ', 'end'], foo.token_stream.map(&:text))
     assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
   end
@@ -2775,8 +2765,6 @@ class C
         #
         # field - A field name.
 
-        find_by_fields :dynamic
-
       end
     RUBY
 
@@ -2796,7 +2784,6 @@ class C
     expected.file = @top_level
 
     assert_equal expected, m.comment.parse
-    assert_equal '  find_by_fields :dynamic', m.token_stream.map(&:text).join
   end
 
   def test_tomdoc_postprocess
diff --git a/test/rdoc/rdoc_rubygems_hook_test.rb b/test/rdoc/rdoc_rubygems_hook_test.rb
index 99fc91b52e..bb122f0e41 100644
--- a/test/rdoc/rdoc_rubygems_hook_test.rb
+++ b/test/rdoc/rdoc_rubygems_hook_test.rb
@@ -11,7 +11,7 @@ def setup
       s.platform    = Gem::Platform::RUBY
       s.name        = "a"
       s.version     = 2
-      s.rdoc_options = %w[--main MyTitle --format=ri]
+      s.rdoc_options = %w[--main MyTitle]
       s.extra_rdoc_files = %w[README]
     end
     @tempdir = File.realpath(Dir.mktmpdir("test_rubygems_hook_"))
@@ -41,9 +41,7 @@ def setup
     File.open(File.join(@tempdir, 'a-2', 'lib', 'a.rb'), 'w') do |f|
       f.puts '# comment'
       f.puts '# :include: include.txt'
-      f.puts 'class A'
-      f.puts '  def visible_method; :visible_body; end'
-      f.puts 'end'
+      f.puts 'class A; end'
     end
     File.open(File.join(@tempdir, 'a-2', 'include.txt'), 'w') do |f|
       f.puts 'included content'
@@ -138,8 +136,6 @@ def test_generate_all
 
     assert @hook.rdoc_installed?
     assert @hook.ri_installed?
-    assert_path_exist File.join(@a.doc_dir('ri'), 'cache.ri')
-    assert_include File.read(File.join(@a.doc_dir('rdoc'), 'A.html')), 'visible_body'
 
     rdoc = @hook.instance_variable_get :@rdoc
 
@@ -246,7 +242,6 @@ def test_generate_no_overwrite
   end
 
   def test_generate_with_ri_opt
-    @a.rdoc_options.delete '--format=ri'
     @a.rdoc_options << '--ri'
     FileUtils.mkdir_p @a.doc_dir
     FileUtils.mkdir_p File.join(@a.gem_dir, 'lib')
diff --git a/test/rdoc/rdoc_token_stream_test.rb b/test/rdoc/rdoc_token_stream_test.rb
index f02096d543..71ac7fc923 100644
--- a/test/rdoc/rdoc_token_stream_test.rb
+++ b/test/rdoc/rdoc_token_stream_test.rb
@@ -3,26 +3,6 @@
 
 class RDocTokenStreamTest < RDoc::TestCase
 
-  class DeferredStream
-    attr_reader :materializations
-
-    def initialize(tokens)
-      @tokens = tokens
-      @materializations = 0
-    end
-
-    def materialize
-      @materializations += 1
-      @tokens
-    end
-  end
-
-  def token_collector
-    Class.new do
-      include RDoc::TokenStream
-    end.new
-  end
-
   def test_class_to_html
     tokens = [
       { kind: :constant, text: 'CONSTANT' },
@@ -98,54 +78,16 @@ def test_add_token
   end
 
   def test_collect_tokens
-    foo = token_collector
+    foo = Class.new do
+      include RDoc::TokenStream
+    end.new
     foo.collect_tokens(:ruby)
     assert_equal [], foo.token_stream
-  end
-
-  def test_collect_deferred_tokens
-    tokens = [{ kind: :identifier, text: 'foo' }]
-    deferred = DeferredStream.new(tokens)
-    foo = token_collector
-
-    foo.collect_tokens(:ruby, deferred)
-
-    assert_equal 0, deferred.materializations
-    assert_same tokens, foo.token_stream
-    assert_same tokens, foo.token_stream
-    assert_equal 1, deferred.materializations
-  end
-
-  def test_collect_existing_token_array
-    tokens = [{ kind: :identifier, text: 'foo' }]
-    foo = token_collector
-
-    foo.collect_tokens(:c, tokens)
-
-    assert_same tokens, foo.token_stream
-  end
-
-  def test_mutating_deferred_tokens
-    deferred = DeferredStream.new([:first])
-    foo = token_collector
-    foo.collect_tokens(:ruby, deferred)
-    foo.add_token(:second)
-    assert_equal [:first, :second], foo.token_stream
-    assert_equal 1, deferred.materializations
 
-    deferred = DeferredStream.new([:first])
-    foo = token_collector
+    deferred = Object.new
+    def deferred.materialize = [:token]
     foo.collect_tokens(:ruby, deferred)
-    foo.add_tokens([:second, :third])
-    assert_equal [:first, :second, :third], foo.token_stream
-    assert_equal 1, deferred.materializations
-
-    deferred = DeferredStream.new([:first])
-    foo = token_collector
-    foo.collect_tokens(:ruby, deferred)
-    assert_equal :first, foo.pop_token
-    assert_equal [], foo.token_stream
-    assert_equal 1, deferred.materializations
+    assert_equal [:token], foo.token_stream
   end
 
   def test_pop_token
@@ -190,12 +132,4 @@ def initialize
     assert_equal "", foo.tokens_to_s
   end
 
-  def test_tokens_to_s_with_deferred_tokens
-    foo = token_collector
-    deferred = DeferredStream.new([{ kind: :identifier, text: 'foo' }])
-    foo.collect_tokens(:ruby, deferred)
-
-    assert_equal 'foo', foo.tokens_to_s
-    assert_equal 1, deferred.materializations
-  end
 end

From d444b1ac2e8ea998231dc3a620e628a4f6529bd8 Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 13:28:50 +0200
Subject: [PATCH 03/17] simplify code

---
 lib/rdoc/generator/markup.rb  | 3 +--
 test/rdoc/parser/ruby_test.rb | 1 -
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/lib/rdoc/generator/markup.rb b/lib/rdoc/generator/markup.rb
index 6a2b235b5e..0fee0da33a 100644
--- a/lib/rdoc/generator/markup.rb
+++ b/lib/rdoc/generator/markup.rb
@@ -138,8 +138,7 @@ def add_location_comment(src)
   # Prepends line numbers if +options.line_numbers+ is true.
 
   def markup_code
-    tokens = token_stream
-    return '' if !tokens
+    return '' if !(tokens = token_stream)
 
     src = RDoc::TokenStream.to_html tokens
 
diff --git a/test/rdoc/parser/ruby_test.rb b/test/rdoc/parser/ruby_test.rb
index 07ad499651..3ebfd5891b 100644
--- a/test/rdoc/parser/ruby_test.rb
+++ b/test/rdoc/parser/ruby_test.rb
@@ -2711,7 +2711,6 @@ def foo
     RUBY
 
     foo, bar = @top_level.classes.first.method_list
-    assert_kind_of RDoc::Parser::RubyColorizer::DeferredTokenStream, foo.instance_variable_get(:@token_stream)
     assert_equal(['  ', 'def', ' ', 'foo', "\n", '    ', '42', "\n", '  ', 'end'], foo.token_stream.map(&:text))
     assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
   end

From be2e4a3eabe3caf00707f3755055638d52dfe5d9 Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 13:35:50 +0200
Subject: [PATCH 04/17] defer code parsing completely, not partially

---
 lib/rdoc/parser/ruby_colorizer.rb       | 35 +++++++++++--------------
 test/rdoc/parser/ruby_colorizer_test.rb | 22 +++++++++++++---
 2 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index 7d94d98fb3..f61d3b0a60 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -12,9 +12,11 @@ module RDoc::Parser::RubyColorizer
 
   # Defers Ruby source colorization until a generator requests the tokens.
   class DeferredTokenStream
-    def initialize(source, start_column)
+    def initialize(source, node)
       @source = source
-      @start_column = start_column
+      @node_type = node.type
+      @start_offset = node.location.start_offset
+      @end_offset = node.location.end_offset
       @tokens = nil
     end
 
@@ -25,11 +27,12 @@ def materialize
       result = Prism.parse_lex(source)
       program_node, unordered_tokens = result.value
       prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
-      node = program_node.statements.body.first.breadth_first_search do |candidate|
-        candidate.location.start_offset == 0 && (candidate.is_a?(Prism::DefNode) || candidate.is_a?(Prism::CallNode))
+      node = program_node.breadth_first_search do |candidate|
+        candidate.type == @node_type &&
+          candidate.location.start_offset == @start_offset &&
+          candidate.location.end_offset == @end_offset
       end
       tokens = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
-      tokens.unshift(ColoredToken.new(:plain, ' ' * @start_column)) if @start_column > 0
       @tokens = tokens
       @source = nil
       @tokens
@@ -83,16 +86,10 @@ def colorize(code)
       partial_colorize(code, program_node, prism_tokens, 0, code.bytesize)
     end
 
-    # Returns a token stream that retains only the source needed to colorize
-    # +node+ when first accessed.
+    # Returns a token stream that defers parsing and lexing +whole_code+ until
+    # +node+ is first accessed.
     def deferred_token_stream(whole_code, node)
-      visitor = NodeColorizeVisitor.new(false)
-      node.accept(visitor)
-      start_offset = node.location.start_offset
-      end_offset = [node.location.end_offset, visitor.effective_end_offset].max
-      source = String.new(capacity: end_offset - start_offset, encoding: whole_code.encoding)
-      source << whole_code.byteslice(start_offset...end_offset)
-      DeferredTokenStream.new(source, node.location.start_column)
+      DeferredTokenStream.new(whole_code, node)
     end
 
     # Colorize partial +node+ in +whole_code+ and returns colored token stream.
@@ -171,11 +168,10 @@ def normal_tokens(tokens)
   # Visitor to determine node colorizing which can't be determined by tokens.
   # STRING_CONTENT/EMBEXPR_BEGIN/EMBEXPR_END in string/regexp/symbol have different colorizing
   class NodeColorizeVisitor < Prism::Visitor # :nodoc:
-    attr_reader :effective_end_offset, :tokens
+    attr_reader :tokens
 
-    def initialize(collect_tokens = true)
-      @effective_end_offset = 0
-      @tokens = [] if collect_tokens
+    def initialize
+      @tokens = []
     end
 
     def visit_symbol_node(node)
@@ -260,8 +256,7 @@ def visit_def_node(node)
     def push_location(kind, location)
       # Only push tokens that have a non-zero length
       if location && location.start_offset < location.end_offset
-        @effective_end_offset = location.end_offset if location.end_offset > @effective_end_offset
-        @tokens << [kind, location.start_offset, location.end_offset] if @tokens
+        @tokens << [kind, location.start_offset, location.end_offset]
       end
     end
 
diff --git a/test/rdoc/parser/ruby_colorizer_test.rb b/test/rdoc/parser/ruby_colorizer_test.rb
index 1977a31377..bc2ca471ce 100644
--- a/test/rdoc/parser/ruby_colorizer_test.rb
+++ b/test/rdoc/parser/ruby_colorizer_test.rb
@@ -9,12 +9,14 @@ def token(kind, text)
 
   def test_deferred_token_stream
     code = <<~'RUBY'
-      method(<<~TEXT) && sibling
-        value
-      TEXT
+      first(<<~ONE); second(<<~TWO) && sibling
+        one
+      ONE
+        two
+      TWO
     RUBY
     program_node, unordered_tokens = Prism.parse_lex(code).value
-    node = program_node.statements.body.first.left
+    node = program_node.statements.body.last.left
     prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
     expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
     deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, node)
@@ -23,6 +25,18 @@ def test_deferred_token_stream
     assert_same deferred.materialize, deferred.materialize
   end
 
+  def test_deferred_token_stream_preserves_lexical_scope
+    code = "x = 1\ny = 2\ni = 3\nx /y/i\n"
+    program_node, unordered_tokens = Prism.parse_lex(code).value
+    node = program_node.statements.body.last
+    prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
+    expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
+
+    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, node)
+
+    assert_equal expected, deferred.materialize
+  end
+
   def test_partial_colorize
     code = <<~RUBY
       class A

From 32b2c54906058832422de8fd0fb79ce2b8e1c0f5 Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 13:38:31 +0200
Subject: [PATCH 05/17] Routed all TokenStream mutators through token_stream

Mutating token_stream before materialization will raise NoMethodError. To avoid that we're routing all mutators through token_stream, this ensure that object is always materialized before mutation.
---
 lib/rdoc/token_stream.rb            |  6 +++---
 test/rdoc/rdoc_token_stream_test.rb | 15 +++++++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/lib/rdoc/token_stream.rb b/lib/rdoc/token_stream.rb
index 274f171183..bd0fabf172 100644
--- a/lib/rdoc/token_stream.rb
+++ b/lib/rdoc/token_stream.rb
@@ -50,14 +50,14 @@ def self.to_html(token_stream)
   # Adds +tokens+ to the collected tokens
 
   def add_tokens(tokens)
-    @token_stream.concat(tokens)
+    token_stream.concat(tokens)
   end
 
   ##
   # Adds one +token+ to the collected tokens
 
   def add_token(token)
-    @token_stream.push(token)
+    token_stream.push(token)
   end
 
   ##
@@ -75,7 +75,7 @@ def collect_tokens(language, initial_stream = [])
   # Remove the last token from the collected tokens
 
   def pop_token
-    @token_stream.pop
+    token_stream.pop
   end
 
   ##
diff --git a/test/rdoc/rdoc_token_stream_test.rb b/test/rdoc/rdoc_token_stream_test.rb
index 71ac7fc923..70be2b0abd 100644
--- a/test/rdoc/rdoc_token_stream_test.rb
+++ b/test/rdoc/rdoc_token_stream_test.rb
@@ -100,6 +100,21 @@ def test_pop_token
     assert_equal [], foo.token_stream
   end
 
+  def test_mutating_deferred_tokens
+    foo = Class.new do
+      include RDoc::TokenStream
+    end.new
+    deferred = Object.new
+    def deferred.materialize = [:first]
+    foo.collect_tokens(:ruby, deferred)
+
+    foo.add_token(:second)
+    foo.add_tokens([:third])
+
+    assert_equal :third, foo.pop_token
+    assert_equal [:first, :second], foo.token_stream
+  end
+
   def test_token_stream
     foo = Class.new do
       include RDoc::TokenStream

From 56b4d94301babb650a087f7a924ae2c9921080eb Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 14:14:13 +0200
Subject: [PATCH 06/17] file-scoped deferred colorization with one parse,
 node-ID dispatch, source release, and thread safety

---
 lib/rdoc/parser/ruby.rb                 |  3 +-
 lib/rdoc/parser/ruby_colorizer.rb       | 76 +++++++++++++++++--------
 test/rdoc/parser/ruby_colorizer_test.rb |  6 +-
 test/rdoc/parser/ruby_test.rb           | 12 +++-
 4 files changed, 68 insertions(+), 29 deletions(-)

diff --git a/lib/rdoc/parser/ruby.rb b/lib/rdoc/parser/ruby.rb
index a3e7a10bf8..259adc6450 100644
--- a/lib/rdoc/parser/ruby.rb
+++ b/lib/rdoc/parser/ruby.rb
@@ -143,6 +143,7 @@ def initialize(top_level, content, options, stats)
     @token_listeners = nil
     content = RDoc::Encoding.remove_magic_comment content
     @content = content
+    @colorizer_context = RDoc::Parser::RubyColorizer::DeferredContext.new(content)
     @markup = @options.markup
     @track_visibility = :nodoc != @options.visibility
     @encoding = @options.encoding
@@ -552,7 +553,7 @@ def extract_section_comment(comment_text, prefix_line_count) # :nodoc:
   # Returns syntax highlighted tokens of the given node
 
   def syntax_highlighted_tokens(node)
-    RDoc::Parser::RubyColorizer.deferred_token_stream(@content, node)
+    @colorizer_context.deferred_token_stream(node)
   end
 
   # Handles `public :foo, :bar` `private :foo, :bar` and `protected :foo, :bar`
diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index f61d3b0a60..1a773d546f 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -10,33 +10,67 @@ module RDoc::Parser::RubyColorizer
 
   ColoredToken = Struct.new(:kind, :text)
 
-  # Defers Ruby source colorization until a generator requests the tokens.
+  # A token stream populated by its file's DeferredContext.
   class DeferredTokenStream
-    def initialize(source, node)
-      @source = source
-      @node_type = node.type
-      @start_offset = node.location.start_offset
-      @end_offset = node.location.end_offset
+    #: (DeferredContext) -> void
+    def initialize(context)
+      @context = context
       @tokens = nil
     end
 
+    #: () -> Array[ColoredToken]
     def materialize
       return @tokens if @tokens
 
-      source = @source
-      result = Prism.parse_lex(source)
-      program_node, unordered_tokens = result.value
-      prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
-      node = program_node.breadth_first_search do |candidate|
-        candidate.type == @node_type &&
-          candidate.location.start_offset == @start_offset &&
-          candidate.location.end_offset == @end_offset
-      end
-      tokens = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
-      @tokens = tokens
-      @source = nil
+      @context&.materialize
       @tokens
     end
+
+    #: (Array[ColoredToken]) -> void
+    def resolve(tokens)
+      @tokens = tokens
+      @context = nil
+    end
+  end
+
+  # Defers colorization for all nodes in one source file until first access.
+  class DeferredContext
+    #: (String) -> void
+    def initialize(source)
+      @source = source
+      @streams = {}
+      @mutex = Mutex.new
+    end
+
+    #: (Prism::Node) -> DeferredTokenStream
+    def deferred_token_stream(node)
+      stream = DeferredTokenStream.new(self)
+      (@streams[node.node_id] ||= []) << stream
+      stream
+    end
+
+    #: () -> void
+    def materialize
+      @mutex.synchronize do
+        return unless @source
+
+        source = @source
+        program_node, unordered_tokens = Prism.parse_lex(source).value
+        prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
+        nodes = [program_node]
+        until nodes.empty? || @streams.empty?
+          node = nodes.pop
+          if (streams = @streams.delete(node.node_id))
+            tokens = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
+            streams.each { |stream| stream.resolve(tokens) }
+          end
+          nodes.concat(node.compact_child_nodes)
+        end
+      ensure
+        @source = nil
+        @streams = nil
+      end
+    end
   end
 
   # Prism operator token types except assignment '='
@@ -86,12 +120,6 @@ def colorize(code)
       partial_colorize(code, program_node, prism_tokens, 0, code.bytesize)
     end
 
-    # Returns a token stream that defers parsing and lexing +whole_code+ until
-    # +node+ is first accessed.
-    def deferred_token_stream(whole_code, node)
-      DeferredTokenStream.new(whole_code, node)
-    end
-
     # Colorize partial +node+ in +whole_code+ and returns colored token stream.
     def partial_colorize(whole_code, node, prism_tokens, start_offset = nil, end_offset = nil)
       start_offset ||= node.location.start_offset
diff --git a/test/rdoc/parser/ruby_colorizer_test.rb b/test/rdoc/parser/ruby_colorizer_test.rb
index bc2ca471ce..942401f3aa 100644
--- a/test/rdoc/parser/ruby_colorizer_test.rb
+++ b/test/rdoc/parser/ruby_colorizer_test.rb
@@ -19,7 +19,8 @@ def test_deferred_token_stream
     node = program_node.statements.body.last.left
     prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
     expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
-    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, node)
+    context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
+    deferred = context.deferred_token_stream(node)
 
     assert_equal expected, deferred.materialize
     assert_same deferred.materialize, deferred.materialize
@@ -32,7 +33,8 @@ def test_deferred_token_stream_preserves_lexical_scope
     prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
     expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
 
-    deferred = RDoc::Parser::RubyColorizer.deferred_token_stream(code, node)
+    context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
+    deferred = context.deferred_token_stream(node)
 
     assert_equal expected, deferred.materialize
   end
diff --git a/test/rdoc/parser/ruby_test.rb b/test/rdoc/parser/ruby_test.rb
index 3ebfd5891b..9d4bd51743 100644
--- a/test/rdoc/parser/ruby_test.rb
+++ b/test/rdoc/parser/ruby_test.rb
@@ -2711,8 +2711,16 @@ def foo
     RUBY
 
     foo, bar = @top_level.classes.first.method_list
-    assert_equal(['  ', 'def', ' ', 'foo', "\n", '    ', '42', "\n", '  ', 'end'], foo.token_stream.map(&:text))
-    assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
+    parse_lex_calls = 0
+    trace = TracePoint.new(:call, :c_call) do |event|
+      parse_lex_calls += 1 if event.self.equal?(Prism) && event.method_id == :parse_lex
+    end
+
+    trace.enable do
+      assert_equal(['  ', 'def', ' ', 'foo', "\n", '    ', '42', "\n", '  ', 'end'], foo.token_stream.map(&:text))
+      assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
+    end
+    assert_equal 1, parse_lex_calls
   end
 
   def test_markup_first_comment

From 99bf9a23334947e219196da7a25f49e9ccd0721d Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 14:25:12 +0200
Subject: [PATCH 07/17] frozen first reads work

---
 lib/rdoc/parser/ruby.rb                 | 32 ++++++++++------------
 lib/rdoc/parser/ruby_colorizer.rb       | 35 +++++--------------------
 lib/rdoc/token_stream.rb                |  7 ++---
 test/rdoc/parser/ruby_colorizer_test.rb | 10 +++----
 test/rdoc/rdoc_token_stream_test.rb     | 18 ++++++++-----
 5 files changed, 41 insertions(+), 61 deletions(-)

diff --git a/lib/rdoc/parser/ruby.rb b/lib/rdoc/parser/ruby.rb
index 259adc6450..99fcb28a72 100644
--- a/lib/rdoc/parser/ruby.rb
+++ b/lib/rdoc/parser/ruby.rb
@@ -366,8 +366,8 @@ def parse_comment_tomdoc(container, comment, line_no, start_line)
     return unless meth.name
 
     node = @line_nodes[line_no]
-    tokens = node ? syntax_highlighted_tokens(node) : []
-    meth.start_collecting_tokens(:ruby, tokens)
+    token_stream_loader = syntax_highlighted_token_loader(node) if node
+    meth.start_collecting_tokens(:ruby, loader: token_stream_loader)
 
     container.add_method meth
     meth.comment = comment
@@ -437,12 +437,8 @@ def handle_meta_method_comment(comment, directives, node)
       end
     elsif line_no || node
       method_name ||= call_node_name_arguments(node).first if is_call_node
-      if node
-        tokens = syntax_highlighted_tokens(node)
-        line_no = node.location.start_line
-      else
-        tokens = []
-      end
+      token_stream_loader = syntax_highlighted_token_loader(node) if node
+      line_no = node.location.start_line if node
       internal_add_method(
         method_name,
         @container,
@@ -455,7 +451,7 @@ def handle_meta_method_comment(comment, directives, node)
         params: nil,
         calls_super: false,
         block_params: nil,
-        tokens: tokens,
+        token_stream_loader: token_stream_loader,
       )
     end
   end
@@ -550,10 +546,10 @@ def extract_section_comment(comment_text, prefix_line_count) # :nodoc:
     comment_text
   end
 
-  # Returns syntax highlighted tokens of the given node
+  # Returns a loader for syntax highlighted tokens of the given node
 
-  def syntax_highlighted_tokens(node)
-    @colorizer_context.deferred_token_stream(node)
+  def syntax_highlighted_token_loader(node)
+    @colorizer_context.token_stream_loader(node)
   end
 
   # Handles `public :foo, :bar` `private :foo, :bar` and `protected :foo, :bar`
@@ -696,7 +692,7 @@ def add_extends(names, line_no) # :nodoc:
 
   # Adds a method defined by `def` syntax
 
-  def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:, singleton:, params:, calls_super:, block_params:, tokens:, start_line:, args_end_line:, end_line:)
+  def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:, singleton:, params:, calls_super:, block_params:, token_stream_loader:, start_line:, args_end_line:, end_line:)
     comment, directives, type_signature_lines = consecutive_comment(start_line)
     apply_document_control_directive(directives) if directives
     handle_code_object_directives(@container, directives) if directives
@@ -716,12 +712,12 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
       params: params,
       calls_super: calls_super,
       block_params: block_params,
-      tokens: tokens,
+      token_stream_loader: token_stream_loader,
       type_signature_lines: type_signature_lines
     )
   end
 
-  private def internal_add_method(method_name, container, comment:, dont_rename_initialize: false, directives:, modifier_comment_lines: nil, line_no:, visibility:, singleton:, params:, calls_super:, block_params:, tokens:, type_signature_lines: nil) # :nodoc:
+  private def internal_add_method(method_name, container, comment:, dont_rename_initialize: false, directives:, modifier_comment_lines: nil, line_no:, visibility:, singleton:, params:, calls_super:, block_params:, token_stream_loader:, type_signature_lines: nil) # :nodoc:
     meth = RDoc::AnyMethod.new(method_name, singleton: singleton)
     meth.comment = comment
     handle_code_object_directives(meth, directives) if directives
@@ -747,7 +743,7 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
     meth.block_params ||= block_params if block_params
     meth.type_signature_lines = type_signature_lines
     record_location(meth)
-    meth.start_collecting_tokens(:ruby, tokens)
+    meth.start_collecting_tokens(:ruby, loader: token_stream_loader)
 
     # Rename after add_method to register duplicated 'new' and 'initialize'
     # defined in c and ruby.
@@ -1169,7 +1165,7 @@ def visit_def_node(node)
       end
       name = node.name.to_s
       params, block_params, calls_super = MethodSignatureVisitor.scan_signature(node)
-      tokens = @scanner.syntax_highlighted_tokens(node)
+      token_stream_loader = @scanner.syntax_highlighted_token_loader(node)
 
       @scanner.add_method(
         name,
@@ -1180,7 +1176,7 @@ def visit_def_node(node)
         params: params,
         block_params: block_params,
         calls_super: calls_super,
-        tokens: tokens,
+        token_stream_loader: token_stream_loader,
         start_line: start_line,
         args_end_line: args_end_line,
         end_line: end_line
diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index 1a773d546f..aaf7d3a7fb 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -10,29 +10,6 @@ module RDoc::Parser::RubyColorizer
 
   ColoredToken = Struct.new(:kind, :text)
 
-  # A token stream populated by its file's DeferredContext.
-  class DeferredTokenStream
-    #: (DeferredContext) -> void
-    def initialize(context)
-      @context = context
-      @tokens = nil
-    end
-
-    #: () -> Array[ColoredToken]
-    def materialize
-      return @tokens if @tokens
-
-      @context&.materialize
-      @tokens
-    end
-
-    #: (Array[ColoredToken]) -> void
-    def resolve(tokens)
-      @tokens = tokens
-      @context = nil
-    end
-  end
-
   # Defers colorization for all nodes in one source file until first access.
   class DeferredContext
     #: (String) -> void
@@ -42,11 +19,11 @@ def initialize(source)
       @mutex = Mutex.new
     end
 
-    #: (Prism::Node) -> DeferredTokenStream
-    def deferred_token_stream(node)
-      stream = DeferredTokenStream.new(self)
-      (@streams[node.node_id] ||= []) << stream
-      stream
+    #: (Prism::Node) -> ^() -> Array[ColoredToken]
+    def token_stream_loader(node)
+      tokens = nil
+      (@streams[node.node_id] ||= []) << ->(resolved_tokens) { tokens = resolved_tokens }
+      -> { tokens || (materialize; tokens) }
     end
 
     #: () -> void
@@ -62,7 +39,7 @@ def materialize
           node = nodes.pop
           if (streams = @streams.delete(node.node_id))
             tokens = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
-            streams.each { |stream| stream.resolve(tokens) }
+            streams.each { |resolve| resolve.call(tokens) }
           end
           nodes.concat(node.compact_child_nodes)
         end
diff --git a/lib/rdoc/token_stream.rb b/lib/rdoc/token_stream.rb
index bd0fabf172..d28893db1d 100644
--- a/lib/rdoc/token_stream.rb
+++ b/lib/rdoc/token_stream.rb
@@ -64,8 +64,10 @@ def add_token(token)
   # Starts collecting tokens
   #
 
-  def collect_tokens(language, initial_stream = [])
+  def collect_tokens(language, initial_stream = [], loader: nil)
     @token_stream = initial_stream
+    loaded_tokens = nil
+    @token_stream_loader = loader && -> { loaded_tokens ||= loader.call }
     @token_stream_language = language
   end
 
@@ -82,8 +84,7 @@ def pop_token
   # Current token stream
 
   def token_stream
-    @token_stream = @token_stream.materialize if @token_stream.respond_to?(:materialize)
-    @token_stream
+    @token_stream_loader ? @token_stream_loader.call : @token_stream
   end
 
   ##
diff --git a/test/rdoc/parser/ruby_colorizer_test.rb b/test/rdoc/parser/ruby_colorizer_test.rb
index 942401f3aa..f512a5e2ee 100644
--- a/test/rdoc/parser/ruby_colorizer_test.rb
+++ b/test/rdoc/parser/ruby_colorizer_test.rb
@@ -20,10 +20,10 @@ def test_deferred_token_stream
     prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
     expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
     context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
-    deferred = context.deferred_token_stream(node)
+    loader = context.token_stream_loader(node)
 
-    assert_equal expected, deferred.materialize
-    assert_same deferred.materialize, deferred.materialize
+    assert_equal expected, loader.call
+    assert_same loader.call, loader.call
   end
 
   def test_deferred_token_stream_preserves_lexical_scope
@@ -34,9 +34,9 @@ def test_deferred_token_stream_preserves_lexical_scope
     expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
 
     context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
-    deferred = context.deferred_token_stream(node)
+    loader = context.token_stream_loader(node)
 
-    assert_equal expected, deferred.materialize
+    assert_equal expected, loader.call
   end
 
   def test_partial_colorize
diff --git a/test/rdoc/rdoc_token_stream_test.rb b/test/rdoc/rdoc_token_stream_test.rb
index 70be2b0abd..da8a0bb2e9 100644
--- a/test/rdoc/rdoc_token_stream_test.rb
+++ b/test/rdoc/rdoc_token_stream_test.rb
@@ -83,11 +83,19 @@ def test_collect_tokens
     end.new
     foo.collect_tokens(:ruby)
     assert_equal [], foo.token_stream
+  end
+
+  def test_collect_tokens_with_loader
+    foo = Class.new do
+      include RDoc::TokenStream
+    end.new
+    loads = 0
+    foo.collect_tokens(:ruby, loader: -> { loads += 1; [:token] })
+    foo.freeze
 
-    deferred = Object.new
-    def deferred.materialize = [:token]
-    foo.collect_tokens(:ruby, deferred)
     assert_equal [:token], foo.token_stream
+    assert_same foo.token_stream, foo.token_stream
+    assert_equal 1, loads
   end
 
   def test_pop_token
@@ -104,9 +112,7 @@ def test_mutating_deferred_tokens
     foo = Class.new do
       include RDoc::TokenStream
     end.new
-    deferred = Object.new
-    def deferred.materialize = [:first]
-    foo.collect_tokens(:ruby, deferred)
+    foo.collect_tokens(:ruby, loader: -> { [:first] })
 
     foo.add_token(:second)
     foo.add_tokens([:third])

From 065d79b7d847ce4319c5e179dd7688d34157f3cb Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 14:46:17 +0200
Subject: [PATCH 08/17] Don't retain every Prism node

---
 lib/rdoc/parser/ruby_colorizer.rb | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index aaf7d3a7fb..d3fe53789e 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -21,8 +21,13 @@ def initialize(source)
 
     #: (Prism::Node) -> ^() -> Array[ColoredToken]
     def token_stream_loader(node)
+      token_stream_loader_for(node.node_id)
+    end
+
+    #: (Integer) -> ^() -> Array[ColoredToken]
+    private def token_stream_loader_for(node_id)
       tokens = nil
-      (@streams[node.node_id] ||= []) << ->(resolved_tokens) { tokens = resolved_tokens }
+      (@streams[node_id] ||= []) << ->(resolved_tokens) { tokens = resolved_tokens }
       -> { tokens || (materialize; tokens) }
     end
 

From 2a8ef559a868a374c304d219984490b54bc1b07c Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 15:01:34 +0200
Subject: [PATCH 09/17] Loader registration now happens only after method
 acceptance

---
 lib/rdoc/parser/ruby.rb                 | 22 ++++++++++------------
 lib/rdoc/parser/ruby_colorizer.rb       |  7 +------
 test/rdoc/parser/ruby_colorizer_test.rb |  4 ++--
 test/rdoc/parser/ruby_test.rb           | 20 ++++++++++++++++++++
 4 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/lib/rdoc/parser/ruby.rb b/lib/rdoc/parser/ruby.rb
index 99fcb28a72..2db8f10e5d 100644
--- a/lib/rdoc/parser/ruby.rb
+++ b/lib/rdoc/parser/ruby.rb
@@ -366,7 +366,7 @@ def parse_comment_tomdoc(container, comment, line_no, start_line)
     return unless meth.name
 
     node = @line_nodes[line_no]
-    token_stream_loader = syntax_highlighted_token_loader(node) if node
+    token_stream_loader = syntax_highlighted_token_loader(node.node_id) if node
     meth.start_collecting_tokens(:ruby, loader: token_stream_loader)
 
     container.add_method meth
@@ -437,7 +437,6 @@ def handle_meta_method_comment(comment, directives, node)
       end
     elsif line_no || node
       method_name ||= call_node_name_arguments(node).first if is_call_node
-      token_stream_loader = syntax_highlighted_token_loader(node) if node
       line_no = node.location.start_line if node
       internal_add_method(
         method_name,
@@ -451,7 +450,7 @@ def handle_meta_method_comment(comment, directives, node)
         params: nil,
         calls_super: false,
         block_params: nil,
-        token_stream_loader: token_stream_loader,
+        node_id: node&.node_id,
       )
     end
   end
@@ -546,10 +545,10 @@ def extract_section_comment(comment_text, prefix_line_count) # :nodoc:
     comment_text
   end
 
-  # Returns a loader for syntax highlighted tokens of the given node
+  # Returns a loader for syntax highlighted tokens of the given node ID
 
-  def syntax_highlighted_token_loader(node)
-    @colorizer_context.token_stream_loader(node)
+  def syntax_highlighted_token_loader(node_id)
+    @colorizer_context.token_stream_loader(node_id)
   end
 
   # Handles `public :foo, :bar` `private :foo, :bar` and `protected :foo, :bar`
@@ -692,7 +691,7 @@ def add_extends(names, line_no) # :nodoc:
 
   # Adds a method defined by `def` syntax
 
-  def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:, singleton:, params:, calls_super:, block_params:, token_stream_loader:, start_line:, args_end_line:, end_line:)
+  def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:, singleton:, params:, calls_super:, block_params:, node_id:, start_line:, args_end_line:, end_line:)
     comment, directives, type_signature_lines = consecutive_comment(start_line)
     apply_document_control_directive(directives) if directives
     handle_code_object_directives(@container, directives) if directives
@@ -712,12 +711,12 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
       params: params,
       calls_super: calls_super,
       block_params: block_params,
-      token_stream_loader: token_stream_loader,
+      node_id: node_id,
       type_signature_lines: type_signature_lines
     )
   end
 
-  private def internal_add_method(method_name, container, comment:, dont_rename_initialize: false, directives:, modifier_comment_lines: nil, line_no:, visibility:, singleton:, params:, calls_super:, block_params:, token_stream_loader:, type_signature_lines: nil) # :nodoc:
+  private def internal_add_method(method_name, container, comment:, dont_rename_initialize: false, directives:, modifier_comment_lines: nil, line_no:, visibility:, singleton:, params:, calls_super:, block_params:, node_id:, type_signature_lines: nil) # :nodoc:
     meth = RDoc::AnyMethod.new(method_name, singleton: singleton)
     meth.comment = comment
     handle_code_object_directives(meth, directives) if directives
@@ -743,6 +742,7 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
     meth.block_params ||= block_params if block_params
     meth.type_signature_lines = type_signature_lines
     record_location(meth)
+    token_stream_loader = syntax_highlighted_token_loader(node_id) if node_id
     meth.start_collecting_tokens(:ruby, loader: token_stream_loader)
 
     # Rename after add_method to register duplicated 'new' and 'initialize'
@@ -1165,8 +1165,6 @@ def visit_def_node(node)
       end
       name = node.name.to_s
       params, block_params, calls_super = MethodSignatureVisitor.scan_signature(node)
-      token_stream_loader = @scanner.syntax_highlighted_token_loader(node)
-
       @scanner.add_method(
         name,
         receiver_name: receiver_name,
@@ -1176,7 +1174,7 @@ def visit_def_node(node)
         params: params,
         block_params: block_params,
         calls_super: calls_super,
-        token_stream_loader: token_stream_loader,
+        node_id: node.node_id,
         start_line: start_line,
         args_end_line: args_end_line,
         end_line: end_line
diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index d3fe53789e..34ee9c398f 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -19,13 +19,8 @@ def initialize(source)
       @mutex = Mutex.new
     end
 
-    #: (Prism::Node) -> ^() -> Array[ColoredToken]
-    def token_stream_loader(node)
-      token_stream_loader_for(node.node_id)
-    end
-
     #: (Integer) -> ^() -> Array[ColoredToken]
-    private def token_stream_loader_for(node_id)
+    def token_stream_loader(node_id)
       tokens = nil
       (@streams[node_id] ||= []) << ->(resolved_tokens) { tokens = resolved_tokens }
       -> { tokens || (materialize; tokens) }
diff --git a/test/rdoc/parser/ruby_colorizer_test.rb b/test/rdoc/parser/ruby_colorizer_test.rb
index f512a5e2ee..21eed21e38 100644
--- a/test/rdoc/parser/ruby_colorizer_test.rb
+++ b/test/rdoc/parser/ruby_colorizer_test.rb
@@ -20,7 +20,7 @@ def test_deferred_token_stream
     prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
     expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
     context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
-    loader = context.token_stream_loader(node)
+    loader = context.token_stream_loader(node.node_id)
 
     assert_equal expected, loader.call
     assert_same loader.call, loader.call
@@ -34,7 +34,7 @@ def test_deferred_token_stream_preserves_lexical_scope
     expected = RDoc::Parser::RubyColorizer.partial_colorize(code, node, prism_tokens)
 
     context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
-    loader = context.token_stream_loader(node)
+    loader = context.token_stream_loader(node.node_id)
 
     assert_equal expected, loader.call
   end
diff --git a/test/rdoc/parser/ruby_test.rb b/test/rdoc/parser/ruby_test.rb
index 9d4bd51743..fb27e5ed2a 100644
--- a/test/rdoc/parser/ruby_test.rb
+++ b/test/rdoc/parser/ruby_test.rb
@@ -2723,6 +2723,26 @@ def foo
     assert_equal 1, parse_lex_calls
   end
 
+  def test_rejected_method_token_stream_is_not_colorized
+    util_parser <<~RUBY
+      class Foo
+        def hidden; end # :nodoc:
+        def visible; end
+      end
+    RUBY
+
+    visible = @top_level.classes.first.method_list.first
+    partial_colorize_calls = 0
+    trace = TracePoint.new(:call) do |event|
+      if event.self.equal?(RDoc::Parser::RubyColorizer) && event.method_id == :partial_colorize
+        partial_colorize_calls += 1
+      end
+    end
+
+    trace.enable { visible.token_stream }
+    assert_equal 1, partial_colorize_calls
+  end
+
   def test_markup_first_comment
     util_parser <<~RUBY
       # :markup: rd

From 06fc9058a5c6d22a404b0fad67ca7a027cdbaee9 Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 15:09:33 +0200
Subject: [PATCH 10/17] TokenStream now uses one @token_stream value: eager
 array or memoizing loader

---
 lib/rdoc/token_stream.rb            | 7 +++----
 test/rdoc/rdoc_token_stream_test.rb | 8 ++++++++
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/lib/rdoc/token_stream.rb b/lib/rdoc/token_stream.rb
index d28893db1d..ee6b8af850 100644
--- a/lib/rdoc/token_stream.rb
+++ b/lib/rdoc/token_stream.rb
@@ -64,10 +64,9 @@ def add_token(token)
   # Starts collecting tokens
   #
 
-  def collect_tokens(language, initial_stream = [], loader: nil)
-    @token_stream = initial_stream
+  def collect_tokens(language, loader: nil)
     loaded_tokens = nil
-    @token_stream_loader = loader && -> { loaded_tokens ||= loader.call }
+    @token_stream = loader ? -> { loaded_tokens ||= loader.call } : []
     @token_stream_language = language
   end
 
@@ -84,7 +83,7 @@ def pop_token
   # Current token stream
 
   def token_stream
-    @token_stream_loader ? @token_stream_loader.call : @token_stream
+    @token_stream.is_a?(Proc) ? @token_stream.call : @token_stream
   end
 
   ##
diff --git a/test/rdoc/rdoc_token_stream_test.rb b/test/rdoc/rdoc_token_stream_test.rb
index da8a0bb2e9..d6ead27af9 100644
--- a/test/rdoc/rdoc_token_stream_test.rb
+++ b/test/rdoc/rdoc_token_stream_test.rb
@@ -85,6 +85,14 @@ def test_collect_tokens
     assert_equal [], foo.token_stream
   end
 
+  def test_collect_tokens_does_not_accept_initial_stream
+    foo = Class.new do
+      include RDoc::TokenStream
+    end.new
+
+    assert_raise(ArgumentError) { foo.collect_tokens(:ruby, [:token]) }
+  end
+
   def test_collect_tokens_with_loader
     foo = Class.new do
       include RDoc::TokenStream

From e452cf92f6591810b55abe4e3cb018fe1b89825c Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 15:15:20 +0200
Subject: [PATCH 11/17] Failures retain source and register atomic retry

---
 lib/rdoc/parser/ruby_colorizer.rb       | 20 +++++++++-----
 test/rdoc/parser/ruby_colorizer_test.rb | 35 +++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index 34ee9c398f..1d99e6c67f 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -23,7 +23,7 @@ def initialize(source)
     def token_stream_loader(node_id)
       tokens = nil
       (@streams[node_id] ||= []) << ->(resolved_tokens) { tokens = resolved_tokens }
-      -> { tokens || (materialize; tokens) }
+      -> { @source ? (materialize; tokens) : tokens }
     end
 
     #: () -> void
@@ -34,18 +34,24 @@ def materialize
         source = @source
         program_node, unordered_tokens = Prism.parse_lex(source).value
         prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
+        staged_tokens = {}
         nodes = [program_node]
-        until nodes.empty? || @streams.empty?
+        until nodes.empty? || staged_tokens.size == @streams.size
           node = nodes.pop
-          if (streams = @streams.delete(node.node_id))
-            tokens = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
-            streams.each { |resolve| resolve.call(tokens) }
+          if @streams.key?(node.node_id)
+            staged_tokens[node.node_id] = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
           end
           nodes.concat(node.compact_child_nodes)
         end
-      ensure
-        @source = nil
+
+        missing_node_ids = @streams.keys - staged_tokens.keys
+        raise KeyError, "Prism nodes not found: #{missing_node_ids.join(', ')}" unless missing_node_ids.empty?
+
+        @streams.each do |node_id, streams|
+          streams.each { |resolve| resolve.call(staged_tokens.fetch(node_id)) }
+        end
         @streams = nil
+        @source = nil
       end
     end
   end
diff --git a/test/rdoc/parser/ruby_colorizer_test.rb b/test/rdoc/parser/ruby_colorizer_test.rb
index 21eed21e38..0cadd99ad3 100644
--- a/test/rdoc/parser/ruby_colorizer_test.rb
+++ b/test/rdoc/parser/ruby_colorizer_test.rb
@@ -39,6 +39,41 @@ def test_deferred_token_stream_preserves_lexical_scope
     assert_equal expected, loader.call
   end
 
+  def test_deferred_token_stream_retries_atomically
+    code = "first\nsecond\n"
+    nodes = Prism.parse(code).value.statements.body
+    context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
+    loaders = nodes.map { |node| context.token_stream_loader(node.node_id) }
+    colorizer = RDoc::Parser::RubyColorizer
+    partial_colorize = colorizer.method(:partial_colorize)
+    calls = 0
+
+    colorizer.define_singleton_method(:partial_colorize) do |*arguments|
+      calls += 1
+      raise 'colorization failed' if calls == 2
+
+      partial_colorize.call(*arguments)
+    end
+    begin
+      assert_raise(RuntimeError) { loaders.first.call }
+    ensure
+      colorizer.define_singleton_method(:partial_colorize, partial_colorize)
+    end
+
+    assert_equal %w[first second], loaders.map { |loader| loader.call.map(&:text).join }
+  end
+
+  def test_deferred_token_stream_rejects_unmatched_node_id
+    code = "first\n"
+    node = Prism.parse(code).value.statements.body.first
+    context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
+    loader = context.token_stream_loader(node.node_id)
+    missing_loader = context.token_stream_loader(node.node_id + 1_000_000)
+
+    2.times { assert_raise(KeyError) { loader.call } }
+    assert_raise(KeyError) { missing_loader.call }
+  end
+
   def test_partial_colorize
     code = <<~RUBY
       class A

From 0ead803da4dc876ea81ef1872e52c84c4fac298c Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 15:42:59 +0200
Subject: [PATCH 12/17] simplifications

---
 lib/rdoc/parser/ruby.rb             | 10 ++-----
 lib/rdoc/parser/ruby_colorizer.rb   | 41 +++++++++++++----------------
 lib/rdoc/token_stream.rb            |  3 +--
 test/rdoc/parser/ruby_test.rb       | 26 ++++--------------
 test/rdoc/rdoc_token_stream_test.rb | 16 +++--------
 5 files changed, 30 insertions(+), 66 deletions(-)

diff --git a/lib/rdoc/parser/ruby.rb b/lib/rdoc/parser/ruby.rb
index 2db8f10e5d..396c5ed00b 100644
--- a/lib/rdoc/parser/ruby.rb
+++ b/lib/rdoc/parser/ruby.rb
@@ -366,7 +366,7 @@ def parse_comment_tomdoc(container, comment, line_no, start_line)
     return unless meth.name
 
     node = @line_nodes[line_no]
-    token_stream_loader = syntax_highlighted_token_loader(node.node_id) if node
+    token_stream_loader = @colorizer_context.token_stream_loader(node.node_id) if node
     meth.start_collecting_tokens(:ruby, loader: token_stream_loader)
 
     container.add_method meth
@@ -545,12 +545,6 @@ def extract_section_comment(comment_text, prefix_line_count) # :nodoc:
     comment_text
   end
 
-  # Returns a loader for syntax highlighted tokens of the given node ID
-
-  def syntax_highlighted_token_loader(node_id)
-    @colorizer_context.token_stream_loader(node_id)
-  end
-
   # Handles `public :foo, :bar` `private :foo, :bar` and `protected :foo, :bar`
 
   def change_method_visibility(names, visibility, singleton: @singleton)
@@ -742,7 +736,7 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
     meth.block_params ||= block_params if block_params
     meth.type_signature_lines = type_signature_lines
     record_location(meth)
-    token_stream_loader = syntax_highlighted_token_loader(node_id) if node_id
+    token_stream_loader = @colorizer_context.token_stream_loader(node_id) if node_id
     meth.start_collecting_tokens(:ruby, loader: token_stream_loader)
 
     # Rename after add_method to register duplicated 'new' and 'initialize'
diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index 1d99e6c67f..db28d81f6b 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -16,7 +16,6 @@ class DeferredContext
     def initialize(source)
       @source = source
       @streams = {}
-      @mutex = Mutex.new
     end
 
     #: (Integer) -> ^() -> Array[ColoredToken]
@@ -28,31 +27,29 @@ def token_stream_loader(node_id)
 
     #: () -> void
     def materialize
-      @mutex.synchronize do
-        return unless @source
-
-        source = @source
-        program_node, unordered_tokens = Prism.parse_lex(source).value
-        prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
-        staged_tokens = {}
-        nodes = [program_node]
-        until nodes.empty? || staged_tokens.size == @streams.size
-          node = nodes.pop
-          if @streams.key?(node.node_id)
-            staged_tokens[node.node_id] = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
-          end
-          nodes.concat(node.compact_child_nodes)
+      return unless @source
+
+      source = @source
+      program_node, unordered_tokens = Prism.parse_lex(source).value
+      prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
+      staged_tokens = {}
+      nodes = [program_node]
+      until nodes.empty? || staged_tokens.size == @streams.size
+        node = nodes.pop
+        if @streams.key?(node.node_id)
+          staged_tokens[node.node_id] = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
         end
+        nodes.concat(node.compact_child_nodes)
+      end
 
-        missing_node_ids = @streams.keys - staged_tokens.keys
-        raise KeyError, "Prism nodes not found: #{missing_node_ids.join(', ')}" unless missing_node_ids.empty?
+      missing_node_ids = @streams.keys - staged_tokens.keys
+      raise KeyError, "Prism nodes not found: #{missing_node_ids.join(', ')}" unless missing_node_ids.empty?
 
-        @streams.each do |node_id, streams|
-          streams.each { |resolve| resolve.call(staged_tokens.fetch(node_id)) }
-        end
-        @streams = nil
-        @source = nil
+      @streams.each do |node_id, streams|
+        streams.each { |resolve| resolve.call(staged_tokens.fetch(node_id)) }
       end
+      @streams = nil
+      @source = nil
     end
   end
 
diff --git a/lib/rdoc/token_stream.rb b/lib/rdoc/token_stream.rb
index ee6b8af850..6c864c9f5b 100644
--- a/lib/rdoc/token_stream.rb
+++ b/lib/rdoc/token_stream.rb
@@ -65,8 +65,7 @@ def add_token(token)
   #
 
   def collect_tokens(language, loader: nil)
-    loaded_tokens = nil
-    @token_stream = loader ? -> { loaded_tokens ||= loader.call } : []
+    @token_stream = loader || []
     @token_stream_language = language
   end
 
diff --git a/test/rdoc/parser/ruby_test.rb b/test/rdoc/parser/ruby_test.rb
index fb27e5ed2a..4960a538fd 100644
--- a/test/rdoc/parser/ruby_test.rb
+++ b/test/rdoc/parser/ruby_test.rb
@@ -2700,6 +2700,8 @@ def test_read_directive_linear_performance
   def test_code_object_token_stream
     util_parser <<~RUBY
       class Foo
+        def hidden; end # :nodoc:
+
         def foo
           42
         end
@@ -2711,9 +2713,10 @@ def foo
     RUBY
 
     foo, bar = @top_level.classes.first.method_list
-    parse_lex_calls = 0
+    parse_lex_calls = partial_colorize_calls = 0
     trace = TracePoint.new(:call, :c_call) do |event|
       parse_lex_calls += 1 if event.self.equal?(Prism) && event.method_id == :parse_lex
+      partial_colorize_calls += 1 if event.self.equal?(RDoc::Parser::RubyColorizer) && event.method_id == :partial_colorize
     end
 
     trace.enable do
@@ -2721,26 +2724,7 @@ def foo
       assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
     end
     assert_equal 1, parse_lex_calls
-  end
-
-  def test_rejected_method_token_stream_is_not_colorized
-    util_parser <<~RUBY
-      class Foo
-        def hidden; end # :nodoc:
-        def visible; end
-      end
-    RUBY
-
-    visible = @top_level.classes.first.method_list.first
-    partial_colorize_calls = 0
-    trace = TracePoint.new(:call) do |event|
-      if event.self.equal?(RDoc::Parser::RubyColorizer) && event.method_id == :partial_colorize
-        partial_colorize_calls += 1
-      end
-    end
-
-    trace.enable { visible.token_stream }
-    assert_equal 1, partial_colorize_calls
+    assert_equal 2, partial_colorize_calls
   end
 
   def test_markup_first_comment
diff --git a/test/rdoc/rdoc_token_stream_test.rb b/test/rdoc/rdoc_token_stream_test.rb
index d6ead27af9..dbedd24b4d 100644
--- a/test/rdoc/rdoc_token_stream_test.rb
+++ b/test/rdoc/rdoc_token_stream_test.rb
@@ -85,25 +85,14 @@ def test_collect_tokens
     assert_equal [], foo.token_stream
   end
 
-  def test_collect_tokens_does_not_accept_initial_stream
-    foo = Class.new do
-      include RDoc::TokenStream
-    end.new
-
-    assert_raise(ArgumentError) { foo.collect_tokens(:ruby, [:token]) }
-  end
-
   def test_collect_tokens_with_loader
     foo = Class.new do
       include RDoc::TokenStream
     end.new
-    loads = 0
-    foo.collect_tokens(:ruby, loader: -> { loads += 1; [:token] })
+    foo.collect_tokens(:ruby, loader: -> { [:token] })
     foo.freeze
 
     assert_equal [:token], foo.token_stream
-    assert_same foo.token_stream, foo.token_stream
-    assert_equal 1, loads
   end
 
   def test_pop_token
@@ -120,7 +109,8 @@ def test_mutating_deferred_tokens
     foo = Class.new do
       include RDoc::TokenStream
     end.new
-    foo.collect_tokens(:ruby, loader: -> { [:first] })
+    tokens = [:first]
+    foo.collect_tokens(:ruby, loader: -> { tokens })
 
     foo.add_token(:second)
     foo.add_tokens([:third])

From 2b9fd2e1ec4842a60585c0fc5dab20b406e27d4b Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 15:52:07 +0200
Subject: [PATCH 13/17] return per-context Mutex to guard against races

---
 lib/rdoc/parser/ruby_colorizer.rb | 41 +++++++++++++++++--------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index db28d81f6b..1d99e6c67f 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -16,6 +16,7 @@ class DeferredContext
     def initialize(source)
       @source = source
       @streams = {}
+      @mutex = Mutex.new
     end
 
     #: (Integer) -> ^() -> Array[ColoredToken]
@@ -27,29 +28,31 @@ def token_stream_loader(node_id)
 
     #: () -> void
     def materialize
-      return unless @source
-
-      source = @source
-      program_node, unordered_tokens = Prism.parse_lex(source).value
-      prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
-      staged_tokens = {}
-      nodes = [program_node]
-      until nodes.empty? || staged_tokens.size == @streams.size
-        node = nodes.pop
-        if @streams.key?(node.node_id)
-          staged_tokens[node.node_id] = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
+      @mutex.synchronize do
+        return unless @source
+
+        source = @source
+        program_node, unordered_tokens = Prism.parse_lex(source).value
+        prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
+        staged_tokens = {}
+        nodes = [program_node]
+        until nodes.empty? || staged_tokens.size == @streams.size
+          node = nodes.pop
+          if @streams.key?(node.node_id)
+            staged_tokens[node.node_id] = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
+          end
+          nodes.concat(node.compact_child_nodes)
         end
-        nodes.concat(node.compact_child_nodes)
-      end
 
-      missing_node_ids = @streams.keys - staged_tokens.keys
-      raise KeyError, "Prism nodes not found: #{missing_node_ids.join(', ')}" unless missing_node_ids.empty?
+        missing_node_ids = @streams.keys - staged_tokens.keys
+        raise KeyError, "Prism nodes not found: #{missing_node_ids.join(', ')}" unless missing_node_ids.empty?
 
-      @streams.each do |node_id, streams|
-        streams.each { |resolve| resolve.call(staged_tokens.fetch(node_id)) }
+        @streams.each do |node_id, streams|
+          streams.each { |resolve| resolve.call(staged_tokens.fetch(node_id)) }
+        end
+        @streams = nil
+        @source = nil
       end
-      @streams = nil
-      @source = nil
     end
   end
 

From de2840f1136c9cc9c9f1f9648f692118f28fbe85 Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 15:57:37 +0200
Subject: [PATCH 14/17] Memoized deferred loader results inside collect_tokens

---
 lib/rdoc/token_stream.rb            | 4 +++-
 test/rdoc/rdoc_token_stream_test.rb | 8 ++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/lib/rdoc/token_stream.rb b/lib/rdoc/token_stream.rb
index 6c864c9f5b..3d8cfb8cd1 100644
--- a/lib/rdoc/token_stream.rb
+++ b/lib/rdoc/token_stream.rb
@@ -63,9 +63,11 @@ def add_token(token)
   ##
   # Starts collecting tokens
   #
+  # The optional +loader+ is called once on first access and its result is reused.
 
   def collect_tokens(language, loader: nil)
-    @token_stream = loader || []
+    loaded_tokens = nil
+    @token_stream = loader ? -> { loaded_tokens ||= loader.call } : []
     @token_stream_language = language
   end
 
diff --git a/test/rdoc/rdoc_token_stream_test.rb b/test/rdoc/rdoc_token_stream_test.rb
index dbedd24b4d..025c152dd9 100644
--- a/test/rdoc/rdoc_token_stream_test.rb
+++ b/test/rdoc/rdoc_token_stream_test.rb
@@ -89,10 +89,14 @@ def test_collect_tokens_with_loader
     foo = Class.new do
       include RDoc::TokenStream
     end.new
-    foo.collect_tokens(:ruby, loader: -> { [:token] })
+    loads = 0
+    foo.collect_tokens(:ruby, loader: -> { loads += 1; [:token] })
     foo.freeze
 
-    assert_equal [:token], foo.token_stream
+    tokens = foo.token_stream
+    assert_equal [:token], tokens
+    assert_same tokens, foo.token_stream
+    assert_equal 1, loads
   end
 
   def test_pop_token

From c1dc957720f52bef1bacf93ebc0d176fb5766f1c Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 16:27:33 +0200
Subject: [PATCH 15/17] Workaround for TruffleRuby tests

---
 test/rdoc/parser/ruby_test.rb | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/test/rdoc/parser/ruby_test.rb b/test/rdoc/parser/ruby_test.rb
index 4960a538fd..bd233a685d 100644
--- a/test/rdoc/parser/ruby_test.rb
+++ b/test/rdoc/parser/ruby_test.rb
@@ -2714,14 +2714,25 @@ def foo
 
     foo, bar = @top_level.classes.first.method_list
     parse_lex_calls = partial_colorize_calls = 0
-    trace = TracePoint.new(:call, :c_call) do |event|
-      parse_lex_calls += 1 if event.self.equal?(Prism) && event.method_id == :parse_lex
-      partial_colorize_calls += 1 if event.self.equal?(RDoc::Parser::RubyColorizer) && event.method_id == :partial_colorize
+    parse_lex = Prism.method(:parse_lex)
+    colorizer = RDoc::Parser::RubyColorizer
+    partial_colorize = colorizer.method(:partial_colorize)
+
+    Prism.define_singleton_method(:parse_lex) do |*arguments|
+      parse_lex_calls += 1
+      parse_lex.call(*arguments)
+    end
+    colorizer.define_singleton_method(:partial_colorize) do |*arguments|
+      partial_colorize_calls += 1
+      partial_colorize.call(*arguments)
     end
 
-    trace.enable do
+    begin
       assert_equal(['  ', 'def', ' ', 'foo', "\n", '    ', '42', "\n", '  ', 'end'], foo.token_stream.map(&:text))
       assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
+    ensure
+      Prism.define_singleton_method(:parse_lex, parse_lex)
+      colorizer.define_singleton_method(:partial_colorize, partial_colorize)
     end
     assert_equal 1, parse_lex_calls
     assert_equal 2, partial_colorize_calls

From e911cf2b7c3affc4a32d0e70f9282b5156a8ff84 Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Fri, 28 Aug 2026 16:41:19 +0200
Subject: [PATCH 16/17] Moved call-count coverage into ruby_colorizer_test.rb

---
 test/rdoc/parser/ruby_colorizer_test.rb | 29 +++++++++++++++++++++++++
 test/rdoc/parser/ruby_test.rb           | 27 ++---------------------
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/test/rdoc/parser/ruby_colorizer_test.rb b/test/rdoc/parser/ruby_colorizer_test.rb
index 0cadd99ad3..45a8329bc6 100644
--- a/test/rdoc/parser/ruby_colorizer_test.rb
+++ b/test/rdoc/parser/ruby_colorizer_test.rb
@@ -26,6 +26,35 @@ def test_deferred_token_stream
     assert_same loader.call, loader.call
   end
 
+  def test_deferred_token_stream_materializes_registered_nodes_once
+    code = "hidden\nfirst\nsecond\n"
+    nodes = Prism.parse(code).value.statements.body
+    context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
+    loaders = nodes.drop(1).map { |node| context.token_stream_loader(node.node_id) }
+    parse_lex_calls = partial_colorize_calls = 0
+    parse_lex = Prism.method(:parse_lex)
+    colorizer = RDoc::Parser::RubyColorizer
+    partial_colorize = colorizer.method(:partial_colorize)
+
+    Prism.define_singleton_method(:parse_lex) do |*arguments|
+      parse_lex_calls += 1
+      parse_lex.call(*arguments)
+    end
+    colorizer.define_singleton_method(:partial_colorize) do |*arguments|
+      partial_colorize_calls += 1
+      partial_colorize.call(*arguments)
+    end
+
+    begin
+      assert_equal %w[first second], loaders.map { |loader| loader.call.map(&:text).join }
+    ensure
+      Prism.define_singleton_method(:parse_lex, parse_lex)
+      colorizer.define_singleton_method(:partial_colorize, partial_colorize)
+    end
+    assert_equal 1, parse_lex_calls
+    assert_equal 2, partial_colorize_calls
+  end
+
   def test_deferred_token_stream_preserves_lexical_scope
     code = "x = 1\ny = 2\ni = 3\nx /y/i\n"
     program_node, unordered_tokens = Prism.parse_lex(code).value
diff --git a/test/rdoc/parser/ruby_test.rb b/test/rdoc/parser/ruby_test.rb
index bd233a685d..3ebfd5891b 100644
--- a/test/rdoc/parser/ruby_test.rb
+++ b/test/rdoc/parser/ruby_test.rb
@@ -2700,8 +2700,6 @@ def test_read_directive_linear_performance
   def test_code_object_token_stream
     util_parser <<~RUBY
       class Foo
-        def hidden; end # :nodoc:
-
         def foo
           42
         end
@@ -2713,29 +2711,8 @@ def foo
     RUBY
 
     foo, bar = @top_level.classes.first.method_list
-    parse_lex_calls = partial_colorize_calls = 0
-    parse_lex = Prism.method(:parse_lex)
-    colorizer = RDoc::Parser::RubyColorizer
-    partial_colorize = colorizer.method(:partial_colorize)
-
-    Prism.define_singleton_method(:parse_lex) do |*arguments|
-      parse_lex_calls += 1
-      parse_lex.call(*arguments)
-    end
-    colorizer.define_singleton_method(:partial_colorize) do |*arguments|
-      partial_colorize_calls += 1
-      partial_colorize.call(*arguments)
-    end
-
-    begin
-      assert_equal(['  ', 'def', ' ', 'foo', "\n", '    ', '42', "\n", '  ', 'end'], foo.token_stream.map(&:text))
-      assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
-    ensure
-      Prism.define_singleton_method(:parse_lex, parse_lex)
-      colorizer.define_singleton_method(:partial_colorize, partial_colorize)
-    end
-    assert_equal 1, parse_lex_calls
-    assert_equal 2, partial_colorize_calls
+    assert_equal(['  ', 'def', ' ', 'foo', "\n", '    ', '42', "\n", '  ', 'end'], foo.token_stream.map(&:text))
+    assert_equal(['          ', 'def', ' ', 'bar', "\n", '    ', 'baz', "\n", '  ', 'end'], bar.token_stream.map(&:text))
   end
 
   def test_markup_first_comment

From cf91fdd6f38136f28ace73f56c9545d671c3fe00 Mon Sep 17 00:00:00 2001
From: "Stanislav (Stas) Katkov" <git@skatkov.com>
Date: Sat, 29 Aug 2026 00:35:30 +0200
Subject: [PATCH 17/17] simplify materialize method

---
 lib/rdoc/parser/ruby_colorizer.rb | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/lib/rdoc/parser/ruby_colorizer.rb b/lib/rdoc/parser/ruby_colorizer.rb
index 1d99e6c67f..c71bd325e2 100644
--- a/lib/rdoc/parser/ruby_colorizer.rb
+++ b/lib/rdoc/parser/ruby_colorizer.rb
@@ -31,22 +31,18 @@ def materialize
       @mutex.synchronize do
         return unless @source
 
-        source = @source
-        program_node, unordered_tokens = Prism.parse_lex(source).value
+        program_node, unordered_tokens = Prism.parse_lex(@source).value
         prism_tokens = unordered_tokens.map(&:first).sort_by! { |token| token.location.start_offset }
         staged_tokens = {}
         nodes = [program_node]
-        until nodes.empty? || staged_tokens.size == @streams.size
+        until nodes.empty?
           node = nodes.pop
           if @streams.key?(node.node_id)
-            staged_tokens[node.node_id] = RDoc::Parser::RubyColorizer.partial_colorize(source, node, prism_tokens)
+            staged_tokens[node.node_id] = RDoc::Parser::RubyColorizer.partial_colorize(@source, node, prism_tokens)
           end
           nodes.concat(node.compact_child_nodes)
         end
 
-        missing_node_ids = @streams.keys - staged_tokens.keys
-        raise KeyError, "Prism nodes not found: #{missing_node_ids.join(', ')}" unless missing_node_ids.empty?
-
         @streams.each do |node_id, streams|
           streams.each { |resolve| resolve.call(staged_tokens.fetch(node_id)) }
         end