Lazy code colorization

lib/rdoc/generator/markup.rb

@@ -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 '' if !(tokens = token_stream)
 
-    src = RDoc::TokenStream.to_html @token_stream
+    src = RDoc::TokenStream.to_html tokens
 
     # dedent the source
     common_indent = src.length

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
@@ -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,
@@ -458,7 +450,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
@@ -553,12 +545,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)
@@ -699,7 +685,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
@@ -719,12 +705,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:, 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:, 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
@@ -750,10 +736,8 @@ 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
+    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'
     # defined in c and ruby.
@@ -1175,8 +1159,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,
@@ -1186,7 +1168,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

lib/rdoc/parser/ruby_colorizer.rb

@@ -10,6 +10,48 @@ module RDoc::Parser::RubyColorizer
 
   ColoredToken = Struct.new(:kind, :text)
 
+  # 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
+
+    #: (Integer) -> ^() -> Array[ColoredToken]
+    def token_stream_loader(node_id)
+      tokens = nil
+      (@streams[node_id] ||= []) << ->(resolved_tokens) { tokens = resolved_tokens }
+      -> { @source ? (materialize; tokens) : tokens }
+    end
+
+    #: () -> void
+    def materialize
+      @mutex.synchronize do
+        return unless @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?
+          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
+
+        @streams.each do |node_id, streams|
+          streams.each { |resolve| resolve.call(staged_tokens.fetch(node_id)) }
+        end
+        @streams = nil
+        @source = nil
+      end
+    end
+  end
+
   # Prism operator token types except assignment '='
   OP_TOKENS = %i[
     AMPERSAND AMPERSAND_AMPERSAND

lib/rdoc/token_stream.rb

@@ -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)
-    @token_stream = []
+  def collect_tokens(language, loader: nil)
+    loaded_tokens = nil
+    @token_stream = loader ? -> { loaded_tokens ||= loader.call } : []
     @token_stream_language = language
   end
 
@@ -75,14 +77,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.is_a?(Proc) ? @token_stream.call : @token_stream
   end
 
   ##

test/rdoc/parser/ruby_colorizer_test.rb

@@ -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)
+
+    2.times { assert_raise(KeyError) { loader.call } }
+    assert_raise(KeyError) { missing_loader.call }
+  end
+
   def test_partial_colorize
     code = <<~RUBY
       class A

test/rdoc/rdoc_token_stream_test.rb

@@ -85,6 +85,20 @@ 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] })
+    foo.freeze
+
+    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 +109,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
@@ -126,4 +154,5 @@ def initialize
     end.new
     assert_equal "", foo.tokens_to_s
   end
+
 end