Add WithEnvironment mixin for visitors

Author: vinistock

lib/syntax_tree.rb

@@ -19,6 +19,8 @@
 require_relative "syntax_tree/visitor/json_visitor"
 require_relative "syntax_tree/visitor/match_visitor"
 require_relative "syntax_tree/visitor/pretty_print_visitor"
+require_relative "syntax_tree/visitor/environment"
+require_relative "syntax_tree/visitor/with_environment"
 
 # Syntax Tree is a suite of tools built on top of the internal CRuby parser. It
 # provides the ability to generate a syntax tree from source, as well as the

lib/syntax_tree/visitor/environment.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module SyntaxTree
+  # The environment class is used to keep track of local variables and arguments inside a particular scope
+  class Environment
+    # [Array[Local]] The local variables and arguments defined in this
+    # environment
+    attr_reader :locals
+
+    # This class tracks the occurrences of a local variable or argument
+    class Local
+      # [Symbol] The type of the local (e.g. :argument, :variable)
+      attr_reader :type
+
+      # [Array[Location]] The locations of all occurrences of this local,
+      # including its definition
+      attr_reader :locations
+
+      def initialize(type)
+        @type = type
+        @locations = []
+      end
+
+      def <<(location)
+        @locations << location
+      end
+    end
+
+    def initialize(parent = nil)
+      @locals = {}
+      @parent = parent
+    end
+
+    # Registering a local will either insert a new entry in the locals hash or
+    # append a new location to an existing local. Notice that it's not possible
+    # to change the type of a local after it has been registered
+    #   find_local: (Ident | Label identifier, Symbol type) -> void
+    def register_local(identifier, type)
+      name = identifier.value.delete_suffix(":")
+
+      @locals[name] ||= Local.new(type)
+      @locals[name] << identifier.location
+    end
+
+    # Try to find the local given its name in this environment or any of its
+    # parents
+    #   find_local: (String name) -> Array[Location] | nil
+    def find_local(name)
+      locations = @locals[name]
+      return locations unless locations.nil?
+
+      @parent&.find_local(name)
+    end
+  end
+end

lib/syntax_tree/visitor/with_environment.rb

@@ -0,0 +1,125 @@
+# frozen_string_literal: true
+
+module SyntaxTree
+  # WithEnvironment is a module intended to be included in classes inheriting
+  # from Visitor. The module overrides a few visit methods to automatically keep
+  # track of local variables and arguments defined in the current environment.
+  # Example usage:
+  #   class MyVisitor < Visitor
+  #     include WithEnvironment
+  #
+  #     def visit_ident(node)
+  #       # Check if we're visiting an identifier for an argument, a local
+  #       variable or something else
+  #       local = current_environment.find_local(node)
+  #
+  #       if local.type == :argument
+  #         # handle identifiers for arguments
+  #       elsif local.type == :variable
+  #         # handle identifiers for variables
+  #       else
+  #         # handle other identifiers, such as method names
+  #       end
+  #   end
+  module WithEnvironment
+    def current_environment
+      @current_environment ||= Environment.new
+    end
+
+    def with_new_environment
+      previous_environment = @current_environment
+      @current_environment = Environment.new(previous_environment)
+      yield
+      @current_environment = previous_environment
+    end
+
+    # Visits for nodes that create new environments, such as classes, modules
+    # and method definitions
+    def visit_class(node)
+      with_new_environment { super }
+    end
+
+    def visit_module(node)
+      with_new_environment { super }
+    end
+
+    def visit_method_add_block(node)
+      with_new_environment { super }
+    end
+
+    def visit_def(node)
+      with_new_environment { super }
+    end
+
+    def visit_defs(node)
+      with_new_environment { super }
+    end
+
+    def visit_def_endless(node)
+      with_new_environment { super }
+    end
+
+    # Visit for keeping track of local arguments, such as method and block
+    # arguments
+    def visit_params(node)
+      node.requireds.each do |param|
+        @current_environment.register_local(param, :argument)
+      end
+
+      node.posts.each do |param|
+        @current_environment.register_local(param, :argument)
+      end
+
+      node.keywords.each do |param|
+        @current_environment.register_local(param.first, :argument)
+      end
+
+      node.optionals.each do |param|
+        @current_environment.register_local(param.first, :argument)
+      end
+
+      super
+    end
+
+    def visit_rest_param(node)
+      name = node.name
+      @current_environment.register_local(name, :argument) if name
+
+      super
+    end
+
+    def visit_kwrest_param(node)
+      name = node.name
+      @current_environment.register_local(name, :argument) if name
+
+      super
+    end
+
+    def visit_blockarg(node)
+      name = node.name
+      @current_environment.register_local(name, :argument) if name
+
+      super
+    end
+
+    # Visits for keeping track of local variables
+    def visit_aref_field(node)
+      name = node.collection.value
+      @current_environment.register_local(name, :variable) if name
+
+      super
+    end
+
+    def visit_var_field(node)
+      value = node.value
+
+      if value.is_a?(SyntaxTree::Ident)
+        @current_environment.register_local(value, :variable)
+      end
+
+      super
+    end
+
+    alias visit_var_ref visit_var_field
+  end
+end