ast_transform 2.1.2 → 3.0.0

data/lib/ast_transform/thunk_lowering.rb

@@ -0,0 +1,240 @@
+# frozen_string_literal: true
+
+require 'ast_transform/node'
+require 'ast_transform/thunk'
+require 'ast_transform/transformation_helper'
+
+module ASTTransform
+  # Lowers Thunk nodes into plain Ruby ahead of emission. Each unique thunk (grouped by id identity) becomes a
+  # hidden proc; each occurrence becomes the proc's call:
+  #
+  #   thunk placed at the execution point
+  #     => x = x; __ast_thunk_<n>__ = proc { body }   (at the body's source lines)
+  #        ...
+  #        __ast_thunk_<n>__.call                     (at the occurrence)
+  #
+  # Placement is inferred, not authored: the proc's text is inserted into the statement sequence enclosing the
+  # occurrence, positioned among its siblings by the body's first source line — the lines the author removed the
+  # statements from. A loc-less body has no textual home and packs immediately before its call. Placements never
+  # escape a scope boundary (def/class/module bodies absorb their own), because the hidden lvar must share the call's
+  # method activation; they DO escape block literals, which close over the defining scope.
+  #
+  # The closure is a non-lambda proc on purpose: `return` inside a proc returns from the method where the proc was
+  # defined, and placement and execution always share one method activation, so a thunked `return` keeps its original
+  # meaning. Jump keywords whose owner lies outside the body keep Ruby's native behavior (`break`/`retry` fail loudly,
+  # `next`/`redo` silently alter flow) — what a transform chooses to thunk is the transform author's call.
+  #
+  # The `x = x` pre-declarations cover every local the body assigns at method scope. A local first assigned inside a
+  # block literal is block-local, so without a textual method-scope assignment before the proc, thunked assignments
+  # would be invisible to the statements that read them after the execution point. Self-assignment registers the name
+  # (nil until the thunk runs — exactly what an unexecuted assignment yields) without clobbering an already-assigned
+  # value.
+  #
+  # Stateless: the thunks encountered during one lowering are tracked in a Registry created at +lower+ entry, so an
+  # instance is a reusable collaborator.
+  class ThunkLowering
+    include TransformationHelper
+
+    # A pending proc definition: +line+ is the body's first source line (nil for fully synthetic bodies),
+    # +statements+ the pre-declarations plus the proc assignment.
+    Placement = Struct.new(:line, :statements)
+
+    # Raised when a thunk cannot be placed: its body's source lines fall after the execution point (the hidden
+    # proc's text IS its assignment, so a call can never textually precede the body), or two occurrences of the
+    # same thunk carry diverging bodies.
+    class PlacementError < StandardError; end
+
+    # The thunks encountered during one lowering, keyed by identity: allocates each thunk's hidden lvar name
+    # on first occurrence and verifies later occurrences carry the same body.
+    class Registry
+      def initialize
+        @names_by_id = {}.compare_by_identity
+        @bodies_by_id = {}.compare_by_identity
+      end
+
+      def known?(id)
+        @names_by_id.key?(id)
+      end
+
+      # @return [Symbol] the hidden lvar name allocated for +id+.
+      def register(id, body)
+        name = :"__ast_thunk_#{@names_by_id.size + 1}__"
+        @names_by_id[id] = name
+        @bodies_by_id[id] = body
+        name
+      end
+
+      def name_for(id)
+        @names_by_id.fetch(id)
+      end
+
+      def verify_same_body!(id, body)
+        return if @bodies_by_id[id] == body
+
+        raise PlacementError,
+          'occurrences of one thunk carry diverging bodies; reuse the same thunk node to multiplex'
+      end
+
+      def hidden_names
+        @names_by_id.values
+      end
+    end
+    private_constant :Registry
+
+    SEQUENCE_TYPES = [:begin, :kwbegin].freeze
+    # Scope-opening containers: the hidden lvar cannot be referenced across these boundaries, so placements arising
+    # inside must land inside.
+    SCOPE_BODY_INDEXES = { def: 2, defs: 3, class: 2, module: 1, sclass: 1 }.freeze
+
+    # @param node [Parser::AST::Node] tree possibly containing Thunk nodes
+    # @return [Parser::AST::Node] tree with thunks lowered to plain Ruby
+    # @raise [PlacementError] when a thunk body's source lines fall after its execution point, or occurrences of
+    #   one thunk diverge
+    def lower(node)
+      lower_body(node, Registry.new)
+    end
+
+    private
+
+    # Lowers a node standing in statement-body position (a container's body or the root), absorbing any placements
+    # that arise within it.
+    def lower_body(node, registry)
+      return node unless node.is_a?(::Parser::AST::Node)
+      return lower_sequence(node, registry) if SEQUENCE_TYPES.include?(node.type)
+
+      lowered, placements = lower_expression(node, registry)
+      return lowered if placements.empty?
+
+      # A loc-less :begin in statement position; the emitter flattens it into the surrounding statement stream.
+      s(:begin, *placements.flat_map(&:statements), lowered)
+    end
+
+    # Lowers a statement sequence, inserting each placement among the statements by the body's source line.
+    def lower_sequence(node, registry)
+      statements = []
+
+      node.children.each_with_index do |child, index|
+        lowered, placements = lower_expression(child, registry)
+        placements.each do |placement|
+          check_placement_precedes_execution!(placement, child, node.children[(index + 1)..])
+          statements.insert(insertion_index(statements, placement), *placement.statements)
+        end
+        statements << lowered
+      end
+
+      node.updated(nil, statements)
+    end
+
+    # Lowers a node in expression position. Returns the lowered node and the placements that must be inserted into
+    # the enclosing statement sequence.
+    #
+    # @return [Array(Parser::AST::Node, Array<Placement>)]
+    def lower_expression(node, registry)
+      return [node, []] unless node.is_a?(::Parser::AST::Node)
+
+      case node.type
+      when :ast_thunk
+        lower_thunk(node, registry)
+      when *SEQUENCE_TYPES
+        [lower_sequence(node, registry), []]
+      when :ensure, :rescue
+        [node.updated(nil, node.children.map { |child| lower_body(child, registry) }), []]
+      when :resbody
+        exceptions, capture, body = node.children
+        [node.updated(nil, [exceptions, capture, lower_body(body, registry)]), []]
+      else
+        lower_generic(node, registry)
+      end
+    end
+
+    def lower_generic(node, registry)
+      scope_body_index = SCOPE_BODY_INDEXES[node.type]
+      pending = []
+
+      children = node.children.each_with_index.map do |child, index|
+        if index == scope_body_index
+          lower_body(child, registry)
+        else
+          lowered, placements = lower_expression(child, registry)
+          pending.concat(placements)
+          lowered
+        end
+      end
+
+      [node.updated(nil, children), pending]
+    end
+
+    # An occurrence of a thunk: the first occurrence of its id yields the placement; every occurrence yields the
+    # call.
+    def lower_thunk(node, registry)
+      id = node.id
+
+      if registry.known?(id)
+        registry.verify_same_body!(id, node.body)
+        return [call_node(id, registry), []]
+      end
+
+      name = registry.register(id, node.body)
+      lowered_body = lower_sequence(s(:begin, *node.body), registry)
+      placement = Placement.new(body_first_line(node.body), placement_statements(name, lowered_body, registry))
+      [call_node(id, registry), [placement]]
+    end
+
+    def call_node(id, registry)
+      s(:send, s(:lvar, registry.name_for(id)), :call)
+    end
+
+    def placement_statements(name, lowered_body, registry)
+      assignment = s(:lvasgn, name, s(:block, s(:send, nil, :proc), s(:args), lowered_body))
+      hidden_names = registry.hidden_names
+      pre_declared = method_scope_assignments(lowered_body).reject { |local| hidden_names.include?(local) }
+      pre_declared.map { |local| s(:lvasgn, local, s(:lvar, local)) } << assignment
+    end
+
+    def body_first_line(body)
+      body.filter_map { |statement| statement.loc&.line }.min
+    end
+
+    # The proc's text must precede its call: a placement whose body lines fall at or after the executing statement
+    # (or any statement after it) cannot be laid out — the assignment would complete after the call.
+    def check_placement_precedes_execution!(placement, executing_statement, following_statements)
+      return if placement.line.nil?
+
+      conflicting = [executing_statement, *following_statements].find do |statement|
+        line = statement.is_a?(::Parser::AST::Node) ? statement.loc&.line : nil
+        line && line < placement.line
+      end
+      return if conflicting.nil?
+
+      raise PlacementError,
+        "thunk body's source lines (from line #{placement.line}) fall after its execution point " \
+          "(statement at line #{conflicting.loc.line}); a thunk can only delay execution, never text"
+    end
+
+    def insertion_index(statements, placement)
+      return statements.size if placement.line.nil?
+
+      statements.index { |statement| statement.loc&.line && statement.loc.line > placement.line } || statements.size
+    end
+
+    # Node types opening a new local-variable scope: assignments inside them were invisible to the method scope in
+    # the original source too, so they get no pre-declaration.
+    NEW_SCOPE_TYPES = [:def, :defs, :class, :module, :sclass].freeze
+    # Block literals: locals first assigned inside them are block-local (the same lexical rule the pre-declarations
+    # exist to work around), but their callee/arguments evaluate at method scope and are still descended.
+    BLOCK_TYPES = [:block, :numblock, :itblock].freeze
+
+    # Locals the thunk body assigns at method scope, in first-assignment order (covers masgn/op_asgn targets — they
+    # all carry :lvasgn nodes).
+    def method_scope_assignments(node, names = [])
+      return names unless node.is_a?(::Parser::AST::Node)
+      return names if NEW_SCOPE_TYPES.include?(node.type)
+
+      names << node.children[0] if node.type == :lvasgn && !names.include?(node.children[0])
+
+      children = BLOCK_TYPES.include?(node.type) ? [node.children[0]] : node.children
+      children.each { |child| method_scope_assignments(child, names) }
+      names
+    end
+  end
+end

data/lib/ast_transform/transformation.rb

@@ -1,8 +1,9 @@
 # frozen_string_literal: true
-require 'ast_transform'
-require 'ast_transform/abstract_transformation'
-require 'ast_transform/transformer'
-require 'unparser'
+
+require "ast_transform"
+require "ast_transform/abstract_transformation"
+require "ast_transform/transformer"
+require "unparser"
 
 module ASTTransform
   class Transformation < ASTTransform::AbstractTransformation
@@ -18,9 +19,9 @@ module ASTTransform
 
       count_before_reject = children.size
 
-      children.reject!.with_index { |child_node, index|
+      children.reject!.with_index do |child_node, index|
         transform_node?(child_node) && transformable_node?(next_child(node, index))
-      }
+      end
 
       processed = process_all(children)
 
@@ -68,19 +69,14 @@ module ASTTransform
 
     def extract_transformation(node)
       return unless node.is_a?(Parser::AST::Node)
-      return unless node.children.count >= 2
-
-      if node.children[1] == :new
-        require_transformation(node)
-        code = Unparser.unparse(node)
+      return if node.children.count < 2
 
-        TOPLEVEL_BINDING.eval(code)
-      else
-        require_transformation(node)
-        code = "#{Unparser.unparse(node)}.new"
+      require_transformation(node)
+      code = Unparser.unparse(node)
+      # A bare constant reference is instantiated; an explicit .new is kept as written.
+      code = "#{code}.new" unless node.children[1] == :new
 
-        TOPLEVEL_BINDING.eval(code)
-      end
+      TOPLEVEL_BINDING.eval(code)
     end
 
     def require_transformation(node)
@@ -100,11 +96,12 @@ module ASTTransform
       acronyms = ASTTransform.acronyms
       acronym_regex = acronyms.empty? ? /(?=a)b/ : /#{acronyms.join("|")}/
       return const_name unless /[A-Z-]|::/.match?(const_name)
-      word = const_name.to_s.gsub("::".freeze, "/".freeze)
-      word.gsub!(/(?:(?<=([A-Za-z\d]))|\b)(#{acronym_regex})(?=\b|[^a-z])/) { "#{$1 && '_'.freeze }#{$2.downcase}" }
-      word.gsub!(/([A-Z\d]+)([A-Z][a-z])/, '\1_\2'.freeze)
-      word.gsub!(/([a-z\d])([A-Z])/, '\1_\2'.freeze)
-      word.tr!("-".freeze, "_".freeze)
+
+      word = const_name.to_s.gsub("::", "/")
+      word.gsub!(/(?:(?<=([A-Za-z\d]))|\b)(#{acronym_regex})(?=\b|[^a-z])/) { "#{::Regexp.last_match(1) && "_"}#{::Regexp.last_match(2).downcase}" }
+      word.gsub!(/([A-Z\d]+)([A-Z][a-z])/, '\1_\2')
+      word.gsub!(/([a-z\d])([A-Z])/, '\1_\2')
+      word.tr!("-", "_")
       word.downcase!
       word
     end

data/lib/ast_transform/transformation_helper.rb

@@ -1,16 +1,127 @@
 # frozen_string_literal: true
+
 require 'parser'
+require 'ast_transform/node'
+require 'ast_transform/thunk'
 
 module ASTTransform
+  # The transform-authoring layer. Three shapes:
+  #
+  # - Constructors (+s+, +s_at+): type + children in, fresh node out.
+  # - The sequence combinator (+run_after+): sequence in, sequence out — the paved road for execution reordering.
+  # - The low-level reordering primitive (+thunk+): statements in, Thunk node out — for execution points inside
+  #   expressions.
+  #
+  # The contract these helpers serve: textual order is source order. The emitter places every loc-carrying statement
+  # at its source line; when execution order must differ from textual order, authors express it as a thunk instead of
+  # moving text.
   module TransformationHelper
-    def self.included(base)
-      base.extend(Methods)
-      base.include(Methods)
+    # Raised by authoring helpers (e.g. +s_at+) when a node that must carry a source location does not have one.
+    class MissingLocationError < StandardError; end
+
+    class << self
+      def included(base)
+        base.extend(Methods)
+        base.include(Methods)
+      end
     end
 
     module Methods
+      # Builds a loc-less node. The emitter packs loc-less nodes onto the current output line — the correct default
+      # for synthetic code, which has no source-line truth to preserve.
+      #
+      # @param type [Symbol] node type
+      # @param children [Array] child nodes / literals
+      # @param properties [Hash] node properties (e.g. location:)
+      # @return [ASTTransform::Node] node routed to its registered class
       def s(type, *children, **properties)
-        Parser::AST::Node.new(type, children, properties)
+        Node.build(type, children, properties)
+      end
+
+      # Builds a fresh node anchored to another node's source location. Use when composing a replacement tree whose
+      # root isn't derived from the node it replaces (otherwise prefer +anchor.updated(...)+). The attached map is a
+      # clean expression-only Source::Map over +anchor.loc.expression+ — no stale typed sub-ranges (selector etc.).
+      # Anchor inheritance is shallow; children keep or lack their own locs.
+      #
+      # @param anchor [Parser::AST::Node] node whose line this code replaces
+      # @param type [Symbol] node type
+      # @param children [Array] child nodes / literals
+      # @return [ASTTransform::Node] node carrying anchor's expression range
+      # @raise [MissingLocationError] if anchor has no expression location
+      def s_at(anchor, type, *children)
+        expression = anchor.loc&.expression
+        raise MissingLocationError, "anchor #{anchor.type} node has no source location" unless expression
+
+        s(type, *children, location: ::Parser::Source::Map.new(expression))
+      end
+
+      # The low-level reordering primitive. Thunking is the one reordering lever: text never moves and execution can
+      # only move later, so "hoist A above B" is expressed as "run B after A". Returns a single Thunk node: splice it
+      # where the statements must RUN — statement position or composed inside an expression, e.g. as an assert_raises
+      # block body. The wrapped statements keep their own locs, and the lowering derives the hidden proc's textual
+      # placement from them, so the body still emits on its source lines even though execution waits. Reuse the same
+      # node to execute one body from several points.
+      #
+      # Semantics are near-transparent (see ThunkLowering): +return+ still returns from the enclosing method
+      # (non-lambda proc), and locals the wrapped statements assign stay method-scope (pre-declared before the proc).
+      # Jump keywords whose owner lies outside the wrapped statements keep Ruby's native behavior — +break+/+retry+
+      # fail loudly at the jump's own source line, +next+/+redo+ silently end or restart the thunk body. Weigh that
+      # when choosing what your surface thunks.
+      #
+      # Prefer +run_after+ when the execution point sits in the same statement sequence as the statements.
+      #
+      # @param statements [Array<Parser::AST::Node>] statements to wrap
+      # @return [ASTTransform::Thunk] the thunk node
+      def thunk(*statements)
+        s(:ast_thunk, Thunk::Id.new, *statements)
+      end
+
+      # The paved road for execution reordering in flat statement sequences. Named for the constraint, not the
+      # mechanism — with text pinned to source lines the only physical lever is delaying execution, so "run X after
+      # Y" is the constraint an author states. Returns a NEW sequence in which the +run+ statements are removed and a
+      # thunk wrapping them is inserted immediately after +after+.
+      #
+      # All membership checks are by identity (equal?), never ==: node equality ignores location, so two textually
+      # identical statements on different lines compare == and value matching could splice the wrong one.
+      #
+      # @param statements [Array<Parser::AST::Node>] the sequence being composed
+      # @param run [Array<Parser::AST::Node>] contiguous run of elements of +statements+ (by identity) whose
+      #   execution must wait
+      # @param after [Parser::AST::Node] element of +statements+ (by identity, not inside +run+) the +run+ statements
+      #   execute after
+      # @return [Array<Parser::AST::Node>] new sequence with the thunk placed
+      # @raise [ArgumentError] if +run+ is not a contiguous identity-run of +statements+, or +after+ is not an
+      #   element (or is inside +run+)
+      def run_after(statements, run:, after:)
+        run_range = contiguous_identity_range(statements, run)
+        raise ArgumentError, "run: must be a contiguous run of elements of statements (by identity)" unless run_range
+
+        after_index = statements.index { |statement| statement.equal?(after) }
+        raise ArgumentError, "after: must be an element of statements (by identity)" unless after_index
+        raise ArgumentError, "after: cannot be inside run:" if run_range.cover?(after_index)
+
+        reordered = statements.dup
+        reordered[run_range] = []
+
+        insertion_index = reordered.index { |statement| statement.equal?(after) }
+        reordered.insert(insertion_index + 1, thunk(*run))
+      end
+
+      private
+
+      # The range +members+ occupies in +sequence+, or nil unless members is a non-empty contiguous identity-run in
+      # order.
+      def contiguous_identity_range(sequence, members)
+        return nil if members.empty?
+
+        start = sequence.index { |element| element.equal?(members.first) }
+        return nil unless start
+
+        contiguous = members.each_with_index.all? do |member, offset|
+          sequence[start + offset]&.equal?(member)
+        end
+
+        contiguous ? (start...(start + members.size)) : nil
       end
     end
   end

data/lib/ast_transform/transformer.rb

@@ -1,17 +1,20 @@
 # frozen_string_literal: true
+
 require 'prism'
 require 'prism/translation/parser'
 require 'unparser'
 require 'ast_transform/kwargs_builder'
-require 'ast_transform/source_map'
+require 'ast_transform/line_aligned_emitter'
 
 module ASTTransform
   class Transformer
     # Constructs a new Transformer instance.
     #
     # @param transformations [Array<ASTTransform::AbstractTransformation>] The transformations to be run.
-    def initialize(*transformations)
+    # @param emitter [ASTTransform::LineAlignedEmitter] The emitter rendering transformed ASTs back to source.
+    def initialize(*transformations, emitter: LineAlignedEmitter.new)
       @transformations = transformations
+      @emitter = emitter
     end
 
     # Builds the AST for the given +source+.
@@ -20,7 +23,7 @@ module ASTTransform
     # @param file_path [String] The file_path. This is important for source mapping in backtraces.
     #
     # @return [Parser::AST::Node] The AST.
-    def build_ast(source, file_path: 'tmp')
+    def build_ast(source, file_path: "tmp")
       buffer = create_buffer(source, file_path)
       parser.parse(buffer)
     end
@@ -39,16 +42,17 @@ module ASTTransform
     #
     # @param source [String] The input source code to be transformed.
     #
-    # @return [String] The transformed code.
+    # @return [String] The transformed code, line-aligned (see #transform_file_source).
     def transform(source)
       ast = build_ast(source)
       transformed_ast = transform_ast(ast)
-      Unparser.unparse(transformed_ast)
+      @emitter.emit(transformed_ast, 'tmp')
     end
 
     # Transforms the give +file_path+.
     #
-    # @param file_path [String] The input file to be transformed. This is required for source mapping in backtraces.
+    # @param file_path [String] The input file to be transformed. Statement placement (and therefore
+    # backtrace and breakpoint line numbers) is derived from this file's source locations.
     # @param transformed_file_path [String] The file path to the transformed file.
     #
     # @return [String] The transformed code.
@@ -60,21 +64,19 @@ module ASTTransform
     # Transforms the given +source+ in +file_path+.
     #
     # @param source [String] The input source code to be transformed.
-    # @param file_path [String] The file path for the input +source+. This is required for source mapping in backtraces.
-    # @param transformed_file_path [String] The file path to the transformed filed. This is required to register the
-    # SourceMap.
+    # @param file_path [String] The file path for the input +source+. Statement placement (and
+    # therefore backtrace and breakpoint line numbers) is derived from the source locations parsed
+    # under this path.
+    # @param transformed_file_path [String] The file path the transformed file will be written to.
     #
-    # @return [String] The transformed code.
-    def transform_file_source(source, file_path, transformed_file_path)
+    # @return [String] The transformed code, line-aligned: every statement carrying a source
+    # location is emitted at its original source line.
+    def transform_file_source(source, file_path, _transformed_file_path)
       source_ast = build_ast(source, file_path: file_path)
-      # At this point, the transformed_ast contains line number mappings for the original +source+.
+      # At this point, the transformed_ast contains source locations for the original +source+.
       transformed_ast = transform_ast(source_ast)
 
-      transformed_source = Unparser.unparse(transformed_ast)
-
-      register_source_map(file_path, transformed_file_path, transformed_ast, transformed_source)
-
-      transformed_source
+      @emitter.emit(transformed_ast, file_path)
     end
 
     # Transforms the given +ast+.
@@ -101,13 +103,5 @@ module ASTTransform
       @parser&.reset
       @parser ||= Prism::Translation::Parser.new(ASTTransform::KwargsBuilder.new)
     end
-
-    def register_source_map(source_file_path, transformed_file_path, transformed_ast, transformed_source)
-      # The transformed_source is re-parsed to get the correct line numbers for the transformed_ast, which is the code
-      # that will run.
-      rewritten_ast = build_ast(transformed_source)
-      source_map = ASTTransform::SourceMap.new(source_file_path, transformed_file_path, transformed_ast, rewritten_ast)
-      ASTTransform::SourceMap.register_source_map(source_map)
-    end
   end
 end

data/lib/ast_transform/version.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module ASTTransform
-  VERSION = "2.1.2"
+  VERSION = "3.0.0"
 end

data/lib/ast_transform.rb

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 require "ast_transform/version"
-require 'ast_transform/instruction_sequence'
-require 'ast_transform/instruction_sequence/mixin'
-require 'ast_transform/instruction_sequence/bootsnap_mixin'
+require "ast_transform/instruction_sequence"
+require "ast_transform/instruction_sequence/mixin"
+require "ast_transform/instruction_sequence/bootsnap_mixin"
 
 module ASTTransform
   DEFAULT_OUTPUT_PATH = Pathname.new("").join("tmp", "ast_transform").to_s
@@ -19,22 +19,18 @@ module ASTTransform
     end
 
     def install
-      @installed ||= begin
-        if defined?(Bootsnap) && ASTTransform::InstructionSequence.using_bootsnap_compilation?
-          class << Bootsnap::CompileCache::ISeq
-            prepend ::ASTTransform::InstructionSequence::BootsnapMixin
-          end
-        else
-          class << RubyVM::InstructionSequence
-            prepend ::ASTTransform::InstructionSequence::Mixin
-          end
+      @installed ||= if defined?(Bootsnap) && ASTTransform::InstructionSequence.using_bootsnap_compilation?
+        class << Bootsnap::CompileCache::ISeq
+          prepend ::ASTTransform::InstructionSequence::BootsnapMixin
+        end
+      else
+        class << RubyVM::InstructionSequence
+          prepend ::ASTTransform::InstructionSequence::Mixin
         end
       end
     end
 
-    def output_path=(path)
-      @output_path = path
-    end
+    attr_writer :output_path
 
     def output_path
       @output_path || DEFAULT_OUTPUT_PATH