Add source map support to track AST node output locations

lib/unparser.rb

@@ -122,6 +122,41 @@ def self.unparse(
   end
   # rubocop:enable Metrics/ParameterLists
 
+  # Unparse an AST node into a string with a source map
+  #
+  # The source map records the mapping from each AST node to its
+  # character range in the generated output string.
+  #
+  # @param [Parser::AST::Node, nil] node
+  # @param [Array] comments
+  # @param [Encoding, nil] explicit_encoding
+  # @param [Set<Symbol>] static_local_variables
+  #
+  # @return [Array(String, SourceMap)]
+  #
+  # @raise InvalidNodeError
+  #   if the node passed is invalid
+  #
+  # @api public
+  #
+  # rubocop:disable Metrics/ParameterLists
+  def self.unparse_with_source_map(
+    node,
+    comments:               EMPTY_ARRAY,
+    explicit_encoding:      nil,
+    static_local_variables: Set.new
+  )
+    unparse_ast_with_source_map(
+      AST.new(
+        comments:               comments,
+        explicit_encoding:      explicit_encoding,
+        node:                   node,
+        static_local_variables: static_local_variables
+      )
+    )
+  end
+  # rubocop:enable Metrics/ParameterLists
+
   # Unparse an AST
   #
   # @param [AST] ast
@@ -136,23 +171,59 @@ def self.unparse(
   #
   # @api public
   def self.unparse_ast(ast)
-    return EMPTY_STRING if ast.node.nil?
+    emit_ast(ast).content
+  end
+
+  # Unparse an AST with source map
+  #
+  # @param [AST] ast
+  #
+  # @return [Array(String, SourceMap)]
+  #
+  # @raise InvalidNodeError
+  #   if the node passed is invalid
+  #
+  # @api public
+  #
+  def self.unparse_ast_with_source_map(ast)
+    source_map = SourceMap.new
+    source     = emit_ast(ast, source_map: source_map).content
+    source_map.freeze
+
+    [source, source_map]
+  end
+
+  # Emit AST into a buffer
+  #
+  # @param [AST] ast
+  # @param [SourceMap, nil] source_map
+  #
+  # @return [Buffer]
+  #
+  # @api private
+  #
+  def self.emit_ast(ast, source_map: nil)
+    buffer = Buffer.new(source_map: source_map)
+    return buffer if ast.node.nil?
 
     local_variable_scope = AST::LocalVariableScope.new(
       node:                   ast.node,
       static_local_variables: ast.static_local_variables
     )
 
-    Buffer.new.tap do |buffer|
+    buffer.record_node(ast.node) do
       Emitter::Root.new(
         buffer:               buffer,
         comments:             Comments.new(ast.comments),
         explicit_encoding:    ast.explicit_encoding,
         local_variable_scope: local_variable_scope,
         node:                 ast.node
       ).write_to_buffer
-    end.content
+    end
+
+    buffer
   end
+  private_class_method :emit_ast
 
   # Unparse AST either
   #
@@ -263,6 +334,7 @@ def self.buffer(source, identification = '(string)')
 require 'unparser/node_helpers'
 require 'unparser/ast'
 require 'unparser/ast/local_variable_scope'
+require 'unparser/source_map'
 require 'unparser/buffer'
 require 'unparser/generation'
 require 'unparser/color'

lib/unparser/buffer.rb

@@ -13,11 +13,28 @@ class Buffer
     #
     # @api private
     #
-    def initialize
-      @content  = +''
-      @heredocs = []
-      @indent   = 0
-      @no_nl    = true
+    def initialize(source_map: nil)
+      @content    = +''
+      @heredocs   = []
+      @indent     = 0
+      @no_nl      = true
+      @source_map = source_map
+    end
+
+    # Return the source map, if any
+    #
+    # @return [SourceMap, nil]
+    #
+    # @api private
+    attr_reader :source_map
+
+    # Return the current write position
+    #
+    # @return [Integer]
+    #
+    # @api private
+    def position
+      @content.length
     end
 
     # Append string
@@ -154,6 +171,22 @@ def write_encoding(encoding)
       write("# -*- encoding: #{encoding} -*-\n")
     end
 
+    # Record a node's output range in the source map
+    #
+    # @param node [Parser::AST::Node]
+    #
+    # @return [Object] the block's return value
+    def record_node(node)
+      unless @source_map
+        return yield
+      end
+
+      start_pos = position
+      result    = yield
+      @source_map.record(node: node, generated_range: start_pos...position)
+      result
+    end
+
   private
 
     INDENT_SPACE = '  '.freeze

lib/unparser/generation.rb

@@ -228,12 +228,13 @@ def writer_with(klass, node:, **attributes)
       klass.new(to_h.merge(node: node, **attributes))
     end
 
+    # mutant:disable
     def visit(node)
-      emitter(node).write_to_buffer
+      buffer.record_node(node) { emitter(node).write_to_buffer }
     end
 
     def visit_deep(node)
-      emitter(node).tap(&:write_to_buffer)
+      buffer.record_node(node) { emitter(node).tap(&:write_to_buffer) }
     end
 
     def first_child

lib/unparser/source_map.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module Unparser
+  # Maps AST nodes to their generated output ranges
+  class SourceMap
+    # Single mapping entry from an AST node to its range in generated output
+    class Entry
+      attr_reader :node, :generated_range
+
+      def initialize(node:, generated_range:)
+        @node            = node
+        @generated_range = generated_range
+        freeze
+      end
+    end # Entry
+
+    attr_reader :entries
+
+    def initialize
+      @entries = []
+    end
+
+    # Record a node mapping
+    #
+    # @param node [Parser::AST::Node]
+    # @param generated_range [Range]
+    #
+    # @return [self]
+    def record(node:, generated_range:)
+      @entries << Entry.new(node: node, generated_range: generated_range)
+      self
+    end
+
+    # Find all entries for a specific node (by identity)
+    #
+    # @param node [Parser::AST::Node]
+    #
+    # @return [Array<Entry>]
+    def for_node(node)
+      @entries.select { |entry| entry.node.equal?(node) }
+    end
+
+    # Freeze the source map and its entries
+    #
+    # @return [self]
+    def freeze
+      @entries.freeze
+      super
+    end
+  end # SourceMap
+end # Unparser