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,49 @@
+# frozen_string_literal: true
+
+module SyntaxTree
+  class Environment
+    # [Array[Local]] The local variables and arguments defined in this environment
+    attr_reader :locals
+
+    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,102 @@
+# 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
+
+    # Visits for nodes that create new environments, such as classes, modules and method definitions
+    def visit_with_new_environment(node)
+      previous_environment = @current_environment
+      @current_environment = Environment.new(previous_environment)
+      super
+      @current_environment = previous_environment
+    end
+
+    alias visit_class visit_with_new_environment
+    alias visit_module visit_with_new_environment
+    alias visit_method_add_block visit_with_new_environment
+    alias visit_def visit_with_new_environment
+    alias visit_defs visit_with_new_environment
+    alias visit_def_endless visit_with_new_environment
+
+    # 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_register_param(node)
+      name = node.name
+      @current_environment.register_local(name, :argument) if name
+
+      super
+    end
+
+    alias visit_rest_param visit_register_param
+    alias visit_kwrest_param visit_register_param
+    alias visit_blockarg visit_register_param
+
+    # Visits for keeping track of local variables
+    def visit_a_ref_field(node)
+      name = node.collection.value
+      @current_environment.register_local(name, :variable) if name
+
+      super
+    end
+
+    def visit_m_assign(node)
+      node.target.parts.each do |var_ref|
+        @current_environment.register_local(var_ref.value, :variable)
+      end
+
+      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