[ruby/yarp] Split up compiler versus visitor

https://github.com/ruby/yarp/commit/2e6baa3f19

10 files changed, 154 insertions, 36 deletions

diff --git a/lib/yarp.rb b/lib/yarp.rb
index 9b9df081d8..4908e57b05 100644
--- a/lib/yarp.rb
+++ b/lib/yarp.rb
@@ -229,24 +229,6 @@ module YARP
     end
   end
 
-  # A class that knows how to walk down the tree. None of the individual visit
-  # methods are implemented on this visitor, so it forces the consumer to
-  # implement each one that they need. For a default implementation that
-  # continues walking the tree, see the Visitor class.
-  class BasicVisitor
-    def visit(node)
-      node&.accept(self)
-    end
-
-    def visit_all(nodes)
-      nodes.map { |node| visit(node) }
-    end
-
-    def visit_child_nodes(node)
-      visit_all(node.child_nodes)
-    end
-  end
-
   # This represents a token from the Ruby source.
   class Token
     attr_reader :type, :value, :location
@@ -539,14 +521,17 @@ module YARP
   # which means the files can end up being quite large. We autoload them to make
   # our require speed faster since consuming libraries are unlikely to use all
   # of these features.
-  autoload :DesugarVisitor, "yarp/desugar_visitor"
+  autoload :BasicVisitor, "yarp/visitor"
+  autoload :Compiler, "yarp/compiler"
+  autoload :DesugarCompiler, "yarp/desugar_compiler"
   autoload :Dispatcher, "yarp/dispatcher"
   autoload :DSL, "yarp/dsl"
-  autoload :MutationVisitor, "yarp/mutation_visitor"
+  autoload :MutationCompiler, "yarp/mutation_compiler"
   autoload :RipperCompat, "yarp/ripper_compat"
   autoload :Pack, "yarp/pack"
   autoload :Pattern, "yarp/pattern"
   autoload :Serialize, "yarp/serialize"
+  autoload :Visitor, "yarp/visitor"
 
   # Load the serialized AST using the source as a reference into a tree.
   def self.load(source, serialized)
diff --git a/lib/yarp/desugar_visitor.rb b/lib/yarp/desugar_compiler.rb
index 6ee5861ac8..2cf5fb701a 100644
--- a/lib/yarp/desugar_visitor.rb
+++ b/lib/yarp/desugar_compiler.rb
@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module YARP
-  class DesugarVisitor < MutationVisitor
+  class DesugarCompiler < MutationCompiler
     # @@foo &&= bar
     #
     # becomes
diff --git a/lib/yarp/yarp.gemspec b/lib/yarp/yarp.gemspec
index 16e8196140..fea48dd47f 100644
--- a/lib/yarp/yarp.gemspec
+++ b/lib/yarp/yarp.gemspec
@@ -59,12 +59,13 @@ 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/compiler.rb",
+    "lib/yarp/desugar_compiler.rb",
     "lib/yarp/dispatcher.rb",
     "lib/yarp/dsl.rb",
     "lib/yarp/ffi.rb",
     "lib/yarp/lex_compat.rb",
-    "lib/yarp/mutation_visitor.rb",
+    "lib/yarp/mutation_compiler.rb",
     "lib/yarp/node.rb",
     "lib/yarp/pack.rb",
     "lib/yarp/pattern.rb",
@@ -72,6 +73,7 @@ Gem::Specification.new do |spec|
     "lib/yarp/serialize.rb",
     "lib/yarp/parse_result/comments.rb",
     "lib/yarp/parse_result/newlines.rb",
+    "lib/yarp/visitor.rb",
     "src/diagnostic.c",
     "src/enc/yp_big5.c",
     "src/enc/yp_euc_jp.c",
diff --git a/test/yarp/compiler_test.rb b/test/yarp/compiler_test.rb
new file mode 100644
index 0000000000..141e183469
--- /dev/null
+++ b/test/yarp/compiler_test.rb
@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require_relative "test_helper"
+
+module YARP
+  class CompilerTest < TestCase
+    class SExpressions < YARP::Compiler
+      def visit_arguments_node(node)
+        [:arguments, super]
+      end
+
+      def visit_call_node(node)
+        [:call, super]
+      end
+
+      def visit_integer_node(node)
+        [:integer]
+      end
+
+      def visit_program_node(node)
+        [:program, super]
+      end
+    end
+
+    def test_compiler
+      expected = [:program, [[[:call, [[:integer], [:arguments, [[:integer]]]]]]]]
+      assert_equal expected, YARP.parse("1 + 2").value.accept(SExpressions.new)
+    end
+  end
+end
diff --git a/test/yarp/desugar_visitor_test.rb b/test/yarp/desugar_compiler_test.rb
index 3966d7bfcb..8d2b207fed 100644
--- a/test/yarp/desugar_visitor_test.rb
+++ b/test/yarp/desugar_compiler_test.rb
@@ -3,7 +3,7 @@
 require_relative "test_helper"
 
 module YARP
-  class DesugarVisitorTest < TestCase
+  class DesugarCompilerTest < TestCase
     def test_and_write
       assert_desugars("(AndNode (ClassVariableReadNode) (ClassVariableWriteNode (CallNode)))", "@@foo &&= bar")
       assert_not_desugared("Foo::Bar &&= baz", "Desugaring would execute Foo twice or need temporary variables")
@@ -72,7 +72,7 @@ module YARP
     end
 
     def assert_desugars(expected, source)
-      ast = YARP.parse(source).value.accept(DesugarVisitor.new)
+      ast = YARP.parse(source).value.accept(DesugarCompiler.new)
       assert_equal expected, ast_inspect(ast.statements.body.last)
 
       ast.accept(EnsureEveryNodeOnceInAST.new)
@@ -80,7 +80,7 @@ module YARP
 
     def assert_not_desugared(source, reason)
       ast = YARP.parse(source).value
-      assert_equal_nodes(ast, ast.accept(DesugarVisitor.new))
+      assert_equal_nodes(ast, ast.accept(DesugarCompiler.new))
     end
   end
 end
diff --git a/yarp/templates/lib/yarp/compiler.rb.erb b/yarp/templates/lib/yarp/compiler.rb.erb
new file mode 100644
index 0000000000..b9ddb4daee
--- /dev/null
+++ b/yarp/templates/lib/yarp/compiler.rb.erb
@@ -0,0 +1,41 @@
+module YARP
+  # A compiler is a visitor that returns the value of each node as it visits.
+  # This is as opposed to a visitor which will only walk the tree. This can be
+  # useful when you are trying to compile a tree into a different format.
+  #
+  # For example, to build a representation of the tree as s-expressions, you
+  # could write:
+  #
+  #     class SExpressions < YARP::Compiler
+  #       def visit_arguments_node(node) = [:arguments, super]
+  #       def visit_call_node(node) = [:call, super]
+  #       def visit_integer_node(node) = [:integer]
+  #       def visit_program_node(node) = [:program, super]
+  #     end
+  #
+  #     YARP.parse("1 + 2").value.accept(SExpressions.new)
+  #     # => [:program, [[[:call, [[:integer], [:arguments, [[:integer]]]]]]]]
+  #
+  class Compiler
+    # Visit an individual node.
+    def visit(node)
+      node&.accept(self)
+    end
+
+    # Visit a list of nodes.
+    def visit_all(nodes)
+      nodes.map { |node| node&.accept(self) }
+    end
+
+    # Visit the child nodes of the given node.
+    def visit_child_nodes(node)
+      node.compact_child_nodes.map { |node| node.accept(self) }
+    end
+
+    <%- nodes.each_with_index do |node, index| -%>
+    <%= "\n" if index != 0 -%>
+    # Compile a <%= node.name %> node
+    alias visit_<%= node.human %> visit_child_nodes
+    <%- end -%>
+  end
+end
diff --git a/yarp/templates/lib/yarp/mutation_compiler.rb.erb b/yarp/templates/lib/yarp/mutation_compiler.rb.erb
new file mode 100644
index 0000000000..a99721f53b
--- /dev/null
+++ b/yarp/templates/lib/yarp/mutation_compiler.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 MutationCompiler < Compiler
+    <%- nodes.each_with_index do |node, index| -%>
+<%= "\n" if index != 0 -%>
+    # Copy a <%= node.name %> node
+    def visit_<%= node.human %>(node)
+      <%- fields = node.fields.select { |field| [YARP::NodeField, YARP::OptionalNodeField, YARP::NodeListField].include?(field.class) } -%>
+      <%- if fields.any? -%>
+      node.copy(<%= fields.map { |field| "#{field.name}: #{field.is_a?(YARP::NodeListField) ? "visit_all" : "visit"}(node.#{field.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 7d52d823da..16717de739 100644
--- a/yarp/templates/lib/yarp/node.rb.erb
+++ b/yarp/templates/lib/yarp/node.rb.erb
@@ -164,7 +164,8 @@ module YARP
   end
 
   <%- end -%>
-  <%- flags.each do |flag| -%>
+  <%- flags.each_with_index do |flag, flag_index| -%>
+<%= "\n" if flag_index > 0 -%>
   module <%= flag.name %>
     <%- flag.values.each_with_index do |value, index| -%>
     # <%= value.comment %>
@@ -172,13 +173,5 @@ module YARP
 <%= "\n" if value != flag.values.last -%>
     <%- end -%>
   end
-
   <%- end -%>
-  class Visitor < BasicVisitor
-    <%- nodes.each do |node| -%>
-    # Visit a <%= node.name %> node
-    alias visit_<%= node.human %> visit_child_nodes
-<%= "\n" if node != nodes.last -%>
-    <%- end -%>
-  end
 end
diff --git a/yarp/templates/lib/yarp/visitor.rb.erb b/yarp/templates/lib/yarp/visitor.rb.erb
new file mode 100644
index 0000000000..0bef47b997
--- /dev/null
+++ b/yarp/templates/lib/yarp/visitor.rb.erb
@@ -0,0 +1,46 @@
+module YARP
+  # A class that knows how to walk down the tree. None of the individual visit
+  # methods are implemented on this visitor, so it forces the consumer to
+  # implement each one that they need. For a default implementation that
+  # continues walking the tree, see the Visitor class.
+  class BasicVisitor
+    def visit(node)
+      node&.accept(self)
+    end
+
+    def visit_all(nodes)
+      nodes.each { |node| node&.accept(self) }
+    end
+
+    def visit_child_nodes(node)
+      node.compact_child_nodes.each { |node| node.accept(self) }
+    end
+  end
+
+  # A visitor is a class that provides a default implementation for every accept
+  # method defined on the nodes. This means it can walk a tree without the
+  # caller needing to define any special handling. This allows you to handle a
+  # subset of the tree, while still walking the whole tree.
+  #
+  # For example, to find all of the method calls that call the `foo` method, you
+  # could write:
+  #
+  #     class FooCalls < YARP::Visitor
+  #       def visit_call_node(node)
+  #         if node.name == "foo"
+  #           # Do something with the node
+  #         end
+  #
+  #         # Call super so that the visitor continues walking the tree
+  #         super
+  #       end
+  #     end
+  #
+  class Visitor < BasicVisitor
+    <%- nodes.each_with_index do |node, index| -%>
+<%= "\n" if index != 0 -%>
+    # Visit a <%= node.name %> node
+    alias visit_<%= node.human %> visit_child_nodes
+    <%- end -%>
+  end
+end
diff --git a/yarp/templates/template.rb b/yarp/templates/template.rb
index ea5946761e..f24c756438 100755
--- a/yarp/templates/template.rb
+++ b/yarp/templates/template.rb
@@ -366,11 +366,13 @@ module YARP
     "java/org/yarp/Loader.java",
     "java/org/yarp/Nodes.java",
     "java/org/yarp/AbstractNodeVisitor.java",
+    "lib/yarp/compiler.rb",
     "lib/yarp/dispatcher.rb",
     "lib/yarp/dsl.rb",
-    "lib/yarp/mutation_visitor.rb",
+    "lib/yarp/mutation_compiler.rb",
     "lib/yarp/node.rb",
     "lib/yarp/serialize.rb",
+    "lib/yarp/visitor.rb",
     "src/node.c",
     "src/prettyprint.c",
     "src/serialize.c",