Handle operator functions better (like arithmetic)

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/AssertTree.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationParent
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrConst
+import org.jetbrains.kotlin.ir.expressions.IrContainerExpression
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.IrWhen
@@ -91,6 +92,13 @@ class RootNode : Node() {
 fun buildAssertTree(expression: IrExpression): RootNode {
   val tree = RootNode()
   expression.accept(object : IrElementVisitor<Unit, Node> {
+    val INCREMENT_DECREMENT_OPERATORS = setOf(
+      IrStatementOrigin.PREFIX_INCR,
+      IrStatementOrigin.PREFIX_DECR,
+      IrStatementOrigin.POSTFIX_INCR,
+      IrStatementOrigin.POSTFIX_DECR
+    )
+
     override fun visitElement(element: IrElement, data: Node) {
       element.acceptChildren(this, data)
     }
@@ -101,6 +109,16 @@ fun buildAssertTree(expression: IrExpression): RootNode {
       expression.acceptChildren(this, node)
     }
 
+    override fun visitContainerExpression(expression: IrContainerExpression, data: Node) {
+      if (expression.origin in INCREMENT_DECREMENT_OPERATORS) {
+        val node = data as? ExpressionNode ?: ExpressionNode(data)
+        node.add(expression)
+        return // Skip the internals of increment/decrement operations
+      }
+
+      super.visitContainerExpression(expression, data)
+    }
+
     override fun visitCall(expression: IrCall, data: Node) {
       if (expression.symbol.owner.name.asString() == "EQEQ" && expression.origin == IrStatementOrigin.EXCLEQ) {
         // Skip the EQEQ part of a EXCLEQ call

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/IrStackVariable.kt

@@ -17,13 +17,14 @@
 package com.bnorm.power
 
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.SourceRangeInfo
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 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.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
@@ -50,51 +51,13 @@ fun IrBuilderWithScope.buildMessage(
 
   val originalInfo = file.info(expression)
   val callIndent = originalInfo.startColumnNumber
-  val callSource = fileSource.substring(expression)
-    .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
 
-  val stackValues = stack.map { (temporary, original) ->
+  val stackValues = stack.map { it.toValueDisplay(fileSource, callIndent, file, originalInfo) }
-    val source = fileSource.substring(original)
-      .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
-
-    val info = file.info(original)
-    var indent = info.startColumnNumber - callIndent
-    var row = info.startLineNumber - originalInfo.startLineNumber
-
-    val columnOffset: Int = when (original) {
-      is IrMemberAccessExpression<*> -> {
-        val owner = original.symbol.owner
-        when {
-          owner is IrSimpleFunction && owner.isInfix -> source.indexOf(owner.name.asString())
-          else -> when (original.origin) {
-            // TODO handle equality and comparison better?
-            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
-          }
-        }
-      }
-      else -> 0
-    }
-
-    val prefix = source.substring(0, columnOffset)
-    val rowShift = prefix.count { it == '\n' }
-    if (rowShift == 0) {
-      indent += columnOffset
-    } else {
-      row += rowShift
-      indent = columnOffset - (prefix.lastIndexOf('\n') + 1)
-    }
-
-    ValueDisplay(temporary, indent, row, source)
-  }
 
   val valuesByRow = stackValues.groupBy { it.row }
-  val rows = callSource.split("\n")
+  val rows = fileSource.substring(expression)
+    .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
+    .split("\n")
 
   return irConcat().apply {
     addArgument(title)
@@ -103,42 +66,130 @@ fun IrBuilderWithScope.buildMessage(
       val rowValues = valuesByRow[row]?.let { values -> values.sortedBy { it.indent } } ?: emptyList()
       val indentations = rowValues.map { it.indent }
 
-      addArgument(irString(buildString {
+      addArgument(irString {
-        newline()
+        appendLine()
         append(rowSource)
-        if (indentations.isNotEmpty()) newline()
+        if (indentations.isNotEmpty()) {
-
+          appendLine()
-        var last = -1
+          var last = -1
-        for (i in indentations) {
+          for (i in indentations) {
-          if (i > last) {
+            if (i > last) indent(i - last - 1).append("|")
-            indent(i - last - 1).append("|")
+            last = i
           }
-          last = i
         }
-      }))
+      })
-
 
       for (tmp in rowValues.asReversed()) {
-        addArgument(irString(buildString {
+        addArgument(irString {
+          appendLine()
           var last = -1
-          newline()
           for (i in indentations) {
             if (i == tmp.indent) break
-            if (i > last) {
+            if (i > last) indent(i - last - 1).append("|")
-              indent(i - last - 1).append("|")
-            }
             last = i
           }
           indent(tmp.indent - last - 1)
-        }))
+        })
         addArgument(irGet(tmp.value))
       }
     }
   }
 }
 
+private fun IrStackVariable.toValueDisplay(
+  fileSource: String,
+  callIndent: Int,
+  file: IrFile,
+  originalInfo: SourceRangeInfo
+): ValueDisplay {
+  val source = fileSource.substring(original)
+    .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
+
+  val info = file.info(original)
+  var indent = info.startColumnNumber - callIndent
+  var row = info.startLineNumber - originalInfo.startLineNumber
+
+  val columnOffset = findDisplayOffset(original, source)
+
+  val prefix = source.substring(0, columnOffset)
+  val rowShift = prefix.count { it == '\n' }
+  if (rowShift == 0) {
+    indent += columnOffset
+  } else {
+    row += rowShift
+    indent = columnOffset - (prefix.lastIndexOf('\n') + 1)
+  }
+
+  return ValueDisplay(temporary, indent, row, source)
+}
+
+/**
+ * Responsible for determining the diagram display offset of the expression
+ * beginning from the startOffset of the expression.
+ *
+ * Equality:
+ * ```
+ * number == 42
+ * | <- startOffset
+ *        | <- display offset: 7
+ * ```
+ *
+ * Arithmetic:
+ * ```
+ * i + 2
+ * | <- startOffset
+ *   | <- display offset: 2
+ * ```
+ *
+ * Infix:
+ * ```
+ * 1 shl 2
+ * | <- startOffset
+ *   | <- display offset: 2
+ * ```
+ *
+ * Standard:
+ * ```
+ * 1.shl(2)
+ *   | <- startOffset
+ *   | <- display offset: 0
+ * ```
+ */
+private fun findDisplayOffset(
+  expression: IrExpression,
+  source: String
+): Int {
+  if (expression !is IrMemberAccessExpression<*>) return 0
+
+  if (expression.origin == IrStatementOrigin.EXCLEQ || expression.origin == IrStatementOrigin.EXCLEQEQ) {
+    // special case to handle `value != null`
+    return source.indexOf("!=")
+  }
+
+  val owner = expression.symbol.owner
+  if (owner !is IrSimpleFunction) return 0
+
+  if (owner.isInfix || owner.isOperator || owner.origin == IrBuiltIns.BUILTIN_OPERATOR) {
+    // Ignore single value operators
+    if (expression.dispatchReceiver != null && expression.valueArgumentsCount == 0) return 0
+
+    // Start after the dispatcher or first argument
+    val dispatcher = expression.dispatchReceiver ?: expression.getValueArgument(0) ?: return 0
+    var offset = dispatcher.endOffset - expression.startOffset
+    if (offset < 0) return 0 // infix function called using non-infix syntax
+
+    // Continue until there is a non-whitespace character
+    while (source[offset].isWhitespace()) offset++
+    return offset
+  }
+
+  return 0
+}
+
 fun String.substring(expression: IrElement) = substring(expression.startOffset, expression.endOffset)
 fun IrFile.info(expression: IrElement) = fileEntry.getSourceRangeInfo(expression.startOffset, expression.endOffset)
 
-fun StringBuilder.indent(indentation: Int): StringBuilder = append(" ".repeat(indentation))
+fun StringBuilder.indent(indentation: Int): StringBuilder {
-fun StringBuilder.newline(): StringBuilder = append("\n")
+  repeat(indentation) { append(" ") }
+  return this
+}

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/PowerAssertGenerator.kt

@@ -25,8 +25,8 @@ import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irTemporary
 import org.jetbrains.kotlin.ir.builders.parent
 import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrGetValue
 import org.jetbrains.kotlin.ir.expressions.IrWhen
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 
 abstract class PowerAssertGenerator {
@@ -87,15 +87,15 @@ abstract class PowerAssertGenerator {
     private val stack: MutableList<IrStackVariable>,
     private val transform: List<IrExpression>
   ) : IrElementTransformerVoid() {
-    private fun push(expression: IrExpression): IrGetValue = with(builder) {
-      val variable = irTemporary(expression)
-      stack.add(IrStackVariable(variable, expression))
-      irGet(variable)
-    }
 
     override fun visitExpression(expression: IrExpression): IrExpression {
       return if (expression in transform) {
-        push(super.visitExpression(expression))
+        with(builder) {
+          val copy = expression.deepCopyWithSymbols(scope.getLocalDeclarationParent())
+          val variable = irTemporary(super.visitExpression(expression))
+          stack.add(IrStackVariable(variable, copy))
+          irGet(variable)
+        }
       } else {
         super.visitExpression(expression)
       }

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/irUtils.kt

@@ -0,0 +1,7 @@
+package com.bnorm.power
+
+import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
+import org.jetbrains.kotlin.ir.builders.irString
+
+fun IrBuilderWithScope.irString(builderAction: StringBuilder.() -> Unit) =
+  irString(buildString { builderAction() })

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

@@ -0,0 +1,146 @@
+package com.bnorm.power
+
+import org.junit.Test
+import kotlin.test.assertEquals
+
+class ArithmeticExpressionTest {
+  @Test
+  fun `assertion with inline addition`() {
+    val actual = executeMainAssertion("assert(1 + 1 == 4)")
+    assertEquals(
+      """
+      Assertion failed
+      assert(1 + 1 == 4)
+               |   |
+               |   false
+               2
+      """.trimIndent(),
+      actual
+    )
+  }
+
+  @Test
+  fun `assertion with inline subtraction`() {
+    val actual = executeMainAssertion("assert(3 - 1 == 4)")
+    assertEquals(
+      """
+      Assertion failed
+      assert(3 - 1 == 4)
+               |   |
+               |   false
+               2
+      """.trimIndent(),
+      actual
+    )
+  }
+
+  @Test
+  fun `assertion with inline multiplication`() {
+    val actual = executeMainAssertion("assert(1 * 2 == 4)")
+    assertEquals(
+      """
+      Assertion failed
+      assert(1 * 2 == 4)
+               |   |
+               |   false
+               2
+      """.trimIndent(),
+      actual
+    )
+  }
+
+  @Test
+  fun `assertion with inline division`() {
+    val actual = executeMainAssertion("assert(2 / 1 == 4)")
+    assertEquals(
+      """
+      Assertion failed
+      assert(2 / 1 == 4)
+               |   |
+               |   false
+               2
+      """.trimIndent(),
+      actual
+    )
+  }
+
+  @Test
+  fun `assertion with inline prefix increment`() {
+    val actual = executeMainAssertion(
+      """
+      var i = 1
+      assert(++i == 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Assertion failed
+      assert(++i == 4)
+             |   |
+             |   false
+             2
+      """.trimIndent(),
+      actual
+    )
+  }
+
+  @Test
+  fun `assertion with inline postfix increment`() {
+    val actual = executeMainAssertion(
+      """
+      var i = 1
+      assert(i++ == 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Assertion failed
+      assert(i++ == 4)
+             |   |
+             |   false
+             1
+      """.trimIndent(),
+      actual
+    )
+  }
+
+  @Test
+  fun `assertion with inline prefix decrement`() {
+    val actual = executeMainAssertion(
+      """
+      var i = 3
+      assert(--i == 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Assertion failed
+      assert(--i == 4)
+             |   |
+             |   false
+             2
+      """.trimIndent(),
+      actual
+    )
+  }
+
+  @Test
+  fun `assertion with inline postfix decrement`() {
+    val actual = executeMainAssertion(
+      """
+      var i = 3
+      assert(i-- == 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Assertion failed
+      assert(i-- == 4)
+             |   |
+             |   false
+             3
+      """.trimIndent(),
+      actual
+    )
+  }
+}

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

@@ -0,0 +1,68 @@
+package com.bnorm.power
+
+import com.tschuchort.compiletesting.KotlinCompilation
+import com.tschuchort.compiletesting.SourceFile
+import org.intellij.lang.annotations.Language
+import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
+import org.jetbrains.kotlin.name.FqName
+import java.lang.reflect.InvocationTargetException
+import kotlin.test.assertEquals
+import kotlin.test.fail
+
+val DEFAULT_COMPONENT_REGISTRARS = arrayOf(
+  PowerAssertComponentRegistrar(setOf(FqName("kotlin.assert"))),
+)
+
+fun compile(
+  list: List<SourceFile>,
+  vararg plugins: ComponentRegistrar = DEFAULT_COMPONENT_REGISTRARS,
+): KotlinCompilation.Result {
+  return KotlinCompilation().apply {
+    sources = list
+    useIR = true
+    messageOutputStream = System.out
+    compilerPlugins = plugins.toList()
+    inheritClassPath = true
+  }.compile()
+}
+
+fun executeAssertion(
+  @Language("kotlin") source: String,
+  vararg plugins: ComponentRegistrar = DEFAULT_COMPONENT_REGISTRARS
+): String {
+  val result = compile(
+    listOf(SourceFile.kotlin("main.kt", source, trimIndent = false)),
+    *plugins,
+  )
+  assertEquals(KotlinCompilation.ExitCode.OK, result.exitCode)
+
+  val kClazz = result.classLoader.loadClass("MainKt")
+  val main = kClazz.declaredMethods.single { it.name == "main" && it.parameterCount == 0 }
+  try {
+    try {
+      main.invoke(null)
+    } catch (t: InvocationTargetException) {
+      throw t.cause!!
+    }
+    fail("should have thrown assertion")
+  } catch (t: Throwable) {
+    return t.message!!
+  }
+}
+
+fun executeMainAssertion(mainBody: String) = executeAssertion(
+  """
+fun main() {
+  $mainBody
+}
+"""
+)
+
+fun assertMessage(
+  @Language("kotlin") source: String,
+  message: String,
+  vararg plugins: ComponentRegistrar = DEFAULT_COMPONENT_REGISTRARS
+) {
+  val actual = executeAssertion(source, *plugins)
+  assertEquals(message, actual)
+}

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

@@ -276,6 +276,20 @@ assert(text?.length?.minus(2) == 1)
     assertMessage(
       """
 fun main() {
+  assert(1.shl(1) == 4)
+}""",
+      """
+Assertion failed
+assert(1.shl(1) == 4)
+         |      |
+         |      false
+         2
+""".trimIndent()
+    )
+
+    assertMessage(
+      """
+fun main() {
   assert(1 shl 1 == 4)
 }""",
       """
@@ -414,32 +428,3 @@ Assertion failed
     )
   }
 }
-
-fun assertMessage(
-  @Language("kotlin") source: String,
-  message: String,
-  vararg plugins: ComponentRegistrar = arrayOf(PowerAssertComponentRegistrar(setOf(FqName("kotlin.assert"))))
-) {
-  val result = KotlinCompilation().apply {
-    sources = listOf(SourceFile.kotlin("main.kt", source, trimIndent = false))
-    useIR = true
-    messageOutputStream = System.out
-    compilerPlugins = plugins.toList()
-    inheritClassPath = true
-  }.compile()
-
-  assertEquals(KotlinCompilation.ExitCode.OK, result.exitCode)
-
-  val kClazz = result.classLoader.loadClass("MainKt")
-  val main = kClazz.declaredMethods.single { it.name == "main" && it.parameterCount == 0 }
-  try {
-    try {
-      main.invoke(null)
-    } catch (t: InvocationTargetException) {
-      throw t.cause!!
-    }
-    fail("should have thrown assertion")
-  } catch (t: Throwable) {
-    assertEquals(message, t.message)
-  }
-}