From 9c5da197ed7dd75207e978ea80fe80d1b21ecbca Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Mon, 15 Feb 2021 15:22:14 +0100
Subject: [PATCH] AST: add Pattern::getRestIndex

---
 ql/src/codeql_ruby/ast/Pattern.qll          | 16 +++++++++++++++-
 ql/src/codeql_ruby/ast/internal/AST.qll     |  2 --
 ql/src/codeql_ruby/ast/internal/Pattern.qll | 12 ++++++++++++
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/ql/src/codeql_ruby/ast/Pattern.qll b/ql/src/codeql_ruby/ast/Pattern.qll
index 4f975857d56..d23ddb0de52 100644
--- a/ql/src/codeql_ruby/ast/Pattern.qll
+++ b/ql/src/codeql_ruby/ast/Pattern.qll
@@ -25,7 +25,12 @@ class VariablePattern extends Pattern {
 /**
  * A tuple pattern.
  *
- * This includes both tuple patterns in parameters and assignments.
+ * This includes both tuple patterns in parameters and assignments. Example patterns:
+ * ```rb
+ * a, self.b = value
+ * (a, b), c[3] = value
+ * a, b, *rest, c, d = value
+ * ```
  */
 class TuplePattern extends Pattern {
   override TuplePattern::Range range;
@@ -35,4 +40,13 @@ class TuplePattern extends Pattern {
 
   /** Gets a sub pattern in this tuple pattern. */
   final Pattern getAnElement() { result = this.getElement(_) }
+
+  /**
+   * Gets the index of the pattern with the `*` marker on it, if it exists.
+   * In the example below the index is `2`.
+   * ```rb
+   * a, b, *rest, c, d = value
+   * ```
+   */
+  final int getRestIndex() { result = range.getRestIndex() }
 }
diff --git a/ql/src/codeql_ruby/ast/internal/AST.qll b/ql/src/codeql_ruby/ast/internal/AST.qll
index ab62aadbfc8..8a8c1a2c168 100644
--- a/ql/src/codeql_ruby/ast/internal/AST.qll
+++ b/ql/src/codeql_ruby/ast/internal/AST.qll
@@ -29,8 +29,6 @@ module AstNode {
       or
       this instanceof Generated::RestAssignment
       or
-      this = any(Generated::RestAssignment ra).getChild()
-      or
       this instanceof Generated::SymbolArray
       or
       this instanceof Generated::Interpolation
diff --git a/ql/src/codeql_ruby/ast/internal/Pattern.qll b/ql/src/codeql_ruby/ast/internal/Pattern.qll
index 808031f751d..fb9deab291a 100644
--- a/ql/src/codeql_ruby/ast/internal/Pattern.qll
+++ b/ql/src/codeql_ruby/ast/internal/Pattern.qll
@@ -20,6 +20,8 @@ predicate explicitAssignmentNode(Generated::AstNode n, Generated::AstNode assign
     parent instanceof Generated::DestructuredLeftAssignment
     or
     parent instanceof Generated::LeftAssignmentList
+    or
+    parent instanceof Generated::RestAssignment
   )
 }
 
@@ -74,6 +76,14 @@ module TuplePattern {
 
   class Range extends Pattern::Range, Range_ {
     Pattern::Range getElement(int i) {
+      exists(Generated::AstNode c | c = getChild(i) |
+        result = c.(Generated::RestAssignment).getChild()
+        or
+        not c instanceof Generated::RestAssignment and result = c
+      )
+    }
+
+    private Generated::AstNode getChild(int i) {
       result = this.(Generated::DestructuredParameter).getChild(i)
       or
       result = this.(Generated::DestructuredLeftAssignment).getChild(i)
@@ -81,6 +91,8 @@ module TuplePattern {
       result = this.(Generated::LeftAssignmentList).getChild(i)
     }
 
+    int getRestIndex() { result = unique(int i | getChild(i) instanceof Generated::RestAssignment) }
+
     override Variable getAVariable() { result = this.getElement(_).getAVariable() }
 
     override string toString() { result = "(..., ...)" }