[ruby/yarp] Add Node#copy and MutationVisitor

ruby/prism@3693091661

lib/yarp.rb

@@ -507,6 +507,7 @@ def self.parse_serialize_file(filepath)
 end
 
 require_relative "yarp/lex_compat"
+require_relative "yarp/mutation_visitor"
 require_relative "yarp/node"
 require_relative "yarp/ripper_compat"
 require_relative "yarp/serialize"

lib/yarp/yarp.gemspec

@@ -61,6 +61,7 @@ Gem::Specification.new do |spec|
     "lib/yarp.rb",
     "lib/yarp/ffi.rb",
     "lib/yarp/lex_compat.rb",
+    "lib/yarp/mutation_visitor.rb",
     "lib/yarp/node.rb",
     "lib/yarp/pack.rb",
     "lib/yarp/ripper_compat.rb",

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

@@ -0,0 +1,19 @@
+module YARP
+  # This visitor walks through the tree and copies each node as it is being
+  # visited. This is useful for consumers that want to mutate the tree, as you
+  # can change subtrees in place without effecting the rest of the tree.
+  class MutationVisitor < BasicVisitor
+    <%- nodes.each_with_index do |node, index| -%>
+<%= "\n" if index != 0 -%>
+    # Copy a <%= node.name %> node
+    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(", ") %>)
+      <%- else -%>
+      node.copy
+      <%- end -%>
+    end
+    <%- end -%>
+  end
+end

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

@@ -48,6 +48,15 @@ module YARP
       }.compact.join(", ") %>]
     end
 
+    # def copy: (**params) -> <%= node.name %>
+    def copy(**params)
+      <%= node.name %>.new(
+        <%- (node.params.map(&:name) + ["location"]).map do |name| -%>
+        <%= name %>: params.fetch(:<%= name %>) { self.<%= name %> },
+        <%- end -%>
+      )
+    end
+
     # def deconstruct: () -> Array[nil | Node]
     alias deconstruct child_nodes
 

yarp/templates/template.rb

@@ -315,6 +315,7 @@ def locals
   "java/org/yarp/Loader.java",
   "java/org/yarp/Nodes.java",
   "java/org/yarp/AbstractNodeVisitor.java",
+  "lib/yarp/mutation_visitor.rb",
   "lib/yarp/node.rb",
   "lib/yarp/serialize.rb",
   "src/node.c",