Ruby: Generalize CfgNodes::ChildMapping

ruby/ql/lib/codeql/ruby/controlflow/CfgNodes.qll

@@ -149,7 +149,7 @@ private AstNode desugar(AstNode n) {
 /**
  * A class for mapping parent-child AST nodes to parent-child CFG nodes.
  */
-abstract private class ExprChildMapping extends Expr {
+abstract private class ChildMapping extends AstNode {
   /**
    * Holds if `child` is a (possibly nested) child of this expression
    * for which we would like to find a matching CFG child.
@@ -157,17 +157,7 @@ abstract private class ExprChildMapping extends Expr {
   abstract predicate relevantChild(AstNode child);
 
   pragma[nomagic]
-  private predicate reachesBasicBlock(AstNode child, CfgNode cfn, BasicBlock bb) {
-    this.relevantChild(child) and
-    cfn = this.getAControlFlowNode() and
-    bb.getANode() = cfn
-    or
-    exists(BasicBlock mid |
-      this.reachesBasicBlock(child, cfn, mid) and
-      bb = mid.getAPredecessor() and
-      not mid.getANode().getNode() = child
-    )
-  }
+  abstract predicate reachesBasicBlock(AstNode child, CfgNode cfn, BasicBlock bb);
 
   /**
    * Holds if there is a control-flow path from `cfn` to `cfnChild`, where `cfn`
@@ -183,6 +173,44 @@ abstract private class ExprChildMapping extends Expr {
   }
 }
 
+/**
+ * A class for mapping parent-child AST nodes to parent-child CFG nodes.
+ */
+abstract private class ExprChildMapping extends Expr, ChildMapping {
+  pragma[nomagic]
+  override predicate reachesBasicBlock(AstNode child, CfgNode cfn, BasicBlock bb) {
+    this.relevantChild(child) and
+    cfn = this.getAControlFlowNode() and
+    bb.getANode() = cfn
+    or
+    exists(BasicBlock mid |
+      this.reachesBasicBlock(child, cfn, mid) and
+      bb = mid.getAPredecessor() and
+      not mid.getANode().getNode() = child
+    )
+  }
+}
+
+/**
+ * A class for mapping parent-child AST nodes to parent-child CFG nodes.
+ */
+abstract private class NonExprChildMapping extends ChildMapping {
+  NonExprChildMapping() { not this instanceof Expr }
+
+  pragma[nomagic]
+  override predicate reachesBasicBlock(AstNode child, CfgNode cfn, BasicBlock bb) {
+    this.relevantChild(child) and
+    cfn.getNode() = this and
+    bb.getANode() = cfn
+    or
+    exists(BasicBlock mid |
+      this.reachesBasicBlock(child, cfn, mid) and
+      bb = mid.getASuccessor() and
+      not mid.getANode().getNode() = child
+    )
+  }
+}
+
 /** Provides classes for control-flow nodes that wrap AST expressions. */
 module ExprNodes {
   private class LiteralChildMapping extends ExprChildMapping, Literal {
@@ -346,10 +374,6 @@ module ExprNodes {
     final ExprCfgNode getBlock() { e.hasCfgChild(e.(MethodCall).getBlock(), this, result) }
   }
 
-  private class CaseExprChildMapping extends ExprChildMapping, CaseExpr {
-    override predicate relevantChild(AstNode e) { e = this.getValue() }
-  }
-
   /** A control-flow node that wraps a `MethodCall` AST expression. */
   class MethodCallCfgNode extends CallCfgNode {
     MethodCallCfgNode() { super.getExpr() instanceof MethodCall }
@@ -357,6 +381,10 @@ module ExprNodes {
     override MethodCall getExpr() { result = super.getExpr() }
   }
 
+  private class CaseExprChildMapping extends ExprChildMapping, CaseExpr {
+    override predicate relevantChild(AstNode e) { e = this.getValue() or e = this.getABranch() }
+  }
+
   /** A control-flow node that wraps a `CaseExpr` AST expression. */
   class CaseExprCfgNode extends ExprCfgNode {
     override CaseExprChildMapping e;
@@ -365,6 +393,47 @@ module ExprNodes {
 
     /** Gets the expression being compared, if any. */
     final ExprCfgNode getValue() { e.hasCfgChild(e.getValue(), this, result) }
+
+    /**
+     * Gets the `n`th branch of this case expression, either a `when` clause, an `in` clause, or an `else` branch.
+     */
+    final AstCfgNode getBranch(int n) { e.hasCfgChild(e.getBranch(n), this, result) }
+  }
+
+  private class InClauseChildMapping extends NonExprChildMapping, InClause {
+    override predicate relevantChild(AstNode e) {
+      e = this.getPattern() or
+      e = this.getCondition() or
+      e = this.getBody()
+    }
+  }
+
+  /** A control-flow node that wraps an `InClause` AST expression. */
+  class InClauseCfgNode extends AstCfgNode {
+    private InClauseChildMapping e;
+
+    /**
+     * Gets the pattern in this `in`-clause.
+     */
+    final AstCfgNode getPattern() { e.hasCfgChild(e.getPattern(), this, result) }
+
+    /** Gets the pattern guard condition in this `in` clause, if any. */
+    final ExprCfgNode getCondition() { e.hasCfgChild(e.getCondition(), this, result) }
+
+    /** Gets the body of this `in`-clause. */
+    final ExprCfgNode getBody() { e.hasCfgChild(e.getBody(), this, result) }
+  }
+
+  private class WhenClauseChildMapping extends NonExprChildMapping, WhenClause {
+    override predicate relevantChild(AstNode e) { e = this.getBody() }
+  }
+
+  /** A control-flow node that wraps a `WhenClause` AST expression. */
+  class WhenClauseCfgNode extends AstCfgNode {
+    private WhenClauseChildMapping e;
+
+    /** Gets the body of this `when`-clause. */
+    final ExprCfgNode getBody() { e.hasCfgChild(e.getBody(), this, result) }
   }
 
   private class ConditionalExprChildMapping extends ExprChildMapping, ConditionalExpr {

ruby/ql/lib/codeql/ruby/dataflow/internal/DataFlowPrivate.qll

@@ -126,15 +126,14 @@ module LocalFlow {
     or
     nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::ConditionalExprCfgNode).getBranch(_)
     or
-    exists(CfgNode n, Stmt stmt, CaseExpr c |
-      c = nodeTo.asExpr().getExpr() and
-      n = nodeFrom.asExpr() and
-      n = nodeTo.asExpr().getAPredecessor() and
-      stmt = n.getNode()
+    exists(CfgNodes::AstCfgNode branch |
+      branch = nodeTo.asExpr().(CfgNodes::ExprNodes::CaseExprCfgNode).getBranch(_)
     |
-      stmt = c.getElseBranch() or
-      stmt = c.getABranch().(InClause).getBody() or
-      stmt = c.getABranch().(WhenClause).getBody()
+      nodeFrom.asExpr() = branch.(CfgNodes::ExprNodes::InClauseCfgNode).getBody()
+      or
+      nodeFrom.asExpr() = branch.(CfgNodes::ExprNodes::WhenClauseCfgNode).getBody()
+      or
+      nodeFrom.asExpr() = branch and branch instanceof CfgNodes::ExprCfgNode
     )
     or
     exists(CfgNodes::ExprCfgNode exprTo, ReturningStatementNode n |