Add WithEnvironment mixin for visitors

README.md

@@ -368,16 +368,16 @@ program = SyntaxTree.parse("1 + 1")
 puts program.construct_keys
 
 # SyntaxTree::Program[
-#   statements: SyntaxTree::Statements[            
-#     body: [                                      
-#       SyntaxTree::Binary[                        
-#         left: SyntaxTree::Int[value: "1"],       
-#         operator: :+,                            
-#         right: SyntaxTree::Int[value: "1"]       
-#       ]                                          
-#     ]                                            
-#   ]                                              
-# ] 
+#   statements: SyntaxTree::Statements[
+#     body: [
+#       SyntaxTree::Binary[
+#         left: SyntaxTree::Int[value: "1"],
+#         operator: :+,
+#         right: SyntaxTree::Int[value: "1"]
+#       ]
+#     ]
+#   ]
+# ]
 ```
 
 ## Visitor
@@ -447,6 +447,28 @@ end
 
 The visitor defined above will error out unless it's only visiting a `SyntaxTree::Int` node. This is useful in a couple of ways, e.g., if you're trying to define a visitor to handle the whole tree but it's currently a work-in-progress.
 
+### WithEnvironment
+
+The `WithEnvironment` module can be included in visitors to automatically keep track of local variables and arguments
+defined inside each environment. A `current_environment` accessor is made availble to the request, allowing it to find
+all usages and definitions of a local.
+
+```ruby
+class MyVisitor < Visitor
+  include WithEnvironment
+
+  def visit_ident(node)
+    # find_local will return a Local for any local variables or arguments present in the current environment or nil if
+    # the identifier is not a local
+    local = current_environment.find_local(node)
+
+    puts local.type # print the type of the local (:variable or :argument)
+    puts local.definitions # print the array of locations where this local is defined
+    puts local.usages # print the array of locations where this local occurs
+  end
+end
+```
+
 ## Language server
 
 Syntax Tree additionally ships with a language server conforming to the [language server protocol](https://microsoft.github.io/language-server-protocol/). It can be invoked through the CLI by running:

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,81 @@
+# 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 definitions and assignments of
+      # this local
+      attr_reader :definitions
+
+      # [Array[Location]] The locations of all usages of this local
+      attr_reader :usages
+
+      #   initialize: (Symbol type) -> void
+      def initialize(type)
+        @type = type
+        @definitions = []
+        @usages = []
+      end
+
+      #   add_definition: (Location location) -> void
+      def add_definition(location)
+        @definitions << location
+      end
+
+      #   add_usage: (Location location) -> void
+      def add_usage(location)
+        @usages << location
+      end
+    end
+
+    #   initialize: (Environment | nil parent) -> void
+    def initialize(parent = nil)
+      @locals = {}
+      @parent = parent
+    end
+
+    # Adding a local definition will either insert a new entry in the locals
+    # hash or append a new definition location to an existing local. Notice that
+    # it's not possible to change the type of a local after it has been
+    # registered
+    #   add_local_definition: (Ident | Label identifier, Symbol type) -> void
+    def add_local_definition(identifier, type)
+      name = identifier.value.delete_suffix(":")
+
+      @locals[name] ||= Local.new(type)
+      @locals[name].add_definition(identifier.location)
+    end
+
+    # Adding a local usage will either insert a new entry in the locals
+    # hash or append a new usage location to an existing local. Notice that
+    # it's not possible to change the type of a local after it has been
+    # registered
+    #   add_local_usage: (Ident | Label identifier, Symbol type) -> void
+    def add_local_usage(identifier, type)
+      name = identifier.value.delete_suffix(":")
+
+      @locals[name] ||= Local.new(type)
+      @locals[name].add_usage(identifier.location)
+    end
+
+    # Try to find the local given its name in this environment or any of its
+    # parents
+    #   find_local: (String name) -> Local | nil
+    def find_local(name)
+      local = @locals[name]
+      return local unless local.nil?
+
+      @parent&.find_local(name)
+    end
+  end
+end

lib/syntax_tree/visitor/with_environment.rb

@@ -0,0 +1,141 @@
+# 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
+    ensure
+      @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.add_local_definition(param, :argument)
+      end
+
+      node.posts.each do |param|
+        @current_environment.add_local_definition(param, :argument)
+      end
+
+      node.keywords.each do |param|
+        @current_environment.add_local_definition(param.first, :argument)
+      end
+
+      node.optionals.each do |param|
+        @current_environment.add_local_definition(param.first, :argument)
+      end
+
+      super
+    end
+
+    def visit_rest_param(node)
+      name = node.name
+      @current_environment.add_local_definition(name, :argument) if name
+
+      super
+    end
+
+    def visit_kwrest_param(node)
+      name = node.name
+      @current_environment.add_local_definition(name, :argument) if name
+
+      super
+    end
+
+    def visit_blockarg(node)
+      name = node.name
+      @current_environment.add_local_definition(name, :argument) if name
+
+      super
+    end
+
+    # Visit for keeping track of local variable definitions
+    def visit_var_field(node)
+      value = node.value
+
+      if value.is_a?(SyntaxTree::Ident)
+        @current_environment.add_local_definition(value, :variable)
+      end
+
+      super
+    end
+
+    alias visit_pinned_var_ref visit_var_field
+
+    # Visits for keeping track of variable and argument usages
+    def visit_aref_field(node)
+      name = node.collection.value
+      @current_environment.add_local_usage(name, :variable) if name
+
+      super
+    end
+
+    def visit_var_ref(node)
+      value = node.value
+
+      if value.is_a?(SyntaxTree::Ident)
+        definition = @current_environment.find_local(value.value)
+
+        if definition
+          @current_environment.add_local_usage(value, definition.type)
+        end
+      end
+
+      super
+    end
+  end
+end