Add AST Visitor

Co-authored-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Co-authored-by: Ufuk Kayserilioglu <ufuk.kayserilioglu@shopify.com>

lib/rbs.rb

@@ -22,6 +22,7 @@
 require "rbs/ast/declarations"
 require "rbs/ast/members"
 require "rbs/ast/annotation"
+require "rbs/ast/visitor"
 require "rbs/environment"
 require "rbs/environment/use_map"
 require "rbs/environment_loader"

lib/rbs/ast/visitor.rb

@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+module RBS
+  module AST
+    # The Visitor class implements the Visitor pattern for traversing the RBS Abstract Syntax Tree (AST).
+    #
+    # It provides methods to visit each type of node in the AST, allowing for custom processing of each node type.
+    #
+    # This class is designed to be subclassed, with specific visit methods overridden to implement custom behavior for
+    # different node types.
+    #
+    # Example usage:
+    #
+    # ~~~rb
+    # class MyVisitor < RBS::AST::Visitor
+    #   def visit_declaration_class(node)
+    #     puts "Visiting class: #{node.name}"
+    #
+    #     super # call `super` to run the default visiting behavior
+    #   end
+    # end
+    #
+    # visitor = MyVisitor.new
+    # visitor.visit(ast_node)
+    # ~~~
+    class Visitor
+      def visit(node)
+        case node
+        when Declarations::Global
+          visit_declaration_global(node)
+        when Declarations::Class
+          visit_declaration_class(node)
+        when Declarations::Module
+          visit_declaration_module(node)
+        when Declarations::Constant
+          visit_declaration_constant(node)
+        when Declarations::TypeAlias
+          visit_declaration_type_alias(node)
+        when Declarations::Interface
+          visit_declaration_interface(node)
+        when Members::Alias
+          visit_member_alias(node)
+        when Members::ClassInstanceVariable
+          visit_member_class_instance_variable(node)
+        when Members::ClassVariable
+          visit_member_class_variable(node)
+        when Members::InstanceVariable
+          visit_member_instance_variable(node)
+        when Members::Private
+          visit_member_private(node)
+        when Members::Public
+          visit_member_public(node)
+        when Members::MethodDefinition
+          visit_member_method_definition(node)
+        when Members::AttrReader
+          visit_member_attr_reader(node)
+        when Members::AttrWriter
+          visit_member_attr_writer(node)
+        when Members::AttrAccessor
+          visit_member_attr_accessor(node)
+        when Members::Include
+          visit_member_include(node)
+        when Members::Prepend
+          visit_member_prepend(node)
+        when Members::Extend
+          visit_member_extend(node)
+        end
+      end
+
+      def visit_all(nodes)
+        nodes.each do |node|
+          visit(node)
+        end
+      end
+
+      def visit_declaration_global(node)
+      end
+
+      def visit_declaration_class(node)
+        visit_all(node.members)
+      end
+
+      def visit_declaration_module(node)
+        visit_all(node.members)
+      end
+
+      def visit_declaration_constant(node)
+      end
+
+      def visit_declaration_type_alias(node)
+      end
+
+      def visit_declaration_interface(node)
+        visit_all(node.members)
+      end
+
+      def visit_member_alias(node)
+      end
+
+      def visit_member_class_instance_variable(node)
+      end
+
+      def visit_member_class_variable(node)
+      end
+
+      def visit_member_instance_variable(node)
+      end
+
+      def visit_member_private(node)
+      end
+
+      def visit_member_public(node)
+      end
+
+      def visit_member_method_definition(node)
+      end
+
+      def visit_member_attr_reader(node)
+      end
+
+      def visit_member_attr_writer(node)
+      end
+
+      def visit_member_attr_accessor(node)
+      end
+
+      def visit_member_include(node)
+      end
+
+      def visit_member_prepend(node)
+      end
+
+      def visit_member_extend(node)
+      end
+    end
+  end
+end

sig/visitor.rbs

@@ -0,0 +1,49 @@
+module RBS
+  module AST
+    module Declarations
+      class Visitor
+        def visit_all: (Enumerable[(Members::Base | Declarations::Base)] nodes) -> void
+
+        def visit: ((Members::Base | Declarations::Base) node) -> void
+
+        def visit_declaration_global: (Declarations::Global node) -> void
+
+        def visit_declaration_class: (Declarations::Class node) -> void
+
+        def visit_declaration_module: (Declarations::Module node) -> void
+
+        def visit_declaration_constant: (Declarations::Constant node) -> void
+
+        def visit_declaration_type_alias: (Declarations::TypeAlias node) -> void
+
+        def visit_declaration_interface: (Declarations::Interface node) -> void
+
+        def visit_member_alias: (Members::Alias node) -> void
+
+        def visit_member_class_instance_variable: (Members::ClassInstanceVariable node) -> void
+
+        def visit_member_class_variable: (Members::ClassVariable node) -> void
+
+        def visit_member_instance_variable: (Members::InstanceVariable node) -> void
+
+        def visit_member_private: (Members::Private node) -> void
+
+        def visit_member_public: (Members::Public node) -> void
+
+        def visit_member_method_definition: (Members::MethodDefinition node) -> void
+
+        def visit_member_attr_reader: (Members::AttrReader node) -> void
+
+        def visit_member_attr_writer: (Members::AttrWriter node) -> void
+
+        def visit_member_attr_accessor: (Members::AttrAccessor node) -> void
+
+        def visit_member_include: (Members::Include node) -> void
+
+        def visit_member_prepend: (Members::Prepend node) -> void
+
+        def visit_member_extend: (Members::Extend node) -> void
+      end
+    end
+  end
+end

test/rbs/ast/visitor_test.rb

@@ -0,0 +1,69 @@
+require "test_helper"
+
+class RBS::AST::VisitorTest < Test::Unit::TestCase
+  RBS_STRING = <<~RBS
+    module Hello
+      attr_reader name: String
+
+      def hello: () -> void
+    end
+  RBS
+
+  def test_visit
+    visitor_class = Class.new(RBS::AST::Visitor) do
+      attr_reader :nodes
+
+      def initialize
+        super()
+        @nodes = []
+      end
+
+      def visit(node)
+        @nodes << node
+        super
+      end
+    end
+
+    declarations = parse_rbs_string(RBS_STRING)
+    visitor = visitor_class.new
+    visitor.visit_all(declarations)
+
+    assert_equal(
+      [
+        "RBS::AST::Declarations::Module",
+        "RBS::AST::Members::AttrReader",
+        "RBS::AST::Members::MethodDefinition"
+      ],
+      visitor.nodes.map(&:class).map(&:name)
+    )
+  end
+
+  def test_visit_node
+    visitor_class = Class.new(RBS::AST::Visitor) do
+      attr_reader :nodes
+
+      def initialize
+        super()
+        @nodes = []
+      end
+
+      def visit_member_method_definition(node)
+        @nodes << node
+      end
+    end
+
+    declarations = parse_rbs_string(RBS_STRING)
+    visitor = visitor_class.new
+    visitor.visit_all(declarations)
+
+    assert_equal(["RBS::AST::Members::MethodDefinition"], visitor.nodes.map(&:class).map(&:name))
+  end
+
+  private
+
+  def parse_rbs_string(rbs)
+    rbs = RBS::Buffer.new(content: rbs, name: "-")
+    buffer, directives, declarations = RBS::Parser.parse_signature(rbs)
+    declarations
+  end
+end