Support multi-line assertion expressions

Fixes #3

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

@@ -17,19 +17,30 @@
 package com.bnorm.power
 
 import org.jetbrains.kotlin.backend.common.lower.irThrow
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irCall
 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.IrVariable
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.IrThrow
 import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
 
 data class IrStackVariable(
-  val variable: IrVariable,
+  val temporary: IrVariable,
-  val startColumnNumber: Int,
+  val original: IrExpression
+)
+
+data class ValueDisplay(
+  val value: IrVariable,
+  val indent: Int,
+  val row: Int,
   val source: String
 )
 
@@ -41,45 +52,104 @@ fun IrBuilderWithScope.buildThrow(
 })
 
 fun IrBuilderWithScope.buildMessage(
+  file: IrFile,
+  fileSource: String,
   title: IrExpression,
-  stack: List<IrStackVariable>,
+  expression: IrExpression,
-  callSource: String,
+  stack: List<IrStackVariable>
-  callIndent: Int = 0
 ): IrExpression {
-  val stack = stack.map { it.copy(startColumnNumber = it.startColumnNumber - callIndent) }
+
+  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 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 descriptor = original.descriptor
+        when {
+          descriptor is FunctionDescriptor && descriptor.isInfix -> source.indexOf(descriptor.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")
+
   return irConcat().apply {
     addArgument(title)
 
-    val sorted = stack.sortedBy { it.startColumnNumber }
+    for ((row, rowSource) in rows.withIndex()) {
-    val indentations = sorted.map { it.startColumnNumber }
+      val rowValues = valuesByRow[row]?.let { values -> values.sortedBy { it.indent } } ?: emptyList()
+      val indentations = rowValues.map { it.indent }
 
-    addArgument(irString(buildString {
-      newline()
-      append(callSource).newline()
-      var last = -1
-      for (i in indentations) {
-        if (i > last) {
-          indent(i - last - 1).append("|")
-        }
-        last = i
-      }
-    }))
-
-
-    for (tmp in sorted.asReversed()) {
       addArgument(irString(buildString {
-        var last = -1
         newline()
+        append(rowSource)
+        if (indentations.isNotEmpty()) newline()
+
+        var last = -1
         for (i in indentations) {
-          if (i == tmp.startColumnNumber) break
           if (i > last) {
             indent(i - last - 1).append("|")
           }
           last = i
         }
-        indent(tmp.startColumnNumber - last - 1)
       }))
-      addArgument(irGet(tmp.variable))
+
+
+      for (tmp in rowValues.asReversed()) {
+        addArgument(irString(buildString {
+          var last = -1
+          newline()
+          for (i in indentations) {
+            if (i == tmp.indent) break
+            if (i > last) {
+              indent(i - last - 1).append("|")
+            }
+            last = i
+          }
+          indent(tmp.indent - last - 1)
+        }))
+        addArgument(irGet(tmp.value))
+      }
     }
   }
 }
+
+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.newline(): StringBuilder = append("\n")

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

@@ -75,9 +75,6 @@ class PowerAssertCallTransformer(
     if (!function.isAssert)
       return super.visitCall(expression)
 
-    val callSource = fileSource.substring(expression.startOffset, expression.endOffset)
-    val callIndent = file.info(expression).startColumnNumber
-
     val assertionArgument = expression.getValueArgument(0)!!
     val lambdaArgument = if (function.valueParameters.size == 2) expression.getValueArgument(1) else null
 
@@ -97,9 +94,9 @@ class PowerAssertCallTransformer(
 //      println(assertionArgument.dump())
 //      println(tree.dump())
 
-      val generator = object : PowerAssertGenerator(file, fileSource) {
+      val generator = object : PowerAssertGenerator() {
         override fun IrBuilderWithScope.buildAssertThrow(subStack: List<IrStackVariable>): IrExpression {
-          return buildThrow(constructor, buildMessage(title, subStack, callSource, callIndent))
+          return buildThrow(constructor, buildMessage(file, fileSource, title, expression, subStack))
         }
       }
 
@@ -113,9 +110,3 @@ class PowerAssertCallTransformer(
 
 val IrFunction.isAssert: Boolean
   get() = name.asString() == "assert" && getPackageFragment()?.fqName == KotlinBuiltIns.BUILT_INS_PACKAGE_FQ_NAME
-
-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.newline(): StringBuilder = append("\n")

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

@@ -18,24 +18,17 @@ package com.bnorm.power
 
 import org.jetbrains.kotlin.backend.common.lower.irIfThen
 import org.jetbrains.kotlin.backend.common.lower.irNot
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.IrStatementsBuilder
 import org.jetbrains.kotlin.ir.builders.irBlock
 import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irTemporary
-import org.jetbrains.kotlin.ir.declarations.IrFile
 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.IrWhen
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 
-abstract class PowerAssertGenerator(
+abstract class PowerAssertGenerator {
-  private val file: IrFile,
-  private val fileSource: String
-) {
   abstract fun IrBuilderWithScope.buildAssertThrow(subStack: List<IrStackVariable>): IrExpression
 
   fun buildAssert(
@@ -95,28 +88,7 @@ abstract class PowerAssertGenerator(
   ) : IrElementTransformerVoid() {
     private fun push(expression: IrExpression): IrGetValue = with(builder) {
       val variable = irTemporary(expression)
-      val source = fileSource.substring(expression)
+      stack.add(IrStackVariable(variable, expression))
-
-      var startColumnNumber = file.info(expression).startColumnNumber
-      if (expression is IrMemberAccessExpression) {
-        val descriptor = expression.descriptor
-        if (descriptor is FunctionDescriptor && descriptor.isInfix) {
-          startColumnNumber += source.indexOf(descriptor.name.asString())
-        }
-
-        // TODO handle equality and comparison better?
-        startColumnNumber += 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(IrStackVariable(variable, startColumnNumber, source))
       irGet(variable)
     }
 

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

@@ -282,6 +282,47 @@ assert(1 shl 1 == 4)
          |     |
          |     false
          2
+""".trimIndent()
+    )
+  }
+
+  @Test
+  fun multiline() {
+    assertMessage(
+      """fun main() {
+  val text: String? = "Hello"
+  assert(
+    text
+        == null ||
+        (
+            text.length == 5 &&
+                text.toLowerCase() == text
+            )
+  )
+}""",
+      """
+Assertion failed
+assert(
+  text
+  |
+  Hello
+      == null ||
+      |
+      false
+      (
+          text.length == 5 &&
+          |    |      |
+          |    |      true
+          |    5
+          Hello
+              text.toLowerCase() == text
+              |    |             |  |
+              |    |             |  Hello
+              |    |             false
+              |    hello
+              Hello
+          )
+)
 """.trimIndent()
     )
   }