ruby
diff --git a/‎lib/rdoc/generator/markup.rb‎
Lines changed: 2 additions & 2 deletions b/‎lib/rdoc/generator/markup.rb‎
Lines changed: 2 additions & 2 deletions
diff --git a/‎lib/rdoc/parser/ruby.rb‎
Lines changed: 13 additions & 32 deletions b/‎lib/rdoc/parser/ruby.rb‎
Lines changed: 13 additions & 32 deletions
diff --git a/‎lib/rdoc/parser/ruby_colorizer.rb‎
Lines changed: 46 additions & 0 deletions b/‎lib/rdoc/parser/ruby_colorizer.rb‎
Lines changed: 46 additions & 0 deletions
diff --git a/‎lib/rdoc/token_stream.rb‎
Lines changed: 10 additions & 4 deletions b/‎lib/rdoc/token_stream.rb‎
Lines changed: 10 additions & 4 deletions
diff --git a/‎test/rdoc/parser/ruby_colorizer_test.rb‎
Lines changed: 96 additions & 0 deletions b/‎test/rdoc/parser/ruby_colorizer_test.rb‎
Lines changed: 96 additions & 0 deletions
diff --git a/‎test/rdoc/rdoc_token_stream_test.rb‎
Lines changed: 27 additions & 0 deletions b/‎test/rdoc/rdoc_token_stream_test.rb‎
Lines changed: 27 additions & 0 deletions
@@ -138,9 +138,9 @@ def add_location_comment(src)
   # Prepends line numbers if +options.line_numbers+ is true.
 
   def markup_code
-    return '' if !@token_stream
+    return '' unless token_stream
 
-    src = RDoc::TokenStream.to_html @token_stream
+    src = RDoc::TokenStream.to_html token_stream
 
     # dedent the source
     common_indent = src.length
 
@@ -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
@@ -251,11 +252,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 +365,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 }
+    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
     meth.comment = comment
@@ -440,12 +437,7 @@ 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
+      line_no = node.location.start_line if node
       internal_add_method(
         method_name,
         @container,
@@ -457,7 +449,7 @@ def handle_meta_method_comment(comment, directives, node)
         params: nil,
         calls_super: false,
         block_params: nil,
-        tokens: tokens,
+        node_id: node&.node_id,
       )
     end
   end
@@ -552,12 +544,6 @@ def extract_section_comment(comment_text, prefix_line_count) # :nodoc:
     comment_text
   end
 
-  # Returns syntax highlighted tokens of the given node
-
-  def syntax_highlighted_tokens(node)
-    RDoc::Parser::RubyColorizer.partial_colorize(@content, node, @prism_tokens)
-  end
-
   # Handles `public :foo, :bar` `private :foo, :bar` and `protected :foo, :bar`
 
   def change_method_visibility(names, visibility, singleton: @singleton)
@@ -696,7 +682,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:, 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
@@ -716,12 +702,12 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
       params: params,
       calls_super: calls_super,
       block_params: block_params,
-      tokens: tokens,
+      node_id: node_id,
       type_signature_lines: type_signature_lines
     )
   end
 
-  private def internal_add_method(method_name, container, comment:, 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:, 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
@@ -745,7 +731,6 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
     meth.calls_super = calls_super
     meth.block_params ||= block_params if block_params
     meth.type_signature_lines = type_signature_lines
-
     # An instance method `initialize` is documented as `::new` unless the
     # :notnew: directive is given
     if method_name == 'initialize' && !singleton
@@ -760,10 +745,8 @@ def add_method(method_name, receiver_name:, receiver_fallback_type:, visibility:
 
     record_location(meth)
     container.add_method(meth)
-    meth.start_collecting_tokens(:ruby)
-    tokens.each do |token|
-      meth.token_stream << token
-    end
+    token_stream_loader = @colorizer_context.token_stream_loader(node_id) if node_id
+    meth.start_collecting_tokens(:ruby, loader: token_stream_loader)
   end
 
   # Find or create module or class from a given module name using Ruby lexical
@@ -1173,8 +1156,6 @@ 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)
-
       @scanner.add_method(
         name,
         receiver_name: receiver_name,
@@ -1184,7 +1165,7 @@ def visit_def_node(node)
         params: params,
         block_params: block_params,
         calls_super: calls_super,
-        tokens: tokens,
+        node_id: node.node_id,
         start_line: start_line,
         args_end_line: args_end_line,
         end_line: end_line
 
@@ -10,6 +10,52 @@ module RDoc::Parser::RubyColorizer
 
   ColoredToken = Struct.new(:kind, :text)
 
+  # Defers colorization for all nodes in one source file until first access.
+  class DeferredContext # :nodoc:
+    #: (String) -> void
+    def initialize(source)
+      @source = source
+      @token_streams = {}
+    end
+
+    #: (Integer) -> ^() -> Array[ColoredToken]
+    def token_stream_loader(node_id)
+      @token_streams[node_id] = nil
+      -> { token_stream_for(node_id) }
+    end
+
+    #: (Integer) -> Array[ColoredToken]
+    def token_stream_for(node_id)
+      materialize unless materialized?
+      @token_streams.fetch(node_id)
+    end
+
+    private
+
+    #: () -> Boolean
+    def materialized?
+      !@source
+    end
+
+    #: () -> void
+    def materialize
+      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?
+        node = nodes.pop
+        if @token_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
+
+      @token_streams = staged_tokens
+      @source = nil
+    end
+  end
+
   # Prism operator token types except assignment '='
   OP_TOKENS = %i[
     AMPERSAND AMPERSAND_AMPERSAND
 
@@ -50,22 +50,24 @@ 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
   #
+  # The optional +loader+ is called once on first access and its result is reused.
 
-  def collect_tokens(language)
+  def collect_tokens(language, loader: nil)
     @token_stream = []
+    @token_stream_loader = loader
     @token_stream_language = language
   end
 
@@ -75,13 +77,17 @@ 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
+    if @token_stream_loader
+      @token_stream = @token_stream_loader.call
+      @token_stream_loader = nil
+    end
     @token_stream
   end
 
 
@@ -7,6 +7,102 @@ def token(kind, text)
     RDoc::Parser::RubyColorizer::ColoredToken.new(kind, text)
   end
 
+  def test_deferred_token_stream
+    code = <<~'RUBY'
+      first(<<~ONE); second(<<~TWO) && sibling
+        one
+      ONE
+        two
+      TWO
+    RUBY
+    program_node, unordered_tokens = Prism.parse_lex(code).value
+    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)
+    context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
+    loader = context.token_stream_loader(node.node_id)
+
+    assert_equal expected, loader.call
+    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
+    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)
+
+    context = RDoc::Parser::RubyColorizer::DeferredContext.new(code)
+    loader = context.token_stream_loader(node.node_id)
+
+    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)
+
+    assert_equal "first", loader.call.map(&:text).join
+    2.times { assert_raise(KeyError) { missing_loader.call } }
+  end
+
   def test_partial_colorize
     code = <<~RUBY
       class A
 
@@ -85,6 +85,19 @@ def test_collect_tokens
     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] })
+
+    tokens = foo.token_stream
+    assert_equal [:token], tokens
+    assert_same tokens, foo.token_stream
+    assert_equal 1, loads
+  end
+
   def test_pop_token
     foo = Class.new do
       include RDoc::TokenStream
@@ -95,6 +108,20 @@ def test_pop_token
     assert_equal [], foo.token_stream
   end
 
+  def test_mutating_deferred_tokens
+    foo = Class.new do
+      include RDoc::TokenStream
+    end.new
+    tokens = [:first]
+    foo.collect_tokens(:ruby, loader: -> { tokens })
+
+    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