From 649aba28f43534c729087238d01362a1f01bcc4a Mon Sep 17 00:00:00 2001
From: Kevin Newton <kddnewton@gmail.com>
Date: Fri, 25 Aug 2023 09:57:39 -0400
Subject: [PATCH] [ruby/yarp] Add Node#copy and MutationVisitor

https://github.com/ruby/yarp/commit/3693091661
---
 lib/yarp.rb                                   |  1 +
 lib/yarp/yarp.gemspec                         |  1 +
 .../lib/yarp/mutation_visitor.rb.erb          | 19 +++++++++++++++++++
 yarp/templates/lib/yarp/node.rb.erb           |  9 +++++++++
 yarp/templates/template.rb                    |  1 +
 5 files changed, 31 insertions(+)
 create mode 100644 yarp/templates/lib/yarp/mutation_visitor.rb.erb

diff --git a/lib/yarp.rb b/lib/yarp.rb
index 15217e80f490a7..a3b0c3b0742fa5 100644
--- a/lib/yarp.rb
+++ b/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"
diff --git a/lib/yarp/yarp.gemspec b/lib/yarp/yarp.gemspec
index db69bcb451b8d7..70b6c7ebaa374f 100644
--- a/lib/yarp/yarp.gemspec
+++ b/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",
diff --git a/yarp/templates/lib/yarp/mutation_visitor.rb.erb b/yarp/templates/lib/yarp/mutation_visitor.rb.erb
new file mode 100644
index 00000000000000..c88cabbdbb030c
--- /dev/null
+++ b/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
diff --git a/yarp/templates/lib/yarp/node.rb.erb b/yarp/templates/lib/yarp/node.rb.erb
index c1f14f537af9ec..9bd9eb54e1d6ac 100644
--- a/yarp/templates/lib/yarp/node.rb.erb
+++ b/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
 
diff --git a/yarp/templates/template.rb b/yarp/templates/template.rb
index 8b8a175172ffea..e34d5b1a7be6ea 100755
--- a/yarp/templates/template.rb
+++ b/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",