CFG: model while, until and variants

ql/src/codeql_ruby/controlflow/internal/AstNodes.qll

@@ -4,6 +4,7 @@
  */
 
 private import codeql_ruby.ast.internal.TreeSitter::Generated
+private import codeql_ruby.controlflow.internal.Completion
 
 class LogicalNotAstNode extends Unary {
   AstNode operand;
@@ -74,6 +75,46 @@ private class ConditionalAstNode extends IfElsifAstNode, Conditional {
   override AstNode getAlternativeNode() { result = this.getAlternative() }
 }
 
+private class CondLoop = @while or @while_modifier or @until or @until_modifier;
+
+class ConditionalLoopAstNode extends AstNode, CondLoop {
+  AstNode getCondition() { none() }
+
+  AstNode getBody() { none() }
+
+  predicate continueLoop(BooleanCompletion c) { c instanceof TrueCompletion }
+
+  final predicate endLoop(BooleanCompletion c) { continueLoop(c.getDual()) }
+}
+
+private class WhileLoop extends ConditionalLoopAstNode, While {
+  override UnderscoreStatement getCondition() { result = While.super.getCondition() }
+
+  override Do getBody() { result = While.super.getBody() }
+}
+
+private class WhileModifierLoop extends ConditionalLoopAstNode, WhileModifier {
+  override AstNode getCondition() { result = WhileModifier.super.getCondition() }
+
+  override UnderscoreStatement getBody() { result = WhileModifier.super.getBody() }
+}
+
+private class UntilLoop extends ConditionalLoopAstNode, Until {
+  override UnderscoreStatement getCondition() { result = Until.super.getCondition() }
+
+  override Do getBody() { result = Until.super.getBody() }
+
+  override predicate continueLoop(BooleanCompletion c) { c instanceof FalseCompletion }
+}
+
+private class UntilModifierLoop extends ConditionalLoopAstNode, UntilModifier {
+  override AstNode getCondition() { result = UntilModifier.super.getCondition() }
+
+  override UnderscoreStatement getBody() { result = UntilModifier.super.getBody() }
+
+  override predicate continueLoop(BooleanCompletion c) { c instanceof FalseCompletion }
+}
+
 class ParenthesizedStatement extends ParenthesizedStatements {
   ParenthesizedStatement() { strictcount(int i | exists(this.getChild(i))) = 1 }
 

ql/src/codeql_ruby/controlflow/internal/Completion.qll

@@ -118,7 +118,7 @@ private predicate mustHaveBooleanCompletion(AstNode n) {
 private predicate inBooleanContext(AstNode n) {
   n = any(IfElsifAstNode parent).getConditionNode()
   or
-  n = any(While parent).getCondition()
+  n = any(ConditionalLoopAstNode parent).getCondition()
   or
   exists(LogicalAndAstNode parent |
     n = parent.getLeft()

ql/src/codeql_ruby/controlflow/internal/ControlFlowGraphImpl.qll

@@ -565,12 +565,12 @@ private module Trees {
     final override AstNode getChildNode(int i) { result = this.getOperand() and i = 0 }
   }
 
-  private class WhileTree extends PreOrderTree, While {
+  private class ConditionalLoopTree extends PreOrderTree, ConditionalLoopAstNode {
     final override predicate propagatesAbnormal(AstNode child) { child = this.getCondition() }
 
     final override predicate last(AstNode last, Completion c) {
       last(this.getCondition(), last, c) and
-      c instanceof FalseCompletion
+      this.endLoop(c)
       or
       last(this.getBody(), last, c) and
       not c.continuesLoop() and
@@ -590,7 +590,7 @@ private module Trees {
       c instanceof SimpleCompletion
       or
       last(this.getCondition(), pred, c) and
-      c instanceof TrueCompletion and
+      this.continueLoop(c) and
       first(this.getBody(), succ)
       or
       last(this.getBody(), pred, c) and