Implement visitor pattern for SyntaxTree::Node

lib/syntax_tree/node.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "syntax_tree/visitor"
+
 module SyntaxTree
   # Represents the location of a node in the tree from the source code.
   class Location
@@ -56,6 +58,29 @@ class Node
     # [Location] the location of this node
     attr_reader :location
 
+    def self.inherited(child)
+      child.class_eval(<<~EOS, __FILE__, __LINE__ + 1)
+        def accept(visitor)
+          visitor.#{child.visit_method_name}(self)
+        end
+      EOS
+
+      Visitor.class_eval(<<~EOS, __FILE__, __LINE__ + 1)
+        def #{child.visit_method_name}(node)
+          visit_all(node.child_nodes)
+        end
+      EOS
+    end
+
+    def self.visit_method_name
+      method_suffix = name.split("::").last
+      method_suffix.gsub!(/([A-Z\d]+)([A-Z][a-z])/, '\1_\2')
+      method_suffix.gsub!(/([a-z\d])([A-Z])/, '\1_\2')
+      method_suffix.tr!("-", "_")
+      method_suffix.downcase!
+      "visit_#{method_suffix}"
+    end
+
     def child_nodes
       raise NotImplementedError
     end

lib/syntax_tree/visitor.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module SyntaxTree
+  class Visitor
+    def visit_all(nodes)
+      nodes.each do |node|
+        visit(node)
+      end
+    end
+
+    def visit(node)
+      node&.accept(self)
+    end
+  end
+end

test/visitor_test.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+require_relative "test_helper"
+require "objspace"
+
+class VisitorTest < Minitest::Test
+  def test_can_visit_all_nodes
+    visitor = SyntaxTree::Visitor.new
+
+    ObjectSpace.each_object(SyntaxTree::Node.singleton_class)
+      .reject { |node| node.singleton_class? || node == SyntaxTree::Node }
+      .each { |node| assert_respond_to(visitor, node.visit_method_name) }
+  end
+
+  def test_node_visit_method_name
+    assert_equal("visit_t_string_end", SyntaxTree::TStringEnd.visit_method_name)
+  end
+
+  def test_visit_tree
+    parsed_tree = SyntaxTree.parse(<<~RUBY)
+      class Foo
+        def foo; end
+
+        class Bar
+          def bar; end
+        end
+      end
+
+      def baz; end
+    RUBY
+
+    visitor = DummyVisitor.new
+    visitor.visit(parsed_tree)
+    assert_equal(["Foo", "foo", "Bar", "bar", "baz"], visitor.visited_nodes)
+  end
+
+  class DummyVisitor < SyntaxTree::Visitor
+    attr_reader :visited_nodes
+
+    def initialize
+      super
+      @visited_nodes = []
+    end
+
+    def visit_class_declaration(node)
+      @visited_nodes << node.constant.constant.value
+      super
+    end
+
+    def visit_def(node)
+      @visited_nodes << node.name.value
+    end
+  end
+end