ast_transform 2.1.2 → 3.0.0

data/lib/ast_transform/line_aligned_emitter.rb

@@ -0,0 +1,224 @@
+# frozen_string_literal: true
+
+require 'ast_transform/node'
+require 'ast_transform/layout'
+require 'ast_transform/statement_renderer'
+require 'ast_transform/thunk_lowering'
+
+module ASTTransform
+  # Emits a transformed AST as text in which every loc-carrying statement occupies its original source line, so that
+  # backtraces, breakpoints and debugger display are correct by construction — CRuby derives line numbers from
+  # physical text position, so placement is our line table.
+  #
+  # Placement policy over statement sequences:
+  #
+  # 1. Statement has loc: target its source line — the Layout pads to reach it, or packs (`; `) when the cursor has
+  #    already passed it. A user statement packing means the transform moved it — the alignment auditor's concern,
+  #    not a runtime failure.
+  # 2. No loc: pack onto the current line — synthetic code has no source-line truth to preserve.
+  #
+  # The emitter owns the Ruby knowledge: walking statement structure, deciding which line each node targets, and
+  # that keywords can never be `;`-packed. The pad-or-pack mechanics live in Layout; statement-to-text rendering
+  # (and its isolation workarounds) in StatementRenderer — both created per emission, so the emitter itself is
+  # stateless and an instance is a reusable collaborator.
+  #
+  # Thunk nodes are lowered (ThunkLowering) before layout; the emitter's postcondition is that no custom node type
+  # (ast_* markers or types registered on ASTTransform::Node) crosses the unparse boundary — they are IR between
+  # stages that understand them.
+  class LineAlignedEmitter
+    # Raised as the emitter's postcondition when a custom node type (ast_* markers or types registered on
+    # ASTTransform::Node) reaches the unparse boundary instead of being lowered by the stage that understands it.
+    class UnloweredNodeTypeError < StandardError; end
+
+    # Containers the emitter recurses into so nested statements align; every other node renders as an Unparser blob
+    # at its head line.
+    RECURSIVE_CONTAINER_TYPES = [:class, :module, :sclass, :def, :defs, :block, :numblock, :itblock, :kwbegin].freeze
+    BODY_INDEXES = {
+      class: 2, module: 1, sclass: 1, def: 2, defs: 3, block: 2, numblock: 2, itblock: 2
+    }.freeze
+    # Assignments whose value is a block (e.g. the lowered thunk proc) recurse into the block so its body statements
+    # align.
+    ASSIGNMENT_TYPES = [:lvasgn, :ivasgn, :gvasgn, :casgn].freeze
+    BLOCK_VALUE_TYPES = [:block, :numblock, :itblock].freeze
+
+    # @param thunk_lowering [ThunkLowering] the lowering run ahead of emission.
+    def initialize(thunk_lowering: ThunkLowering.new)
+      @thunk_lowering = thunk_lowering
+    end
+
+    # @param ast [Parser::AST::Node] transformed AST
+    # @param source_path [String] original file path (for error messages)
+    # @return [String] transformed source, line-aligned
+    # @raise [ThunkLowering::PlacementError] if a thunk cannot be textually placed
+    # @raise [UnloweredNodeTypeError] if a custom node type survived to emission
+    def emit(ast, source_path)
+      lowered = @thunk_lowering.lower(ast)
+      assert_no_custom_types(lowered, source_path)
+
+      layout = Layout.new
+      renderer = StatementRenderer.for_tree(lowered)
+      emit_statements(statements_of(lowered), layout, renderer)
+      layout.to_source
+    end
+
+    private
+
+    def emit_statements(statements, layout, renderer)
+      statements.each { |statement| emit_statement(statement, layout, renderer) }
+    end
+
+    def emit_statement(node, layout, renderer)
+      if recursive_container?(node)
+        emit_container(node, layout, renderer)
+      else
+        layout.place(node.loc&.line, renderer.aligned_render(node), column: node.loc&.column)
+      end
+    end
+
+    # Emits a container body that may be a bare :ensure/:rescue node (their begin/end context comes from the
+    # surrounding def/block/kwbegin, so the keywords must be emitted inline, aligned like statements).
+    def emit_body(body, layout, renderer)
+      case body&.type
+      when :ensure then emit_ensure(body, layout, renderer)
+      when :rescue then emit_rescue(body, layout, renderer)
+      else emit_statements(statements_of(body), layout, renderer)
+      end
+    end
+
+    def emit_ensure(node, layout, renderer)
+      *body, ensurer = node.children
+      body.each { |statement| emit_body(statement, layout, renderer) }
+      place_keyword(layout, keyword_line(node), 'ensure', column: keyword_column(node))
+      emit_statements(statements_of(ensurer), layout, renderer)
+    end
+
+    def emit_rescue(node, layout, renderer)
+      body, *resbodies, else_body = node.children
+      emit_body(body, layout, renderer)
+      resbodies.each { |resbody| emit_resbody(resbody, layout, renderer) }
+      return if else_body.nil?
+
+      else_range = node.loc.else if node.loc.respond_to?(:else)
+      place_keyword(layout, else_range&.line, 'else', column: else_range&.column)
+      emit_statements(statements_of(else_body), layout, renderer)
+    end
+
+    def emit_resbody(node, layout, renderer)
+      exceptions, capture, body = node.children
+      header = ['rescue']
+      header << " #{renderer.unparse(exceptions).delete_prefix('[').delete_suffix(']')}" if exceptions
+      header << " => #{capture.children[0]}" if capture
+      place_keyword(layout, node.loc&.line, header.join, column: node.loc&.column)
+      emit_statements(statements_of(body), layout, renderer)
+    end
+
+    # Keywords (rescue/ensure/else) cannot be `;`-packed after a statement; when their line is taken they go on a
+    # fresh line instead.
+    def place_keyword(layout, target_line, keyword, column: nil)
+      if target_line && target_line > layout.cursor
+        layout.place(target_line, keyword, column: column)
+      else
+        layout.place_on_fresh_line(keyword)
+      end
+    end
+
+    def keyword_line(node)
+      loc = node.loc
+      loc.keyword.line if loc.respond_to?(:keyword) && loc.keyword
+    end
+
+    def keyword_column(node)
+      loc = node.loc
+      loc.keyword.column if loc.respond_to?(:keyword) && loc.keyword
+    end
+
+    def recursive_container?(node)
+      RECURSIVE_CONTAINER_TYPES.include?(node.type) || block_assignment?(node)
+    end
+
+    def block_assignment?(node)
+      ASSIGNMENT_TYPES.include?(node.type) && node.children.last.is_a?(::Parser::AST::Node) &&
+        BLOCK_VALUE_TYPES.include?(node.children.last.type)
+    end
+
+    # Renders a container's opener and closer from the node with its body emptied, then recurses into the body so
+    # nested statements align.
+    def emit_container(node, layout, renderer)
+      opener, closer = container_delimiters(node, renderer)
+      layout.place(node.loc&.line, opener, column: node.loc&.column)
+      emit_body(container_body(node), layout, renderer)
+      layout.place(closer_line(node), closer, column: closer_column(node))
+    end
+
+    def container_delimiters(node, renderer)
+      rendered = renderer.unparse(empty_container(node)).split("\n").reject(&:empty?)
+      opener = rendered[0..-2].join("\n")
+      closer = rendered.last
+
+      # Unparser renders empty blocks with braces, but brace blocks cannot hold rescue/ensure bodies; do/end always
+      # can.
+      if opener.end_with?(' {') && closer == '}'
+        [opener.sub(/ \{\z/, ' do'), 'end']
+      else
+        [opener, closer]
+      end
+    end
+
+    def empty_container(node)
+      case node.type
+      when :kwbegin
+        node.updated(nil, [])
+      when *ASSIGNMENT_TYPES
+        block_node = node.children.last
+        emptied_block = block_node.updated(nil, [*block_node.children[0..-2], nil])
+        node.updated(nil, [*node.children[0..-2], emptied_block])
+      else
+        children = node.children.dup
+        children[BODY_INDEXES.fetch(node.type)] = nil
+        node.updated(nil, children)
+      end
+    end
+
+    def container_body(node)
+      case node.type
+      when :kwbegin then node.children.size == 1 ? node.children.first : node.updated(:begin, node.children)
+      when *ASSIGNMENT_TYPES then node.children.last.children[2]
+      else node.children[BODY_INDEXES.fetch(node.type)]
+      end
+    end
+
+    # The line the container's `end`/`}` occupies in the source, when known.
+    def closer_line(node)
+      loc = node.loc
+      loc.end.line if loc.respond_to?(:end) && loc.end
+    end
+
+    def closer_column(node)
+      loc = node.loc
+      loc.end.column if loc.respond_to?(:end) && loc.end
+    end
+
+    # Loc-less :begin nodes in statement position (e.g. a lowered thunk in a single-statement container body, carried
+    # with its placement) are grouping, not structure: flatten them so each inner statement is laid out independently.
+    def statements_of(body)
+      return [] if body.nil?
+      return [body] unless body.type == :begin
+
+      body.children.flat_map do |child|
+        child.is_a?(::Parser::AST::Node) && child.type == :begin && child.loc.nil? ? statements_of(child) : [child]
+      end
+    end
+
+    def assert_no_custom_types(node, source_path)
+      return unless node.is_a?(::Parser::AST::Node)
+
+      if node.type.start_with?('ast_') || Node.registry.key?(node.type)
+        raise UnloweredNodeTypeError,
+          "custom node type :#{node.type} reached emission in #{source_path}; custom types are " \
+            "IR between transformation stages and must be lowered by the stage that understands them"
+      end
+
+      node.children.each { |child| assert_no_custom_types(child, source_path) }
+    end
+  end
+end

data/lib/ast_transform/node.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require 'parser'
+
+module ASTTransform
+  # Base class for custom IR nodes. Transform authors subclass it and register a custom node type to get type-routed
+  # construction from +s+ with domain accessors:
+  #
+  #   class InteractionNode < ASTTransform::Node
+  #     register :rspock_interaction
+  #
+  #     def cardinality = children[0]
+  #   end
+  #
+  #   s(:rspock_interaction, ...) # => InteractionNode
+  #
+  # Custom node *types* are IR between stages that understand them and must be lowered before emission (the emitter
+  # enforces this). Standard-typed nodes deliberately stay plain Parser::AST::Node everywhere — parsed and +s+-built
+  # alike: AST::Node#eql? compares class, and Unparser verifies dynamic-string emission by re-parsing and
+  # eql?-comparing, so a custom class on a standard type breaks emission.
+  class Node < ::Parser::AST::Node
+    class << self
+      # Registers +self+ as the class to construct for +type+ nodes.
+      #
+      # @param type [Symbol] the custom node type routed to this class
+      # @return [void]
+      def register(type)
+        Node.registry[type] = self
+      end
+
+      # Builds a node of +type+: registered types construct their custom class, everything else a plain
+      # Parser::AST::Node.
+      #
+      # @param type [Symbol] node type
+      # @param children [Array] child nodes / literals
+      # @param properties [Hash] node properties (e.g. location:)
+      # @return [Parser::AST::Node]
+      def build(type, children, properties = {})
+        klass = Node.registry.fetch(type, ::Parser::AST::Node)
+        klass.new(type, children, properties)
+      end
+
+      def registry
+        @registry ||= {}
+      end
+    end
+  end
+end

data/lib/ast_transform/statement_renderer.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+require 'unparser'
+
+module ASTTransform
+  # Renders individual statements to text via Unparser, working around the two consequences of unparsing them in
+  # isolation (line-aligned emission places each statement independently, so each is ripped out of its context):
+  #
+  # * Isolation loses the surrounding scope's local variables, making Unparser re-parse identifiers as method calls
+  #   and fail its dstr round-trip verification — so a renderer is built once per tree with every local bound
+  #   anywhere in it, and feeds Unparser that set on each render.
+  # * Unparser normalizes some single-line constructs into multi-line form, which would push following statements
+  #   off their lines — so renders taller than their source are compressed back to one line when safely possible.
+  #
+  # Immutable: configured with the tree's locals at construction, no per-render state.
+  class StatementRenderer
+    # Node types that bind a local variable name: assignments plus every method/block parameter flavor.
+    LOCAL_BINDING_TYPES = [:lvasgn, :arg, :optarg, :restarg, :kwarg, :kwoptarg, :blockarg, :shadowarg].freeze
+
+    class << self
+      # Builds a renderer for statements of +node+'s tree, holding every local bound anywhere in it — an
+      # over-approximation that is safe because the set only informs Unparser's re-parse verification, never the
+      # rendered text.
+      def for_tree(node)
+        new(local_variables: collect_local_variables(node))
+      end
+
+      private
+
+      def collect_local_variables(node, names = Set.new)
+        return names unless node.is_a?(::Parser::AST::Node)
+
+        names << node.children[0] if LOCAL_BINDING_TYPES.include?(node.type) && node.children[0]
+        node.children.each { |child| collect_local_variables(child, names) }
+        names
+      end
+    end
+
+    # @param local_variables [Set<Symbol>] every local bound in the tree the statements come from.
+    def initialize(local_variables:)
+      @local_variables = local_variables
+    end
+
+    # @param node [Parser::AST::Node] the statement to render.
+    # @return [String] Unparser's render, informed of the tree's locals.
+    def unparse(node)
+      Unparser.unparse(node, static_local_variables: @local_variables)
+    end
+
+    # Renders +node+ no taller than its source when safely possible. Unparser normalizes some single-line
+    # constructs into multi-line form (e.g. modifier-if into if/end); when the render is taller than the
+    # statement's source, compress it back to one line — verified by re-parse so a statement that cannot be safely
+    # single-lined (e.g. containing a heredoc) falls back to its multi-line render.
+    def aligned_render(node)
+      render = unparse(node)
+      loc = node.loc
+      return render unless loc.respond_to?(:last_line) && loc.line
+
+      source_height = loc.last_line - loc.line + 1
+      return render if render.count("\n") < source_height
+
+      compress_to_single_line(render) || render
+    end
+
+    private
+
+    def compress_to_single_line(render)
+      candidate = render.split("\n").map(&:strip).join('; ')
+      # Both sides parsed without scope context, so lvar/send ambiguity cancels out; equality means the newline
+      # join preserved structure.
+      Unparser.parse(candidate) == Unparser.parse(render) ? candidate : nil
+    rescue Parser::SyntaxError
+      nil
+    end
+  end
+end

data/lib/ast_transform/testing/assertions.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require 'ast_transform/transformer'
+require 'ast_transform/instruction_sequence'
+
+module ASTTransform
+  module Testing
+    # Assertions for transform authors' own test suites — the enforcement arm of the authoring contract ("textual
+    # order is source order"). Minitest-flavored. Never loaded in production; require it from test code:
+    #
+    #   require "ast_transform/testing/assertions"
+    #
+    #   class MyTransformationTest < Minitest::Test
+    #     include ASTTransform::Testing::Assertions
+    #   end
+    module Assertions
+      # Transforms +source+ through the real pipeline (transform + line-aligned emission), re-parses both sides,
+      # matches surviving statements by location, and asserts each one's emitted line equals its source line.
+      # Statements the transform deletes (e.g. description strings) are exempt; statements the transform rewrites in
+      # place keep their anchor and are checked.
+      #
+      # @param source [String] fixture source
+      # @param transformations [Array<ASTTransform::AbstractTransformation>]
+      # @param path [String] pseudo-path used for parsing and messages
+      # @return [void]
+      def assert_line_aligned(source, *transformations, path: 'fixture.rb')
+        transformer = Transformer.new(*transformations)
+        emitted = transformer.transform_file_source(source, path, path)
+
+        source_lines_by_statement = statement_lines(transformer.build_ast(source, file_path: path))
+        emitted_lines_by_statement = statement_lines(transformer.build_ast(emitted, file_path: path))
+
+        misaligned = source_lines_by_statement.filter_map do |render, source_line|
+          emitted_line = emitted_lines_by_statement[render]
+          next if emitted_line.nil? || emitted_line == source_line
+
+          format('  MISALIGNED %s: source line %d, emitted line %d', render, source_line, emitted_line)
+        end
+
+        assert misaligned.empty?, <<~MESSAGE
+          expected every surviving statement at its source line in #{path}:
+          #{misaligned.join("\n")}
+
+          emitted:
+          #{numbered_listing(emitted)}
+        MESSAGE
+      end
+
+      # Runtime complement of assert_line_aligned: compiles +source+ through the full pipeline under +path+, executes
+      # it, and asserts the raw first backtrace frame — no filtering of any kind — is "<path>:<raise_at>".
+      #
+      # @param source [String] fixture that raises when executed
+      # @param path [String] pseudo source path to compile under
+      # @param raise_at [Integer] expected source line of the raise
+      # @return [void]
+      def assert_backtrace_lines(source, path:, raise_at:)
+        iseq = InstructionSequence.source_to_transformed_iseq(source, path)
+
+        error = assert_raises(StandardError, "fixture at #{path} should raise when executed") do
+          iseq.eval
+        end
+
+        location = error.backtrace_locations.first
+        assert_equal "#{location.path}:#{raise_at}", "#{location.path}:#{location.lineno}",
+          "raw backtrace should cite source line #{raise_at} of #{path}"
+      end
+
+      private
+
+      # Flat statement renders and their first line, keyed by unparsed text so source and emitted sides can be matched
+      # without location identity. Duplicate renders keep their first occurrence — good enough for fixtures, which
+      # authors control.
+      def statement_lines(ast, lines = {})
+        return lines unless ast.is_a?(::Parser::AST::Node)
+
+        if statement_sequence?(ast)
+          ast.children.each do |statement|
+            next unless statement.is_a?(::Parser::AST::Node) && statement.loc&.expression
+
+            lines[Unparser.unparse(statement)] ||= statement.loc.line
+          end
+        end
+
+        ast.children.each { |child| statement_lines(child, lines) }
+        lines
+      end
+
+      def statement_sequence?(node)
+        [:begin, :kwbegin].include?(node.type)
+      end
+
+      def numbered_listing(source)
+        source.lines.map.with_index(1) { |line, number| format('%3d| %s', number, line) }.join
+      end
+    end
+  end
+end

data/lib/ast_transform/thunk.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+require 'ast_transform/node'
+
+module ASTTransform
+  # The reordering primitive: an eagerly built wrapper node, spliced wherever the wrapped statements must EXECUTE —
+  # statement position or composed inside an expression (e.g. an assert_raises block body). Its body keeps its own
+  # source locations, and the lowering derives the wrapper's textual placement from them (see ThunkLowering), so the
+  # statements still emit on their original lines even though execution waits.
+  #
+  # Children are +[id, *body_statements]+ and the invariants are enforced here in +initialize+, which every
+  # construction path shares — +s+ routing, the +thunk+ helper, and Processor rebuilds (+updated+ re-initializes).
+  # Build thunks with +TransformationHelper#thunk+; reuse the same node to execute one body from several points
+  # (multiplexing).
+  #
+  # Runtime semantics are near-transparent (proc lowering): +return+ still returns from the enclosing method, and
+  # locals the body assigns stay method-scope. See ThunkLowering for the full contract.
+  class Thunk < Node
+    register :ast_thunk
+
+    # The identity of a Thunk across transformation passes: Processor and Node#updated rebuilds create new node
+    # objects, so node identity does not survive — but children DO (carried by reference through every rebuild). Every
+    # rebuild of a thunk therefore carries this same id object, and the lowering groups occurrences by its object
+    # identity: one proc, one call per occurrence. Minted internally by the +thunk+ helper, never handled by authors.
+    # No behavior — a named class over a bare Object.new only for self-documenting AST dumps and greppability.
+    class Id; end
+
+    def initialize(type, children, properties = {})
+      id, *body = children
+      unless id.is_a?(ASTTransform::Thunk::Id)
+        raise ArgumentError,
+          "a Thunk's first child must be its #{Thunk::Id} (got #{id.class}); build thunks with the " \
+            "thunk(*statements) helper"
+      end
+      raise ArgumentError, "a Thunk must wrap at least one statement" if body.empty?
+
+      # Captured before super (which freezes the node); frozen so the shared array cannot be mutated out from under
+      # +children+. Rebuilds via +updated+ re-run initialize, so the capture can never go stale.
+      @id = id
+      @body = body.freeze
+
+      super
+    end
+
+    # Retrieves the Thunk's id.
+    # @return [ASTTransform::Thunk::Id] The Id.
+    attr_reader(:id)
+
+    # Retrieves the Thunk's body.
+    #
+    # @note Same as the +Parser::AST::Node#children+
+    # @return [Array<Parser::AST::Node>] The nodes forming the body of the Thunk.
+    attr_reader(:body)
+  end
+end