Merge pull request github#13937 from hvitved/ruby/for-loop-desugar

Ruby: Improve desugaring of `for` loops

Author: hvitved

ruby/ql/lib/codeql/ruby/ast/Operation.qll

@@ -48,7 +48,7 @@ class UnaryLogicalOperation extends UnaryOperation, TUnaryLogicalOperation { }
  * not params.empty?
  * ```
  */
-class NotExpr extends UnaryLogicalOperation, TNotExpr {
+class NotExpr extends UnaryLogicalOperation instanceof NotExprImpl {
   final override string getAPrimaryQlClass() { result = "NotExpr" }
 }
 
@@ -118,7 +118,7 @@ class ComplementExpr extends UnaryBitwiseOperation, TComplementExpr {
  * defined? some_method
  * ```
  */
-class DefinedExpr extends UnaryOperation, TDefinedExpr {
+class DefinedExpr extends UnaryOperation instanceof DefinedExprImpl {
   final override string getAPrimaryQlClass() { result = "DefinedExpr" }
 }
 

ruby/ql/lib/codeql/ruby/ast/internal/AST.qll

@@ -123,7 +123,8 @@ private module Cached {
     TConstantWriteAccessSynth(Ast::AstNode parent, int i, string value) {
       mkSynthChild(ConstantWriteAccessKind(value), parent, i)
     } or
-    TDefinedExpr(Ruby::Unary g) { g instanceof @ruby_unary_definedquestion } or
+    TDefinedExprReal(Ruby::Unary g) { g instanceof @ruby_unary_definedquestion } or
+    TDefinedExprSynth(Ast::AstNode parent, int i) { mkSynthChild(DefinedExprKind(), parent, i) } or
     TDelimitedSymbolLiteral(Ruby::DelimitedSymbol g) or
     TDestructuredLeftAssignment(Ruby::DestructuredLeftAssignment g) {
       not strictcount(int i | exists(g.getParent().(Ruby::LeftAssignmentList).getChild(i))) = 1
@@ -228,7 +229,8 @@ private module Cached {
     TNilLiteralReal(Ruby::Nil g) or
     TNilLiteralSynth(Ast::AstNode parent, int i) { mkSynthChild(NilLiteralKind(), parent, i) } or
     TNoRegExpMatchExpr(Ruby::Binary g) { g instanceof @ruby_binary_bangtilde } or
-    TNotExpr(Ruby::Unary g) { g instanceof @ruby_unary_bang or g instanceof @ruby_unary_not } or
+    TNotExprReal(Ruby::Unary g) { g instanceof @ruby_unary_bang or g instanceof @ruby_unary_not } or
+    TNotExprSynth(Ast::AstNode parent, int i) { mkSynthChild(NotExprKind(), parent, i) } or
     TOptionalParameter(Ruby::OptionalParameter g) or
     TPair(Ruby::Pair g) or
     TParenthesizedExpr(Ruby::ParenthesizedStatements g) or
@@ -354,21 +356,21 @@ private module Cached {
         TBitwiseOrExprReal or TBitwiseXorExprReal or TBlockArgument or TBlockParameter or
         TBraceBlockReal or TBreakStmt or TCaseEqExpr or TCaseExpr or TCaseMatchReal or
         TCharacterLiteral or TClassDeclaration or TClassVariableAccessReal or TComplementExpr or
-        TComplexLiteral or TDefinedExpr or TDelimitedSymbolLiteral or TDestructuredLeftAssignment or
-        TDestructuredParameter or TDivExprReal or TDo or TDoBlock or TElementReference or
-        TElseReal or TElsif or TEmptyStmt or TEncoding or TEndBlock or TEnsure or TEqExpr or
-        TExponentExprReal or TFalseLiteral or TFile or TFindPattern or TFloatLiteral or TForExpr or
-        TForwardParameter or TForwardArgument or TGEExpr or TGTExpr or TGlobalVariableAccessReal or
-        THashKeySymbolLiteral or THashLiteral or THashPattern or THashSplatExpr or
-        THashSplatNilParameter or THashSplatParameter or THereDoc or TIdentifierMethodCall or
-        TIfReal or TIfModifierExpr or TInClauseReal or TInstanceVariableAccessReal or
-        TIntegerLiteralReal or TKeywordParameter or TLEExpr or TLShiftExprReal or TLTExpr or
-        TLambda or TLeftAssignmentList or TLine or TLocalVariableAccessReal or
-        TLogicalAndExprReal or TLogicalOrExprReal or TMethod or TMatchPattern or
-        TModuleDeclaration or TModuloExprReal or TMulExprReal or TNEExpr or TNextStmt or
-        TNilLiteralReal or TNoRegExpMatchExpr or TNotExpr or TOptionalParameter or TPair or
-        TParenthesizedExpr or TParenthesizedPattern or TRShiftExprReal or TRangeLiteralReal or
-        TRationalLiteral or TRedoStmt or TRegExpLiteral or TRegExpMatchExpr or
+        TComplexLiteral or TDefinedExprReal or TDelimitedSymbolLiteral or
+        TDestructuredLeftAssignment or TDestructuredParameter or TDivExprReal or TDo or TDoBlock or
+        TElementReference or TElseReal or TElsif or TEmptyStmt or TEncoding or TEndBlock or
+        TEnsure or TEqExpr or TExponentExprReal or TFalseLiteral or TFile or TFindPattern or
+        TFloatLiteral or TForExpr or TForwardParameter or TForwardArgument or TGEExpr or TGTExpr or
+        TGlobalVariableAccessReal or THashKeySymbolLiteral or THashLiteral or THashPattern or
+        THashSplatExpr or THashSplatNilParameter or THashSplatParameter or THereDoc or
+        TIdentifierMethodCall or TIfReal or TIfModifierExpr or TInClauseReal or
+        TInstanceVariableAccessReal or TIntegerLiteralReal or TKeywordParameter or TLEExpr or
+        TLShiftExprReal or TLTExpr or TLambda or TLeftAssignmentList or TLine or
+        TLocalVariableAccessReal or TLogicalAndExprReal or TLogicalOrExprReal or TMethod or
+        TMatchPattern or TModuleDeclaration or TModuloExprReal or TMulExprReal or TNEExpr or
+        TNextStmt or TNilLiteralReal or TNoRegExpMatchExpr or TNotExprReal or TOptionalParameter or
+        TPair or TParenthesizedExpr or TParenthesizedPattern or TRShiftExprReal or
+        TRangeLiteralReal or TRationalLiteral or TRedoStmt or TRegExpLiteral or TRegExpMatchExpr or
         TRegularArrayLiteral or TRegularMethodCall or TRegularStringLiteral or TRegularSuperCall or
         TRescueClause or TRescueModifierExpr or TRetryStmt or TReturnStmt or
         TScopeResolutionConstantAccess or TSelfReal or TSimpleParameterReal or
@@ -438,7 +440,7 @@ private module Cached {
     n = TClassVariableAccessReal(result, _) or
     n = TComplementExpr(result) or
     n = TComplexLiteral(result) or
-    n = TDefinedExpr(result) or
+    n = TDefinedExprReal(result) or
     n = TDelimitedSymbolLiteral(result) or
     n = TDestructuredLeftAssignment(result) or
     n = TDivExprReal(result) or
@@ -495,7 +497,7 @@ private module Cached {
     n = TNextStmt(result) or
     n = TNilLiteralReal(result) or
     n = TNoRegExpMatchExpr(result) or
-    n = TNotExpr(result) or
+    n = TNotExprReal(result) or
     n = TOptionalParameter(result) or
     n = TPair(result) or
     n = TParenthesizedExpr(result) or
@@ -585,6 +587,8 @@ private module Cached {
     or
     result = TConstantWriteAccessSynth(parent, i, _)
     or
+    result = TDefinedExprSynth(parent, i)
+    or
     result = TDivExprSynth(parent, i)
     or
     result = TElseSynth(parent, i)
@@ -617,6 +621,8 @@ private module Cached {
     or
     result = TNilLiteralSynth(parent, i)
     or
+    result = TNotExprSynth(parent, i)
+    or
     result = TRangeLiteralSynth(parent, i, _)
     or
     result = TRShiftExprSynth(parent, i)
@@ -789,10 +795,14 @@ class TNamespace = TClassDeclaration or TModuleDeclaration;
 
 class TOperation = TUnaryOperation or TBinaryOperation or TAssignment;
 
+class TDefinedExpr = TDefinedExprReal or TDefinedExprSynth;
+
 class TUnaryOperation =
   TUnaryLogicalOperation or TUnaryArithmeticOperation or TUnaryBitwiseOperation or TDefinedExpr or
       TSplatExpr or THashSplatExpr;
 
+class TNotExpr = TNotExprReal or TNotExprSynth;
+
 class TUnaryLogicalOperation = TNotExpr;
 
 class TUnaryArithmeticOperation = TUnaryPlusExpr or TUnaryMinusExpr;

ruby/ql/lib/codeql/ruby/ast/internal/Operation.qll

@@ -35,6 +35,16 @@ class UnaryOperationGenerated extends UnaryOperationImpl {
   final override string getOperatorImpl() { result = g.getOperator() }
 }
 
+abstract class NotExprImpl extends UnaryOperationImpl, TNotExpr { }
+
+class NotExprReal extends NotExprImpl, UnaryOperationGenerated, TNotExprReal { }
+
+class NotExprSynth extends NotExprImpl, TNotExprSynth {
+  final override string getOperatorImpl() { result = "!" }
+
+  final override Expr getOperandImpl() { synthChild(this, 0, result) }
+}
+
 class SplatExprReal extends UnaryOperationImpl, TSplatExprReal {
   private Ruby::SplatArgument g;
 
@@ -67,6 +77,16 @@ class HashSplatExprImpl extends UnaryOperationImpl, THashSplatExpr {
   final override string getOperatorImpl() { result = "**" }
 }
 
+abstract class DefinedExprImpl extends UnaryOperationImpl, TDefinedExpr { }
+
+class DefinedExprReal extends DefinedExprImpl, UnaryOperationGenerated, TDefinedExprReal { }
+
+class DefinedExprSynth extends DefinedExprImpl, TDefinedExprSynth {
+  final override string getOperatorImpl() { result = "defined?" }
+
+  final override Expr getOperandImpl() { synthChild(this, 0, result) }
+}
+
 abstract class BinaryOperationImpl extends OperationImpl, MethodCallImpl, TBinaryOperation {
   abstract Stmt getLeftOperandImpl();
 

ruby/ql/lib/codeql/ruby/ast/internal/Synthesis.qll

@@ -21,6 +21,7 @@ newtype SynthKind =
   BraceBlockKind() or
   CaseMatchKind() or
   ClassVariableAccessKind(ClassVariable v) or
+  DefinedExprKind() or
   DivExprKind() or
   ElseKind() or
   ExponentExprKind() or
@@ -40,6 +41,7 @@ newtype SynthKind =
   ModuloExprKind() or
   MulExprKind() or
   NilLiteralKind() or
+  NotExprKind() or
   RangeLiteralKind(boolean inclusive) { inclusive in [false, true] } or
   RShiftExprKind() or
   SimpleParameterKind() or
@@ -1258,6 +1260,7 @@ private module HashLiteralDesugar {
  * ```
  * desugars to, roughly,
  * ```rb
+ * if not defined? x then x = nil end
  * xs.each { |__synth__0| x = __synth__0; <loop_body> }
  * ```
  *
@@ -1267,58 +1270,160 @@ private module HashLiteralDesugar {
  * scoped to the synthesized block.
  */
 private module ForLoopDesugar {
+  private Ruby::AstNode getForLoopPatternChild(Ruby::For for) {
+    result = for.getPattern()
+    or
+    result.getParent() = getForLoopPatternChild(for)
+  }
+
+  /** Holds if `n` is an access to variable `v` in the pattern of `for`. */
+  pragma[nomagic]
+  private predicate forLoopVariableAccess(Ruby::For for, Ruby::AstNode n, VariableReal v) {
+    n = getForLoopPatternChild(for) and
+    access(n, v)
+  }
+
+  /** Holds if `v` is the `i`th iteration variable of `for`. */
+  private predicate forLoopVariable(Ruby::For for, VariableReal v, int i) {
+    v =
+      rank[i + 1](VariableReal v0, Ruby::AstNode n, Location l |
+        forLoopVariableAccess(for, n, v0) and
+        l = n.getLocation()
+      |
+        v0 order by l.getStartLine(), l.getStartColumn()
+      )
+  }
+
+  /** Gets the number of iteration variables of `for`. */
+  private int forLoopVariableCount(Ruby::For for) {
+    result = count(int j | forLoopVariable(for, _, j))
+  }
+
+  private Ruby::For toTsFor(ForExpr for) { for = TForExpr(result) }
+
+  /**
+   * Synthesizes an assignment
+   * ```rb
+   * if not defined? v then v = nil end
+   * ```
+   * anchored at index `rootIndex` of `root`.
+   */
+  bindingset[root, rootIndex, v]
+  private predicate nilAssignUndefined(
+    AstNode root, int rootIndex, AstNode parent, int i, Child child, VariableReal v
+  ) {
+    parent = root and
+    i = rootIndex and
+    child = SynthChild(IfKind())
+    or
+    exists(AstNode if_ | if_ = TIfSynth(root, rootIndex) |
+      parent = if_ and
+      i = 0 and
+      child = SynthChild(NotExprKind())
+      or
+      exists(AstNode not_ | not_ = TNotExprSynth(if_, 0) |
+        parent = not_ and
+        i = 0 and
+        child = SynthChild(DefinedExprKind())
+        or
+        parent = TDefinedExprSynth(not_, 0) and
+        i = 0 and
+        child = SynthChild(LocalVariableAccessRealKind(v))
+      )
+      or
+      parent = if_ and
+      i = 1 and
+      child = SynthChild(AssignExprKind())
+      or
+      parent = TAssignExprSynth(if_, 1) and
+      (
+        i = 0 and
+        child = SynthChild(LocalVariableAccessRealKind(v))
+        or
+        i = 1 and
+        child = SynthChild(NilLiteralKind())
+      )
+    )
+  }
+
   pragma[nomagic]
   private predicate forLoopSynthesis(AstNode parent, int i, Child child) {
     exists(ForExpr for |
-      // each call
       parent = for and
       i = -1 and
-      child = SynthChild(MethodCallKind("each", false, 0))
+      child = SynthChild(StmtSequenceKind())
       or
-      exists(MethodCall eachCall | eachCall = TMethodCallSynth(for, -1, "each", false, 0) |
-        // receiver
-        parent = eachCall and
-        i = 0 and
-        child = childRef(for.getValue()) // value is the Enumerable
-        or
-        parent = eachCall and
-        i = 1 and
-        child = SynthChild(BraceBlockKind())
+      exists(AstNode seq | seq = TStmtSequenceSynth(for, -1) |
+        exists(VariableReal v, int j | forLoopVariable(toTsFor(for), v, j) |
+          nilAssignUndefined(seq, j, parent, i, child, v)
+        )
         or
-        exists(Block block | block = TBraceBlockSynth(eachCall, 1) |
-          // block params
-          parent = block and
-          i = 0 and
-          child = SynthChild(SimpleParameterKind())
+        exists(int numberOfVars | numberOfVars = forLoopVariableCount(toTsFor(for)) |
+          // each call
+          parent = seq and
+          i = numberOfVars and
+          child = SynthChild(MethodCallKind("each", false, 0))
           or
-          exists(SimpleParameter param | param = TSimpleParameterSynth(block, 0) |
-            parent = param and
+          exists(MethodCall eachCall |
+            eachCall = TMethodCallSynth(seq, numberOfVars, "each", false, 0)
+          |
+            // receiver
+            parent = eachCall and
             i = 0 and
-            child = SynthChild(LocalVariableAccessSynthKind(TLocalVariableSynth(param, 0)))
+            child = childRef(for.getValue()) // value is the Enumerable
             or
-            // assignment to pattern from for loop to synth parameter
-            parent = block and
+            parent = eachCall and
             i = 1 and
-            child = SynthChild(AssignExprKind())
+            child = SynthChild(BraceBlockKind())
             or
-            parent = TAssignExprSynth(block, 1) and
-            (
+            exists(Block block | block = TBraceBlockSynth(eachCall, 1) |
+              // block params
+              parent = block and
               i = 0 and
-              child = childRef(for.getPattern())
+              child = SynthChild(SimpleParameterKind())
               or
-              i = 1 and
-              child = SynthChild(LocalVariableAccessSynthKind(TLocalVariableSynth(param, 0)))
+              exists(SimpleParameter param | param = TSimpleParameterSynth(block, 0) |
+                parent = param and
+                i = 0 and
+                child = SynthChild(LocalVariableAccessSynthKind(TLocalVariableSynth(param, 0)))
+                or
+                // assignment to pattern from for loop to synth parameter
+                parent = block and
+                i = 1 and
+                child = SynthChild(AssignExprKind())
+                or
+                parent = TAssignExprSynth(block, 1) and
+                (
+                  i = 0 and
+                  child = childRef(for.getPattern())
+                  or
+                  i = 1 and
+                  child = SynthChild(LocalVariableAccessSynthKind(TLocalVariableSynth(param, 0)))
+                )
+              )
+              or
+              // rest of block body
+              parent = block and
+              child = childRef(for.getBody().(Do).getStmt(i - 2))
             )
           )
-          or
-          // rest of block body
-          parent = block and
-          child = childRef(for.getBody().(Do).getStmt(i - 2))
         )
       )
     )
   }
 
+  pragma[nomagic]
+  private predicate isDesugaredInitNode(ForExpr for, Variable v, AstNode n) {
+    exists(StmtSequence seq, AssignExpr ae |
+      seq = for.getDesugared() and
+      n = seq.getStmt(_) and
+      ae = n.(IfExpr).getThen() and
+      v = ae.getLeftOperand().getAVariable()
+    )
+    or
+    isDesugaredInitNode(for, v, n.getParent())
+  }
+
   private class ForLoopSynthesis extends Synthesis {
     final override predicate child(AstNode parent, int i, Child child) {
       forLoopSynthesis(parent, i, child)
@@ -1338,6 +1443,14 @@ private module ForLoopDesugar {
     final override predicate excludeFromControlFlowTree(AstNode n) {
       n = any(ForExpr for).getBody()
     }
+
+    final override predicate location(AstNode n, Location l) {
+      exists(ForExpr for, Ruby::AstNode access, Variable v |
+        forLoopVariableAccess(toTsFor(for), access, v) and
+        isDesugaredInitNode(for, v, n) and
+        l = access.getLocation()
+      )
+    }
   }
 }