Extract WhenBranch as Stmt

java/ql/lib/config/semmlecode.dbscheme

@@ -601,7 +601,7 @@ stmts(
   int bodydecl: @callable ref
 );
 
-@stmtparent = @callable | @stmt | @switchexpr | @whenbranch | @stmtexpr;
+@stmtparent = @callable | @stmt | @switchexpr | @whenexpr| @whenbranch | @stmtexpr;
 
 case @stmt.kind of
   0 = @block
@@ -630,6 +630,7 @@ case @stmt.kind of
 | 23 = @yieldstmt
 | 24 = @errorstmt
 | 25 = @anonymousclassdeclstmt
+| 26 = @whenbranch
 ;
 
 #keyset[parent,idx]
@@ -755,13 +756,6 @@ case @expr.kind of
 /** Holds if this `when` expression was written as an `if` expression. */
 when_if(unique int id: @whenexpr ref);
 
-#keyset[whenid,idx]
-when_branch(
-    int id: @whenbranch,
-    int whenid: @whenexpr ref,
-    int idx: int ref
-);
-
 /** Holds if this `when` branch was written as an `else` branch. */
 when_branch_else(unique int id: @whenbranch ref);
 

java/ql/lib/semmle/code/java/ControlFlowGraph.qll

@@ -88,16 +88,14 @@ class ControlFlowNode extends Top, @exprparent {
   /** Gets the statement containing this node, if any. */
   Stmt getEnclosingStmt() {
     result = this or
-    result = this.(Expr).getEnclosingStmt() or
-    result = this.(WhenBranch).getWhenExpr().getEnclosingStmt()
+    result = this.(Expr).getEnclosingStmt()
   }
 
   /** Gets the immediately enclosing callable whose body contains this node. */
   Callable getEnclosingCallable() {
     result = this or
     result = this.(Stmt).getEnclosingCallable() or
-    result = this.(Expr).getEnclosingCallable() or
-    result = this.(WhenBranch).getWhenExpr().getEnclosingCallable()
+    result = this.(Expr).getEnclosingCallable()
   }
 
   /** Gets an immediate successor of this node. */

java/ql/lib/semmle/code/java/Expr.qll

@@ -523,9 +523,7 @@ class AssignExpr extends Assignment, @assignexpr {
  * For example, `class X { val y = 1 }`
  */
 class KtInitializerAssignExpr extends AssignExpr {
-  KtInitializerAssignExpr() {
-    ktInitializerAssignment(this)
-  }
+  KtInitializerAssignExpr() { ktInitializerAssignment(this) }
 
   override string getAPrimaryQlClass() { result = "KtInitializerAssignExpr" }
 }
@@ -2331,7 +2329,7 @@ class Argument extends Expr {
 }
 
 /** A Kotlin `when` expression. */
-class WhenExpr extends Expr, @whenexpr {
+class WhenExpr extends Expr, StmtParent, @whenexpr {
   override string toString() { result = "when ..." }
 
   override string getHalsteadID() { result = "WhenExpr" }
@@ -2339,14 +2337,14 @@ class WhenExpr extends Expr, @whenexpr {
   override string getAPrimaryQlClass() { result = "WhenExpr" }
 
   /** Gets the `i`th branch. */
-  WhenBranch getBranch(int i) { when_branch(result, this, i) }
+  WhenBranch getBranch(int i) { result.isNthChildOf(this, i) }
 
   /** Holds if this was written as an `if` expression. */
   predicate isIf() { when_if(this) }
 }
 
 /** A Kotlin `when` branch. */
-class WhenBranch extends Top, @whenbranch {
+class WhenBranch extends Stmt, @whenbranch {
   /** Gets the condition of this branch. */
   Expr getCondition() { result.isNthChildOf(this, 0) }