Visit IR tree to transform and gather all assertion values

kotlin-power-assert/build.gradle.kts

@@ -13,6 +13,7 @@ dependencies {
   compileOnly("org.jetbrains.kotlin:kotlin-compiler-embeddable:1.3.60")
 
   testImplementation(kotlin("test-junit"))
+  testImplementation("org.jetbrains.kotlin:kotlin-compiler-embeddable:1.3.60")
   testImplementation("com.github.tschuchortdev:kotlin-compile-testing:1.2.5")
 }
 

kotlin-power-assert/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt

@@ -17,7 +17,8 @@
 package com.bnorm.power
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.ir.asSimpleLambda
+import org.jetbrains.kotlin.backend.common.ir.inline
 import org.jetbrains.kotlin.backend.common.lower.at
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irIfThen
@@ -29,19 +30,29 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.SourceRangeInfo
 import org.jetbrains.kotlin.ir.builders.irBlock
 import org.jetbrains.kotlin.ir.builders.irCall
+import org.jetbrains.kotlin.ir.builders.irCallOp
 import org.jetbrains.kotlin.ir.builders.irConcat
 import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irString
 import org.jetbrains.kotlin.ir.builders.irTemporary
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrGetValue
+import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.IrStringConcatenationImpl
+import org.jetbrains.kotlin.ir.types.getClass
+import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.getPackageFragment
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
+import org.jetbrains.kotlin.util.OperatorNameConventions
 import java.io.File
 
 fun FileLoweringPass.runOnFileInOrder(irFile: IrFile) {
@@ -59,7 +70,7 @@ fun FileLoweringPass.runOnFileInOrder(irFile: IrFile) {
 
 class PowerAssertCallTransformer(
   private val context: JvmBackendContext
-) : IrElementTransformerVoidWithContext(), FileLoweringPass {
+) : IrElementTransformerVoid(), FileLoweringPass {
   private lateinit var file: IrFile
   private lateinit var fileSource: String
 
@@ -72,6 +83,12 @@ class PowerAssertCallTransformer(
     irFile.transformChildrenVoid()
   }
 
+  private data class IrTemporaryVariable(
+    val variable: IrVariable,
+    val indentation: Int,
+    val source: String
+  )
+
   override fun visitCall(expression: IrCall): IrExpression {
     val function = expression.symbol.owner
     if (!function.isAssert)
@@ -84,37 +101,129 @@ class PowerAssertCallTransformer(
       at(expression)
 
       return irBlock {
-        val assertCondition = expression.getValueArgument(0)!!
+        val stack = mutableListOf<IrTemporaryVariable>()
-//        val argumentSource = fileSource.substring(assertCondition.startOffset, assertCondition.endOffset)
-        val indent = file.info(assertCondition).startColumnNumber - callIndent
 
-//        println(buildString {
+        fun push(expression: IrExpression): IrGetValue {
+          val variable = irTemporary(expression)
+          val source = fileSource.substring(expression.startOffset, expression.endOffset)
+
+          var indentation = file.info(expression).startColumnNumber - callIndent
+          if (expression is IrMemberAccessExpression) {
+            // TODO Is this the best way to fix indentation of infix operators?
+            indentation += when (expression.origin) {
+              IrStatementOrigin.EQEQ, IrStatementOrigin.EQEQEQ -> source.indexOf("==")
+              IrStatementOrigin.EXCLEQ, IrStatementOrigin.EXCLEQEQ -> source.indexOf("!=")
+              IrStatementOrigin.LT -> source.indexOf("<") // TODO What about generics?
+              IrStatementOrigin.GT -> source.indexOf(">") // TODO What about generics?
+              IrStatementOrigin.LTEQ -> source.indexOf("<=")
+              IrStatementOrigin.GTEQ -> source.indexOf(">=")
+              else -> 0
+            }
+          }
+
+          stack.add(IrTemporaryVariable(variable, indentation, source))
+          return irGet(variable)
+        }
+
+        val assertCondition = expression.getValueArgument(0)!!.transform(object : IrElementTransformerVoid() {
+          override fun visitExpression(expression: IrExpression): IrExpression {
+            return when (val transformed = super.visitExpression(expression)) {
+              is IrGetValue -> push(transformed)
+              is IrCall -> push(transformed)
+              // TODO what else needs to get pushed in the stack?
+              else -> transformed
+            }
+          }
+        }, null)
+        require(assertCondition is IrGetValue)
+
+//        print(buildString {
 //          append(callSource).newline()
-//          indent(indent).append("|").newline()
+//          val sorted = stack.sortedBy { it.indentation }
-//          indent(indent).append(argumentSource)
+//
+//          val indentations = sorted.map { it.indentation }
+//          var last = -1
+//          for (i in indentations) {
+//            if (i > last) {
+//              indent(i - last - 1).append("|")
+//            }
+//            last = i
+//          }
+//          newline()
+//
+//          for (tmp in sorted.asReversed()) {
+//
+//            last = -1
+//            for (i in indentations) {
+//              if (i == tmp.indentation) break
+//              if (i > last) {
+//                indent(i - last - 1).append("|")
+//              }
+//              last = i
+//            }
+//
+//            indent(tmp.indentation - last - 1)
+//            append(tmp.source).newline()
+//          }
 //        })
 
-        // TODO transform tree of expressions and create irTemporary for each
+        val lambdaArgument = if (function.valueParameters.size == 2) expression.getValueArgument(1) else null
-        val temp = irTemporary(assertCondition)
+        val lambda = lambdaArgument?.asSimpleLambda()
-        val throwError = irThrow(irCall(constructor).apply {
+        val invokeVar = if (lambda == null && lambdaArgument != null) irTemporary(lambdaArgument) else null
+
+        // Build assertion message
+        val throwError = irThrow(irCall(constructor).apply {
+          putValueArgument(0, irConcat().apply {
+
+            addArgument(
+              when {
+                lambda != null -> lambda.inline()
+                lambdaArgument != null -> {
+                  val invoke = lambdaArgument.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
+                  irCallOp(invoke.symbol, invoke.returnType, irGet(invokeVar!!))
+                }
+                else -> irString("Assertion failed")
+              }
+            )
+
+            val sorted = stack.sortedBy { it.indentation }
+            val indentations = sorted.map { it.indentation }
 
-          val message = irConcat().apply {
             addArgument(irString(buildString {
-              append("Assertion failed:").newline()
+              append(":").newline()
               append(callSource).newline()
-              indent(indent).append("|").newline()
+              var last = -1
-              indent(indent)
+              for (i in indentations) {
+                if (i > last) {
+                  indent(i - last - 1).append("|")
+                }
+                last = i
+              }
             }))
-            addArgument(irGet(temp))
+
-          }
+
-          putValueArgument(0, message)
+            for (tmp in sorted.asReversed()) {
+              addArgument(irString(buildString {
+                var last = -1
+                newline()
+                for (i in indentations) {
+                  if (i == tmp.indentation) break
+                  if (i > last) {
+                    indent(i - last - 1).append("|")
+                  }
+                  last = i
+                }
+                indent(tmp.indentation - last - 1)
+              }))
+              addArgument(irGet(tmp.variable))
+            }
+          })
         })
-        +irIfThen(irNot(irGet(temp)), throwError)
+
+        +irIfThen(irNot(assertCondition), throwError)
       }
     }
   }
-
-
 }
 
 val IrFunction.isAssert: Boolean

kotlin-power-assert/src/main/kotlin/com/bnorm/power/PowerAssertComponentRegistrar.kt

@@ -20,9 +20,6 @@ import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.com.intellij.mock.MockProject
 import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
-import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.SourceRangeInfo
-import org.jetbrains.kotlin.ir.declarations.IrFile
 
 class PowerAssertComponentRegistrar : ComponentRegistrar {
   override fun registerProjectComponents(

kotlin-power-assert/src/main/resources/META-INF/services/org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar

@@ -0,0 +1 @@
+com.bnorm.power.PowerAssertComponentRegistrar

kotlin-power-assert/src/test/kotlin/com/bnorm/power/test.kt

@@ -26,22 +26,61 @@ import kotlin.test.fail
 
 class CompilerTest {
   @Test
-  fun testMyCompilerPlugin() {
+  fun memberFunctions() {
     assertMessage(
       """
 fun main() {
-  val hello = "Brian"
+  val hello = "Hello"
-  assert(hello == "World")
+  assert(hello.length == "World".substring(1, 4).length)
 }""",
       """
 Assertion failed:
-assert(hello == "World")
+assert(hello.length == "World".substring(1, 4).length)
-       |
+       |     |      |          |               |
-       false
+       |     |      |          |               3
+       |     |      |          orl
+       |     |      false
+       |     5
+       Hello
 """.trimIndent()
     )
   }
 
+  @Test
+  fun transformations() {
+    assertMessage(
+      """
+fun main() {
+  val hello = listOf("Hello", "World")
+  assert(hello.reversed() == emptyList<String>())
+}""",
+      """
+Assertion failed:
+assert(hello.reversed() == emptyList<String>())
+       |     |          |  |
+       |     |          |  []
+       |     |          false
+       |     [World, Hello]
+       [Hello, World]
+""".trimIndent()
+    )
+  }
+
+  @Test
+  fun customMessage() {
+    assertMessage(
+      """
+fun main() {
+  assert(1 == 2) { "Not equal" }
+}""",
+      """
+Not equal:
+assert(1 == 2) { "Not equal" }
+         |
+         false
+""".trimIndent()
+    )
+  }
 }
 
 fun assertMessage(@Language("kotlin") source: String, message: String) {