From 7c0ea7b3bce5e769e2acb0e02d40288cbdf92f9e Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Wed, 24 Feb 2021 17:02:18 +0100
Subject: [PATCH] CFG: add AstNode for @in

---
 .../controlflow/internal/Completion.qll           |  3 +--
 .../controlflow/internal/ControlFlowGraphImpl.qll | 15 ++++++++++++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/ql/src/codeql_ruby/controlflow/internal/Completion.qll b/ql/src/codeql_ruby/controlflow/internal/Completion.qll
index 95f7553b59c..0c918a56dc9 100644
--- a/ql/src/codeql_ruby/controlflow/internal/Completion.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/Completion.qll
@@ -5,7 +5,6 @@
  */
 
 private import codeql_ruby.AST
-private import codeql_ruby.ast.internal.Control
 private import codeql_ruby.controlflow.ControlFlowGraph
 private import ControlFlowGraphImpl
 private import NonReturning
@@ -54,7 +53,7 @@ private predicate nestedEnsureCompletion(Completion outer, int nestLevel) {
 
 pragma[noinline]
 private predicate completionIsValidForStmt(AstNode n, Completion c) {
-  n instanceof In::Range and
+  n instanceof InRange and
   c instanceof EmptinessCompletion
   or
   n instanceof BreakStmt and
diff --git a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
index da10ca2282a..4e7b48c3214 100644
--- a/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
+++ b/ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll
@@ -31,7 +31,8 @@
  * caught up by its surrounding loop and turned into a `NormalCompletion`.
  */
 
-private import codeql_ruby.AST as AST
+private import codeql_ruby.ast.internal.AST as ASTInternal
+private import codeql_ruby.ast.internal.Control as Control
 private import codeql_ruby.ast.internal.TreeSitter::Generated
 private import AstNodes
 private import codeql_ruby.ast.internal.Variable
@@ -307,6 +308,18 @@ abstract private class PreOrderTree extends ControlFlowTree {
   final override predicate first(AstNode first) { first = this }
 }
 
+class InRange extends ASTInternal::AstNode::Range, @in {
+  final override string toString() { result = "In" }
+}
+
+class ForRange extends Control::ForExpr::Range, @for {
+  override predicate child(string label, ASTInternal::AstNode::Range child) {
+    Control::ForExpr::Range.super.child(label, child)
+    or
+    label = "<in>" and this.(AstNode).getAFieldOrChild().(In) = child
+  }
+}
+
 // TODO: remove this predicate
 predicate isValidFor(Completion c, ControlFlowTree node) {
   c instanceof SimpleCompletion and isHidden(node)