Add the visitor method on reflection

Author: kddnewton

lib/syntax_tree/reflection.rb

@@ -138,12 +138,13 @@ def initialize(name, comment)
     # as a placeholder for collecting all of the various places that nodes are
     # used.
     class Node
-      attr_reader :name, :comment, :attributes
+      attr_reader :name, :comment, :attributes, :visitor_method
 
-      def initialize(name, comment, attributes)
+      def initialize(name, comment, attributes, visitor_method)
         @name = name
         @comment = comment
         @attributes = attributes
+        @visitor_method = visitor_method
       end
     end
 
@@ -196,6 +197,10 @@ def parse_comments(statements, index)
           Attribute.new(:location, "[Location] the location of this node")
       }
 
+      # This is the name of the method tha gets called on the given visitor when
+      # the accept method is called on this node.
+      visitor_method = nil
+
       statements = main_statement.bodystmt.statements.body
       statements.each_with_index do |statement, statement_index|
         case statement
@@ -225,16 +230,25 @@ def parse_comments(statements, index)
           end
 
           attributes[attribute.name] = attribute
+        when SyntaxTree::DefNode
+          if statement.name.value == "accept"
+            call_node = statement.bodystmt.statements.body.first
+            visitor_method = call_node.message.value.to_sym
+          end
         end
       end
 
+      # If we never found a visitor method, then we have an error.
+      raise if visitor_method.nil?
+
       # Finally, set it up in the hash of nodes so that we can use it later.
       comments = parse_comments(main_statements, main_statement_index)
       node =
         Node.new(
           main_statement.constant.constant.value.to_sym,
           "#{comments.join("\n")}\n",
-          attributes
+          attributes,
+          visitor_method
         )
 
       @nodes[node.name] = node

tasks/sorbet.rake

@@ -20,6 +20,22 @@ module SyntaxTree
       generate_parent
       Reflection.nodes.sort.each { |(_, node)| generate_node(node) }
 
+      body << ClassDeclaration(
+        ConstPathRef(VarRef(Const("SyntaxTree")), Const("BasicVisitor")),
+        nil,
+        BodyStmt(
+          Statements(generate_visitor("overridable")), nil, nil, nil, nil
+        ),
+        location
+      )
+
+      body << ClassDeclaration(
+        ConstPathRef(VarRef(Const("SyntaxTree")), Const("Visitor")),
+        ConstPathRef(VarRef(Const("SyntaxTree")), Const("BasicVisitor")),
+        BodyStmt(Statements(generate_visitor("override")), nil, nil, nil, nil),
+        location
+      )
+
       Formatter.format(nil, Program(Statements(body)))
     end
 
@@ -228,6 +244,42 @@ module SyntaxTree
       )
     end
 
+    def generate_visitor(override)
+      body = []
+
+      Reflection.nodes.each do |name, node|
+        body << sig_block do
+          CallNode(
+            CallNode(
+              Ident(override),
+              Period("."),
+              sig_params do
+                BareAssocHash([
+                  Assoc(Label("node:"), 
+                  sig_type_for(SyntaxTree.const_get(name)))
+                ])
+              end,
+              nil
+            ),
+            Period("."),
+            sig_returns do
+              CallNode(VarRef(Const("T")), Period("."), Ident("untyped"), nil)
+            end,
+            nil
+          )
+        end
+
+        body << generate_def_node(node.visitor_method, Paren(
+          LParen("("),
+          Params.new(requireds: [Ident("node")], location: location)
+        ))
+
+        @line += 2
+      end
+
+      body
+    end
+
     def sig_block
       MethodAddBlock(
         CallNode(nil, nil, Ident("sig"), nil),