[PowerAssert] Capture safe cast expressions

Previously the only type operations captured by power-assert where type
checks. This caused safe casts to not be included in the diagram, and
more alarmingly, ignored if they were the root expression of a
parameter. This resulted in incorrect behavior of power-assert
transformed code in certain situations.

To fix, capture safe cast expressions and include all other type
operations as a constant expression.

^KT-69401 Fixed

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/diagram/ExpressionTree.kt

@@ -183,9 +183,15 @@ fun buildTree(expression: IrExpression): Node? {
 
                 expression.acceptChildren(this, chainNode)
 
-                if (expression.operator in setOf(IrTypeOperator.INSTANCEOF, IrTypeOperator.NOT_INSTANCEOF)) {
-                    // Only include `is` and `!is` checks
-                    chainNode.addChild(ExpressionNode(expression))
+                when (expression.operator) {
+                    // Only include `is` and `!is` checks and `as?` casts in the diagram.
+                    IrTypeOperator.INSTANCEOF,
+                    IrTypeOperator.NOT_INSTANCEOF,
+                    IrTypeOperator.SAFE_CAST,
+                        -> chainNode.addChild(ExpressionNode(expression))
+
+                    // Do not diagram other type operations.
+                    else -> chainNode.addChild(ConstantNode(expression))
                 }
             }
 

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/diagram/IrDiagram.kt

@@ -189,7 +189,11 @@ private fun typeOperatorOffset(
     source: String,
 ): Int {
     return when (expression.operator) {
-        IrTypeOperator.INSTANCEOF, IrTypeOperator.NOT_INSTANCEOF -> binaryOperatorOffset(expression.argument, sourceRangeInfo, source)
+        IrTypeOperator.INSTANCEOF,
+        IrTypeOperator.NOT_INSTANCEOF,
+        IrTypeOperator.SAFE_CAST,
+            -> binaryOperatorOffset(expression.argument, sourceRangeInfo, source)
+
         else -> 0
     }
 }

plugins/power-assert/testData/codegen/cast/SafeCast.box.txt

@@ -0,0 +1,22 @@
+---
+assert((a as? String)?.length == 5)
+        | |            |      |
+        | |            |      false
+        | |            4
+        | test
+        test
+---
+---
+assert((a as? String)?.length == 5)
+        | |            |      |
+        | |            |      false
+        | |            null
+        | null
+        1
+---
+---
+requireNotNull(a as? String) { "" }
+               | |
+               | null
+               1
+---

plugins/power-assert/testData/codegen/cast/SafeCast.fir.kt.txt

@@ -0,0 +1,43 @@
+fun box(): String {
+  return test1(a = "test").plus(other = test1(a = 1)).plus(other = test2(a = 1))
+}
+
+fun test1(a: Any): String {
+  return expectThrowableMessage(block = local fun <anonymous>() {
+    { // BLOCK
+      val tmp_0: Any = a
+      val tmp_1: String? = tmp_0 as? String
+      val tmp_2: Int? = { // BLOCK
+        val tmp_3: String? = tmp_1
+        when {
+          EQEQ(arg0 = tmp_3, arg1 = null) -> null
+          else -> tmp_3.<get-length>()
+        }
+      }
+      val tmp_4: Boolean = EQEQ(arg0 = tmp_2, arg1 = 5)
+      assert(value = tmp_4, lazyMessage = local fun <anonymous>(): String {
+        return "\nassert((a as? String)?.length == 5)\n        | |            |      |" + "\n        | |            |      " + tmp_4 + "\n        | |            " + tmp_2 + "\n        | " + tmp_1 + "\n        " + tmp_0 + "\n"
+      }
+)
+    }
+  }
+)
+}
+
+fun test2(a: Any): String {
+  return expectThrowableMessage(block = local fun <anonymous>() {
+    { // BLOCK
+      val tmp_5: Any = a
+      val tmp_6: String? = tmp_5 as? String
+      requireNotNull<String>(value = tmp_6, lazyMessage = local fun <anonymous>(): String {
+        return local fun <anonymous>(): Any {
+          return ""
+        }
+.invoke() + "\nrequireNotNull(a as? String) { \"\" }\n               | |" + "\n               | " + tmp_6 + "\n               " + tmp_5 + "\n"
+      }
+)
+    } /*~> Unit */
+  }
+)
+}
+

plugins/power-assert/testData/codegen/cast/SafeCast.kt

@@ -0,0 +1,17 @@
+// FUNCTION: kotlin.assert
+// FUNCTION: kotlin.requireNotNull
+// DUMP_KT_IR
+
+fun box(): String {
+    return test1("test") +
+            test1(1) +
+            test2(1)
+}
+
+fun test1(a: Any) = expectThrowableMessage {
+    assert((a as? String)?.length == 5)
+}
+
+fun test2(a: Any) = expectThrowableMessage {
+    requireNotNull(a as? String) { "" }
+}

plugins/power-assert/testData/codegen/cast/SafeCast.kt.txt

@@ -0,0 +1,43 @@
+fun box(): String {
+  return test1(a = "test").plus(other = test1(a = 1)).plus(other = test2(a = 1))
+}
+
+fun test1(a: Any): String {
+  return expectThrowableMessage(block = local fun <anonymous>() {
+    { // BLOCK
+      val tmp_0: Any = a
+      val tmp_1: String? = tmp_0 as? String
+      val tmp_2: Int? = { // BLOCK
+        val tmp_3: String? = tmp_1
+        when {
+          EQEQ(arg0 = tmp_3, arg1 = null) -> null
+          else -> tmp_3.<get-length>()
+        }
+      }
+      val tmp_4: Boolean = EQEQ(arg0 = tmp_2, arg1 = 5)
+      assert(value = tmp_4, lazyMessage = local fun <anonymous>(): String {
+        return "\nassert((a as? String)?.length == 5)\n        | |            |      |" + "\n        | |            |      " + tmp_4 + "\n        | |            " + tmp_2 + "\n        | " + tmp_1 + "\n        " + tmp_0 + "\n"
+      }
+)
+    }
+  }
+)
+}
+
+fun test2(a: Any): String {
+  return expectThrowableMessage(block = local fun <anonymous>() {
+    { // BLOCK
+      val tmp_5: Any = a
+      val tmp_6: String? = tmp_5 as? String
+      requireNotNull<String>(value = tmp_6, lazyMessage = local fun <anonymous>(): String {
+        return local fun <anonymous>(): Any {
+          return ""
+        }
+.invoke() + "\nrequireNotNull(a as? String) { \"\" }\n               | |" + "\n               | " + tmp_6 + "\n               " + tmp_5 + "\n"
+      }
+)
+    } /*~> Unit */
+  }
+)
+}
+