Add MasgnNode class for masgn nodes

changelog/new_add_masgnnode_class_for_masgn_nodes.md

@@ -0,0 +1 @@
+* [#203](https://github.com/rubocop-hq/rubocop-ast/pull/203): Add classes for `masgn` and `mlhs` nodes. ([@dvandersluis][])

docs/modules/ROOT/pages/node_types.adoc

@@ -152,9 +152,9 @@ The following fields are given when relevant to nodes in the source code:
 
 |lvasgn|Local variable assignment|Two children: The variable name (symbol) and the expression.|a = some_thing|https://rubydoc.info/github/rubocop/rubocop-ast/RuboCop/AST/AsgnNode[AsgnNode]
 
-|masgn|Multiple assigment.|First set of children are all `mlhs` nodes, and the rest of the children must be expression nodes corresponding to the values in the `mlhs` nodes.|a, b, = [1, 2]|N/A
+|masgn|Multiple assigment.|First set of children are all `mlhs` nodes, and the rest of the children must be expression nodes corresponding to the values in the `mlhs` nodes.|a, b, = [1, 2]|a = some_thing|https://rubydoc.info/github/rubocop/rubocop-ast/RuboCop/AST/MasgnNode[MasgnNode]
 
-|mlhs|Multiple left-hand side. Used inside a `masgn` and block argument destructuring.|Children must all be assignment nodes. Represents the left side of a multiple assignment (`a, b` in the example).|a, b = 5, 6|N/A
+|mlhs|Multiple left-hand side. Used inside a `masgn` and block argument destructuring.|Children must all be assignment nodes (or `send` nodes). Represents the left side of a multiple assignment (`a, b` in the example).|a, b = 5, 6|https://rubydoc.info/github/rubocop/rubocop-ast/RuboCop/AST/MlhsNode[MlhsNode]
 
 |module|Module definition|Two children. First child is a `const` node for the module name. Second child is a body statement.|module Foo < Bar; end|https://rubydoc.info/github/rubocop/rubocop-ast/RuboCop/AST/ModuleNode[ModuleNode]
 

lib/rubocop/ast.rb

@@ -61,6 +61,8 @@
 require_relative 'ast/node/int_node'
 require_relative 'ast/node/keyword_splat_node'
 require_relative 'ast/node/lambda_node'
+require_relative 'ast/node/masgn_node'
+require_relative 'ast/node/mlhs_node'
 require_relative 'ast/node/module_node'
 require_relative 'ast/node/next_node'
 require_relative 'ast/node/op_asgn_node'

lib/rubocop/ast/builder.rb

@@ -64,6 +64,8 @@ class Builder < Parser::Builders::Default
         kwargs:       HashNode,
         kwsplat:      KeywordSplatNode,
         lambda:       LambdaNode,
+        masgn:        MasgnNode,
+        mlhs:         MlhsNode,
         module:       ModuleNode,
         next:         NextNode,
         op_asgn:      OpAsgnNode,

lib/rubocop/ast/node/masgn_node.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module RuboCop
+  module AST
+    # A node extension for `masgn` nodes.
+    # This will be used in place of a plain node when the builder constructs
+    # the AST, making its methods available to all assignment nodes within RuboCop.
+    class MasgnNode < Node
+      # @return [MlhsNode] the `mlhs` node
+      def lhs
+        # The first child is a `mlhs` node
+        node_parts[0]
+      end
+
+      # @return [Array<Node>] the assignment nodes of the multiple assignment
+      def assignments
+        lhs.assignments
+      end
+
+      # @return [Array<Symbol>] names of all the variables being assigned
+      def names
+        assignments.map do |assignment|
+          if assignment.send_type? || assignment.indexasgn_type?
+            assignment.source
+          else
+            assignment.name
+          end
+        end
+      end
+
+      # The expression being assigned to the variable.
+      #
+      # @return [Node] the expression being assigned.
+      def expression
+        node_parts[1]
+      end
+      alias rhs expression
+
+      # @return [Array<Node>] values being assigned on the RHS of the multiple assignment
+      def values
+        array? ? expression.children : [expression]
+      end
+
+      # @return [Boolean] whether the expression has multiple values
+      def array?
+        expression.array_type?
+      end
+    end
+  end
+end

lib/rubocop/ast/node/mlhs_node.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module RuboCop
+  module AST
+    # A node extension for `mlhs` nodes.
+    # This will be used in place of a plain node when the builder constructs
+    # the AST, making its methods available to all assignment nodes within RuboCop.
+    class MlhsNode < Node
+      # Returns all the assignment nodes on the left hand side (LHS) of a multiple assigment.
+      # These are generally assigment nodes (`lvasgn`, `ivasgn`, `cvasgn`, `gvasgn`, `casgn`)
+      # but can also be `send` nodes in case of `foo.bar, ... =` or `foo[:bar], ... =`.
+      #
+      # @return [Array<Node>] the assignment nodes of the multiple assignment LHS
+      def assignments
+        child_nodes.flat_map do |node|
+          if node.splat_type?
+            node.child_nodes.first
+          elsif node.mlhs_type?
+            node.assignments
+          else
+            node
+          end
+        end
+      end
+    end
+  end
+end

spec/rubocop/ast/masgn_node_spec.rb

@@ -0,0 +1,118 @@
+# frozen_string_literal: true
+
+RSpec.describe RuboCop::AST::MasgnNode do
+  let(:masgn_node) { parse_source(source).ast }
+  let(:source) { 'x, y = z' }
+
+  describe '.new' do
+    context 'with a `masgn` node' do
+      it { expect(masgn_node).to be_a(described_class) }
+    end
+  end
+
+  describe '#names' do
+    subject { masgn_node.names }
+
+    let(:source) { 'a, @b, @@c, $d, E, *f = z' }
+
+    it { is_expected.to eq(%i[a @b @@c $d E f]) }
+
+    context 'with nested `mlhs` nodes' do
+      let(:source) { 'a, (b, c) = z' }
+
+      it { is_expected.to eq(%i[a b c]) }
+    end
+
+    context 'with nested assignment on LHS' do
+      let(:source) { 'a, b[c+=1] = z' }
+
+      it { is_expected.to eq([:a, 'b[c+=1]']) }
+    end
+
+    context 'with a method chain on LHS' do
+      let(:source) { 'a, b.c = z' }
+
+      it { is_expected.to eq([:a, 'b.c']) }
+    end
+  end
+
+  describe '#expression' do
+    include AST::Sexp
+
+    subject { masgn_node.expression }
+
+    context 'with variables' do
+      it { is_expected.to eq(s(:send, nil, :z)) }
+    end
+
+    context 'with a LHS splat' do
+      let(:source) { 'x, *y = z' }
+
+      it { is_expected.to eq(s(:send, nil, :z)) }
+    end
+
+    context 'with multiple RHS values' do
+      let(:source) { 'x, y = 1, 2' }
+
+      it { is_expected.to eq(s(:array, s(:int, 1), s(:int, 2))) }
+    end
+
+    context 'with an RHS splat' do
+      let(:source) { 'x, y = *z' }
+
+      it { is_expected.to eq(s(:array, s(:splat, s(:send, nil, :z)))) }
+    end
+
+    context 'with assignment on RHS' do
+      let(:source) { 'x, y = 1, z+=2' }
+
+      it { is_expected.to eq(s(:array, s(:int, 1), s(:op_asgn, s(:lvasgn, :z), :+, s(:int, 2)))) }
+    end
+  end
+
+  describe '#values' do
+    include AST::Sexp
+
+    subject { masgn_node.values }
+
+    context 'when the RHS has a single value' do
+      let(:source) { 'x, y = z' }
+
+      it { is_expected.to eq([s(:send, nil, :z)]) }
+    end
+
+    context 'when the RHS has a multiple values' do
+      let(:source) { 'x, y = u, v' }
+
+      it { is_expected.to eq([s(:send, nil, :u), s(:send, nil, :v)]) }
+    end
+
+    context 'when the RHS has a splat' do
+      let(:source) { 'x, y = *z' }
+
+      it { is_expected.to eq([s(:splat, s(:send, nil, :z))]) }
+    end
+  end
+
+  describe '#array?' do
+    subject { masgn_node.array? }
+
+    context 'when the RHS has a single value' do
+      let(:source) { 'x, y = z' }
+
+      it { is_expected.to eq(false) }
+    end
+
+    context 'when the RHS has a multiple values' do
+      let(:source) { 'x, y = u, v' }
+
+      it { is_expected.to eq(true) }
+    end
+
+    context 'when the RHS has a splat' do
+      let(:source) { 'x, y = *z' }
+
+      it { is_expected.to eq(true) }
+    end
+  end
+end

spec/rubocop/ast/mlhs_node_spec.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+RSpec.describe RuboCop::AST::MlhsNode do
+  let(:mlhs_node) { parse_source(source).ast.node_parts[0] }
+
+  describe '.new' do
+    context 'with a `masgn` node' do
+      let(:source) { 'x, y = z' }
+
+      it { expect(mlhs_node).to be_a(described_class) }
+    end
+  end
+
+  describe '#assignments' do
+    include AST::Sexp
+
+    subject { mlhs_node.assignments }
+
+    context 'with variables' do
+      let(:source) { 'x, y = z' }
+
+      it { is_expected.to eq([s(:lvasgn, :x), s(:lvasgn, :y)]) }
+    end
+
+    context 'with a splat' do
+      let(:source) { 'x, *y = z' }
+
+      it { is_expected.to eq([s(:lvasgn, :x), s(:lvasgn, :y)]) }
+    end
+
+    context 'with nested `mlhs` nodes' do
+      let(:source) { 'a, (b, c) = z' }
+
+      it { is_expected.to eq([s(:lvasgn, :a), s(:lvasgn, :b), s(:lvasgn, :c)]) }
+    end
+
+    context 'with different variable types' do
+      let(:source) { 'a, @b, @@c, $d, E, *f = z' }
+      let(:expected_nodes) do
+        [
+          s(:lvasgn, :a),
+          s(:ivasgn, :@b),
+          s(:cvasgn, :@@c),
+          s(:gvasgn, :$d),
+          s(:casgn, nil, :E),
+          s(:lvasgn, :f)
+        ]
+      end
+
+      it { is_expected.to eq(expected_nodes) }
+    end
+
+    context 'with assignment on RHS' do
+      let(:source) { 'x, y = 1, z += 2' }
+
+      it { is_expected.to eq([s(:lvasgn, :x), s(:lvasgn, :y)]) }
+    end
+
+    context 'with nested assignment on LHS' do
+      let(:source) { 'a, b[c+=1] = z' }
+
+      if RuboCop::AST::Builder.emit_index
+        let(:expected_nodes) do
+          [
+            s(:lvasgn, :a),
+            s(:indexasgn,
+              s(:send, nil, :b),
+              s(:op_asgn,
+                s(:lvasgn, :c), :+, s(:int, 1)))
+          ]
+        end
+      else
+        let(:expected_nodes) do
+          [
+            s(:lvasgn, :a),
+            s(:send,
+              s(:send, nil, :b), :[]=,
+              s(:op_asgn,
+                s(:lvasgn, :c), :+, s(:int, 1)))
+          ]
+        end
+      end
+
+      it { is_expected.to eq(expected_nodes) }
+    end
+  end
+end