From 3abd9a755ee79805b7738b2c72075448221dac40 Mon Sep 17 00:00:00 2001
From: Tamas Vajk <tamasvajk@github.com>
Date: Fri, 22 Nov 2024 16:22:39 +0100
Subject: [PATCH] Code quality improvements

---
 .../src/main/kotlin/entities/Expression.kt    | 15 ++---
 java/ql/lib/semmle/code/java/Expr.qll         | 60 +++++++++++++++++--
 2 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt b/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
index ab93906be72..8a92e204457 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/Expression.kt
@@ -1734,16 +1734,13 @@ private fun KotlinFileExtractor.extractWhen(
     tw.writeExprsKotlinType(id, type.kotlinResult.id)
     extractExprContext(id, locId, callable, exprParent.enclosingStmt)
 
-    /* OLD: KE1, Should we remove this `when_if` DB construct?
-    if (e.origin == IrStatementOrigin.IF) {
-        tw.writeWhen_if(id)
-    }
-     */
+    val subjectVariable = e.subjectVariable
+    val subjectExpression = e.subjectExpression
 
-    if (e.subjectVariable != null) {
-        extractVariableExpr(e.subjectVariable!!, callable, id, -1, exprParent.enclosingStmt)
-    } else if (e.subjectExpression != null) {
-        extractExpressionExpr(e.subjectExpression!!, callable, id, -1, exprParent.enclosingStmt)
+    if (subjectVariable != null) {
+        extractVariableExpr(subjectVariable, callable, id, -1, exprParent.enclosingStmt)
+    } else if (subjectExpression != null) {
+        extractExpressionExpr(subjectExpression, callable, id, -1, exprParent.enclosingStmt)
     }
 
     e.entries.forEachIndexed { i, b ->
diff --git a/java/ql/lib/semmle/code/java/Expr.qll b/java/ql/lib/semmle/code/java/Expr.qll
index 22d955c183d..7bf70d7f447 100644
--- a/java/ql/lib/semmle/code/java/Expr.qll
+++ b/java/ql/lib/semmle/code/java/Expr.qll
@@ -2571,10 +2571,28 @@ class WhenExpr extends Expr, StmtParent, @whenexpr {
   /** Holds if this was written as an `if` expression. */
   predicate isIf() { when_if(this) }
 
-  /** Gets the expression of this `when` expression, if any. */
+  /**
+   * Gets the expression of this `when` expression, if any; such as `foo()` in the below sample.
+   *
+   * ```
+   * when (foo()) {
+   *   1 -> ...
+   *   2 -> ...
+   * }
+   */
   Expr getExpr() { result.isNthChildOf(this, -1) }
 
-  /** Gets the local variable declaration of this `when` expression, if any. */
+  /**
+   * Gets the local variable declaration of this `when` expression, if any; such as
+   * `val x = foo()` in the below sample.
+   *
+   * ```
+   * when (val x = foo()) {
+   *   1 -> ...
+   *   2 -> ...
+   * }
+   * ```
+   */
   LocalVariableDeclExpr getAVariableDeclExpr() { result.isNthChildOf(this, -1) }
 }
 
@@ -2589,13 +2607,32 @@ class WhenBranch extends Stmt, @whenbranch {
     result = this.getCondition(0).(WhenBranchConditionWithExpression).getExpression()
   }
 
-  /** Gets the `i`th condition of this branch. */
+  /**
+   * Gets the `i`th condition of this branch. The first branch in the below sample has two conditions:
+   *
+   * ```
+   * when (foo()) {
+   *   1, !in 4..10 -> ...
+   *   3 -> ...
+   * }
+   * ```
+   */
   WhenBranchCondition getCondition(int i) { i <= 0 and result.isNthChildOf(this, i) }
 
   /** Gets a condition of this branch. */
   WhenBranchCondition getACondition() { result = this.getCondition(_) }
 
-  /** Gets the guard applicable to this branch, if any. */
+  /**
+   * Gets the guard applicable to this branch, if any. Guards are currently experimental Kotlin features.
+   * In the below sample, the first branch has a guard: `bar() == 42`.
+   *
+   * ```
+   * when (foo()) {
+   *   1 if bar() == 42 -> ...
+   *   else -> ..
+   * }
+   * ```
+   */
   Expr getGuard() { result.isNthChildOf(this, 2) }
 
   /** Gets the result of this branch. */
@@ -2615,7 +2652,20 @@ class WhenBranch extends Stmt, @whenbranch {
   override string getAPrimaryQlClass() { result = "WhenBranch" }
 }
 
-/** A Kotlin `when` branch condition. */
+/**
+ * A Kotlin `when` branch condition. Sample conditions are shown below:
+ *
+ * ```
+ * fun foo(): Number = ...
+ *
+ * when (foo()) {
+ *   1 -> ...
+ *   in 2..10 -> ...
+ *   is Int -> ...
+ *   !is Int -> ...
+ * }
+ * ```
+ */
 abstract class WhenBranchCondition extends Stmt, @whenbranchcondition { }
 
 /** A Kotlin `when` branch condition with an expression. */