[ruby/yarp] Provide a desugar visitor

https://github.com/ruby/yarp/commit/9fad513089

lib/yarp.rb

@@ -52,6 +52,16 @@ module YARP
       @length = length
     end
 
+    # Create a new location object with the given options.
+    def copy(**options)
+      Location.new(
+        options.fetch(:source) { source },
+        options.fetch(:start_offset) { start_offset },
+        options.fetch(:length) { length }
+      )
+    end
+
+    # Returns a string representation of this location.
     def inspect
       "#<YARP::Location @start_offset=#{@start_offset} @length=#{@length}>"
     end
@@ -508,6 +518,7 @@ end
 
 require_relative "yarp/lex_compat"
 require_relative "yarp/mutation_visitor"
+require_relative "yarp/desugar_visitor"
 require_relative "yarp/node"
 require_relative "yarp/ripper_compat"
 require_relative "yarp/serialize"

lib/yarp/desugar_visitor.rb

@@ -0,0 +1,267 @@
+# frozen_string_literal: true
+
+module YARP
+  class DesugarVisitor < MutationVisitor
+    # @@foo &&= bar
+    #
+    # becomes
+    #
+    # @@foo && @@foo = bar
+    def visit_class_variable_and_write_node(node)
+      AndNode.new(
+        ClassVariableReadNode.new(node.name_loc),
+        ClassVariableWriteNode.new(node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # @@foo ||= bar
+    #
+    # becomes
+    #
+    # @@foo || @@foo = bar
+    def visit_class_variable_or_write_node(node)
+      OrNode.new(
+        ClassVariableReadNode.new(node.name_loc),
+        ClassVariableWriteNode.new(node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # @@foo += bar
+    #
+    # becomes
+    #
+    # @@foo = @@foo + bar
+    def visit_class_variable_operator_write_node(node)
+      desugar_operator_write_node(node, ClassVariableWriteNode, ClassVariableReadNode)
+    end
+
+    # Foo &&= bar
+    #
+    # becomes
+    #
+    # Foo && Foo = bar
+    def visit_constant_and_write_node(node)
+      AndNode.new(
+        ConstantReadNode.new(node.name_loc),
+        ConstantWriteNode.new(node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # Foo ||= bar
+    #
+    # becomes
+    #
+    # Foo || Foo = bar
+    def visit_constant_or_write_node(node)
+      OrNode.new(
+        ConstantReadNode.new(node.name_loc),
+        ConstantWriteNode.new(node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # Foo += bar
+    #
+    # becomes
+    #
+    # Foo = Foo + bar
+    def visit_constant_operator_write_node(node)
+      desugar_operator_write_node(node, ConstantWriteNode, ConstantReadNode)
+    end
+
+    # Foo::Bar &&= baz
+    #
+    # becomes
+    #
+    # Foo::Bar && Foo::Bar = baz
+    def visit_constant_path_and_write_node(node)
+      AndNode.new(
+        node.target,
+        ConstantPathWriteNode.new(node.target, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # Foo::Bar ||= baz
+    #
+    # becomes
+    #
+    # Foo::Bar || Foo::Bar = baz
+    def visit_constant_path_or_write_node(node)
+      OrNode.new(
+        node.target,
+        ConstantPathWriteNode.new(node.target, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # Foo::Bar += baz
+    #
+    # becomes
+    #
+    # Foo::Bar = Foo::Bar + baz
+    def visit_constant_path_operator_write_node(node)
+      ConstantPathWriteNode.new(
+        node.target,
+        CallNode.new(
+          node.target,
+          nil,
+          node.operator_loc.copy(length: node.operator_loc.length - 1),
+          nil,
+          ArgumentsNode.new([node.value], node.value.location),
+          nil,
+          nil,
+          0,
+          node.operator_loc.slice.chomp("="),
+          node.location
+        ),
+        node.operator_loc.copy(start_offset: node.operator_loc.end_offset - 1, length: 1),
+        node.location
+      )
+    end
+
+    # $foo &&= bar
+    #
+    # becomes
+    #
+    # $foo && $foo = bar
+    def visit_global_variable_and_write_node(node)
+      AndNode.new(
+        GlobalVariableReadNode.new(node.name_loc),
+        GlobalVariableWriteNode.new(node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # $foo ||= bar
+    #
+    # becomes
+    #
+    # $foo || $foo = bar
+    def visit_global_variable_or_write_node(node)
+      OrNode.new(
+        GlobalVariableReadNode.new(node.name_loc),
+        GlobalVariableWriteNode.new(node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # $foo += bar
+    #
+    # becomes
+    #
+    # $foo = $foo + bar
+    def visit_global_variable_operator_write_node(node)
+      desugar_operator_write_node(node, GlobalVariableWriteNode, GlobalVariableReadNode)
+    end
+
+    # @foo &&= bar
+    #
+    # becomes
+    #
+    # @foo && @foo = bar
+    def visit_instance_variable_and_write_node(node)
+      AndNode.new(
+        InstanceVariableReadNode.new(node.name_loc),
+        InstanceVariableWriteNode.new(node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # @foo ||= bar
+    #
+    # becomes
+    #
+    # @foo || @foo = bar
+    def visit_instance_variable_or_write_node(node)
+      OrNode.new(
+        InstanceVariableReadNode.new(node.name_loc),
+        InstanceVariableWriteNode.new(node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # @foo += bar
+    #
+    # becomes
+    #
+    # @foo = @foo + bar
+    def visit_instance_variable_operator_write_node(node)
+      desugar_operator_write_node(node, InstanceVariableWriteNode, InstanceVariableReadNode)
+    end
+
+    # foo &&= bar
+    #
+    # becomes
+    #
+    # foo && foo = bar
+    def visit_local_variable_and_write_node(node)
+      AndNode.new(
+        LocalVariableReadNode.new(node.constant_id, node.depth, node.name_loc),
+        LocalVariableWriteNode.new(node.constant_id, node.depth, node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # foo ||= bar
+    #
+    # becomes
+    #
+    # foo || foo = bar
+    def visit_local_variable_or_write_node(node)
+      OrNode.new(
+        LocalVariableReadNode.new(node.constant_id, node.depth, node.name_loc),
+        LocalVariableWriteNode.new(node.constant_id, node.depth, node.name_loc, node.value, node.operator_loc, node.location),
+        node.operator_loc,
+        node.location
+      )
+    end
+
+    # foo += bar
+    #
+    # becomes
+    #
+    # foo = foo + bar
+    def visit_local_variable_operator_write_node(node)
+      desugar_operator_write_node(node, LocalVariableWriteNode, LocalVariableReadNode, arguments: [node.constant_id, node.depth])
+    end
+
+    private
+
+    # Desugar `x += y` to `x = x + y`
+    def desugar_operator_write_node(node, write_class, read_class, arguments: [])
+      write_class.new(
+        *arguments,
+        node.name_loc,
+        CallNode.new(
+          read_class.new(*arguments, node.name_loc),
+          nil,
+          node.operator_loc.copy(length: node.operator_loc.length - 1),
+          nil,
+          ArgumentsNode.new([node.value], node.value.location),
+          nil,
+          nil,
+          0,
+          node.operator_loc.slice.chomp("="),
+          node.location
+        ),
+        node.operator_loc.copy(start_offset: node.operator_loc.end_offset - 1, length: 1),
+        node.location
+      )
+    end
+  end
+end

lib/yarp/yarp.gemspec

@@ -59,6 +59,7 @@ Gem::Specification.new do |spec|
     "include/yarp/util/yp_strpbrk.h",
     "include/yarp/version.h",
     "lib/yarp.rb",
+    "lib/yarp/desugar_visitor.rb",
     "lib/yarp/ffi.rb",
     "lib/yarp/lex_compat.rb",
     "lib/yarp/mutation_visitor.rb",

test/desugar_visitor_test.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require "yarp_test_helper"
+
+class DesugarVisitorTest < Test::Unit::TestCase
+  def test_and_write
+    assert_desugars("(AndNode (ClassVariableReadNode) (ClassVariableWriteNode (CallNode)))", "@@foo &&= bar")
+    assert_desugars("(AndNode (ConstantPathNode (ConstantReadNode) (ConstantReadNode)) (ConstantPathWriteNode (ConstantPathNode (ConstantReadNode) (ConstantReadNode)) (CallNode)))", "Foo::Bar &&= baz")
+    assert_desugars("(AndNode (ConstantReadNode) (ConstantWriteNode (CallNode)))", "Foo &&= bar")
+    assert_desugars("(AndNode (GlobalVariableReadNode) (GlobalVariableWriteNode (CallNode)))", "$foo &&= bar")
+    assert_desugars("(AndNode (InstanceVariableReadNode) (InstanceVariableWriteNode (CallNode)))", "@foo &&= bar")
+    assert_desugars("(AndNode (LocalVariableReadNode) (LocalVariableWriteNode (CallNode)))", "foo &&= bar")
+    assert_desugars("(AndNode (LocalVariableReadNode) (LocalVariableWriteNode (CallNode)))", "foo = 1; foo &&= bar")
+  end
+
+  def test_or_write
+    assert_desugars("(OrNode (ClassVariableReadNode) (ClassVariableWriteNode (CallNode)))", "@@foo ||= bar")
+    assert_desugars("(OrNode (ConstantPathNode (ConstantReadNode) (ConstantReadNode)) (ConstantPathWriteNode (ConstantPathNode (ConstantReadNode) (ConstantReadNode)) (CallNode)))", "Foo::Bar ||= baz")
+    assert_desugars("(OrNode (ConstantReadNode) (ConstantWriteNode (CallNode)))", "Foo ||= bar")
+    assert_desugars("(OrNode (GlobalVariableReadNode) (GlobalVariableWriteNode (CallNode)))", "$foo ||= bar")
+    assert_desugars("(OrNode (InstanceVariableReadNode) (InstanceVariableWriteNode (CallNode)))", "@foo ||= bar")
+    assert_desugars("(OrNode (LocalVariableReadNode) (LocalVariableWriteNode (CallNode)))", "foo ||= bar")
+    assert_desugars("(OrNode (LocalVariableReadNode) (LocalVariableWriteNode (CallNode)))", "foo = 1; foo ||= bar")
+  end
+
+  def test_operator_write
+    assert_desugars("(ClassVariableWriteNode (CallNode (ClassVariableReadNode) (ArgumentsNode (CallNode))))", "@@foo += bar")
+    assert_desugars("(ConstantPathWriteNode (ConstantPathNode (ConstantReadNode) (ConstantReadNode)) (CallNode (ConstantPathNode (ConstantReadNode) (ConstantReadNode)) (ArgumentsNode (CallNode))))", "Foo::Bar += baz")
+    assert_desugars("(ConstantWriteNode (CallNode (ConstantReadNode) (ArgumentsNode (CallNode))))", "Foo += bar")
+    assert_desugars("(GlobalVariableWriteNode (CallNode (GlobalVariableReadNode) (ArgumentsNode (CallNode))))", "$foo += bar")
+    assert_desugars("(InstanceVariableWriteNode (CallNode (InstanceVariableReadNode) (ArgumentsNode (CallNode))))", "@foo += bar")
+    assert_desugars("(LocalVariableWriteNode (CallNode (LocalVariableReadNode) (ArgumentsNode (CallNode))))", "foo += bar")
+    assert_desugars("(LocalVariableWriteNode (CallNode (LocalVariableReadNode) (ArgumentsNode (CallNode))))", "foo = 1; foo += bar")
+  end
+
+  private
+
+  def ast_inspect(node)
+    parts = [node.class.name.split("::").last]
+
+    node.deconstruct_keys(nil).each do |_, value|
+      case value
+      when YARP::Node
+        parts << ast_inspect(value)
+      when Array
+        parts.concat(value.map { |element| ast_inspect(element) })
+      end
+    end
+
+    "(#{parts.join(" ")})"
+  end
+
+  def assert_desugars(expected, source)
+    ast = YARP.parse(source).value.accept(YARP::DesugarVisitor.new)
+    assert_equal expected, ast_inspect(ast.statements.body.last)
+  end
+end

yarp/templates/lib/yarp/mutation_visitor.rb.erb

@@ -9,7 +9,7 @@ module YARP
     def visit_<%= node.human %>(node)
       <%- params = node.params.select { |param| [NodeParam, OptionalNodeParam, NodeListParam].include?(param.class) } -%>
       <%- if params.any? -%>
-      node.copy(<%= params.map { |param| "#{param.name}: visit(node.#{param.name})" }.join(", ") %>)
+      node.copy(<%= params.map { |param| "#{param.name}: #{param.is_a?(NodeListParam) ? "visit_all" : "visit"}(node.#{param.name})" }.join(", ") %>)
       <%- else -%>
       node.copy
       <%- end -%>

yarp/templates/lib/yarp/node.rb.erb

@@ -52,7 +52,7 @@ module YARP
     def copy(**params)
       <%= node.name %>.new(
         <%- (node.params.map(&:name) + ["location"]).map do |name| -%>
-        <%= name %>: params.fetch(:<%= name %>) { self.<%= name %> },
+        params.fetch(:<%= name %>) { <%= name %> },
         <%- end -%>
       )
     end