Support for simple boolean conjunction and disjunction

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

@@ -0,0 +1,159 @@
+/*
+ * Copyright (C) 2020 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.bnorm.power
+
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.ir.expressions.IrConst
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.IrWhen
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
+
+sealed class Node {
+  abstract val parent: Node?
+  val mutableChildren: MutableList<Node> = mutableListOf()
+  val children: List<Node> get() = mutableChildren
+
+  protected fun dump(builder: StringBuilder, indent: Int) {
+    builder.append("  ".repeat(indent)).append(this).appendln()
+    for (child in children) {
+      child.dump(builder, indent + 1)
+    }
+  }
+}
+
+class AndNode(override val parent: Node) : Node() {
+  init {
+    parent.mutableChildren.add(this)
+  }
+
+  override fun toString() = "AndNode"
+}
+
+class OrNode(override val parent: Node) : Node() {
+  init {
+    parent.mutableChildren.add(this)
+  }
+
+  override fun toString() = "OrNode"
+}
+
+class ExpressionNode(
+  override val parent: Node
+) : Node() {
+  init {
+    parent.mutableChildren.add(this)
+  }
+
+  val expressions: MutableList<IrExpression> = mutableListOf()
+
+  override fun toString() = "ExpressionNode($expressions)"
+}
+
+class RootNode : Node() {
+  override val parent: Node? = null
+
+  fun dump(): String = buildString {
+    dump(this, 0)
+  }
+
+  override fun toString() = "RootNode"
+}
+
+fun buildAssertTree(expression: IrExpression): RootNode {
+  val tree = RootNode()
+  expression.accept(object : IrElementVisitor<Unit, Node> {
+    override fun visitElement(element: IrElement, data: Node) {
+      element.acceptChildren(this, data)
+    }
+
+    override fun visitExpression(expression: IrExpression, data: Node) {
+      val node = data as? ExpressionNode ?: ExpressionNode(data)
+      node.expressions += expression
+      expression.acceptChildren(this, node)
+    }
+
+    override fun visitCall(expression: IrCall, data: Node) {
+      if (expression.descriptor.name.asString() == "EQEQ" && expression.origin == IrStatementOrigin.EXCLEQ) {
+        // Skip the EQEQ part of a EXCLEQ call
+        expression.acceptChildren(this, data)
+      } else {
+        super.visitCall(expression, data)
+      }
+    }
+
+    override fun <T> visitConst(expression: IrConst<T>, data: Node) {
+      // Do not include constants
+    }
+
+    override fun visitWhen(expression: IrWhen, data: Node) {
+      when (expression.origin) {
+        IrStatementOrigin.ANDAND -> {
+          // flatten `&&` expressions to be at the same level
+          val node = data as? AndNode ?: AndNode(data)
+
+          require(expression.branches.size == 2)
+          val thenBranch = expression.branches[0]
+
+          thenBranch.condition.accept(this, node)
+          thenBranch.result.accept(this, node)
+
+          val elseBranchCondition = expression.branches[1].condition
+          val elseBranchResult = expression.branches[1].result
+
+          if (elseBranchCondition !is IrConst<*> || elseBranchCondition.value != true) {
+            elseBranchCondition.accept(this, node)
+          }
+
+          if (elseBranchResult !is IrConst<*> || elseBranchResult.value != false) {
+            elseBranchResult.accept(this, node)
+          }
+        }
+        IrStatementOrigin.OROR -> {
+          // flatten `||` expressions to be at the same level
+          val node = data as? OrNode ?: OrNode(data)
+
+          require(expression.branches.size == 2)
+          val thenBranchCondition = expression.branches[0].condition
+          val thenBranchResult = expression.branches[0].result
+          val elseBranchCondition = expression.branches[1].condition
+          val elseBranchResult = expression.branches[1].result
+
+          thenBranchCondition.accept(this, node)
+
+          if (thenBranchResult !is IrConst<*> || thenBranchResult.value != true) {
+            thenBranchResult.accept(this, node)
+          }
+
+          if (elseBranchCondition !is IrConst<*> || elseBranchCondition.value != true) {
+            elseBranchCondition.accept(this, node)
+          }
+
+          if (elseBranchResult !is IrConst<*> || elseBranchResult.value != false) {
+            elseBranchResult.accept(this, node)
+          }
+        }
+        else -> {
+          super.visitWhen(expression, data)
+        }
+      }
+    }
+  }, tree)
+
+  return tree
+}

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

@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 2020 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.bnorm.power
+
+import org.jetbrains.kotlin.backend.common.lower.irThrow
+import org.jetbrains.kotlin.ir.builders.IrBlockBuilder
+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.IrVariable
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrThrow
+import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
+
+data class IrStackVariable(
+  val variable: IrVariable,
+  val indentation: Int,
+  val source: String
+)
+
+fun IrBlockBuilder.buildThrow(
+  constructor: IrConstructorSymbol,
+  message: IrExpression
+): IrThrow = irThrow(irCall(constructor).apply {
+  putValueArgument(0, message)
+})
+
+fun IrBlockBuilder.buildMessage(
+  title: IrExpression,
+  stack: List<IrStackVariable>,
+  callSource: String
+): IrExpression {
+  return irConcat().apply {
+    addArgument(title)
+
+    val sorted = stack.sortedBy { it.indentation }
+    val indentations = sorted.map { it.indentation }
+
+    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()
+        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))
+    }
+  }
+}

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

@@ -23,27 +23,25 @@ import org.jetbrains.kotlin.backend.common.lower.at
 import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.irIfThen
 import org.jetbrains.kotlin.backend.common.lower.irNot
-import org.jetbrains.kotlin.backend.common.lower.irThrow
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.SourceRangeInfo
+import org.jetbrains.kotlin.ir.builders.IrBlockBuilder
 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.IrWhen
+import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
 import org.jetbrains.kotlin.ir.types.getClass
 import org.jetbrains.kotlin.ir.util.functions
 import org.jetbrains.kotlin.ir.util.getPackageFragment
@@ -67,15 +65,15 @@ fun FileLoweringPass.runOnFileInOrder(irFile: IrFile) {
   })
 }
 
+fun String.substring(expression: IrElement) = substring(expression.startOffset, expression.endOffset)
+fun IrFile.info(expression: IrElement) = fileEntry.getSourceRangeInfo(expression.startOffset, expression.endOffset)
+
 class PowerAssertCallTransformer(
   private val context: JvmBackendContext
 ) : IrElementTransformerVoid(), FileLoweringPass {
   private lateinit var file: IrFile
   private lateinit var fileSource: String
 
-  // TODO this is the only thing keeping this project from being multiplatform
-  private val constructor = this@PowerAssertCallTransformer.context.ir.symbols.assertionErrorConstructor
-
   override fun lower(irFile: IrFile) {
     file = irFile
     fileSource = File(irFile.path).readText()
@@ -83,12 +81,6 @@ 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)
@@ -97,131 +89,118 @@ class PowerAssertCallTransformer(
     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
+
     context.createIrBuilder(expression.symbol).run {
       at(expression)
 
-      return irBlock {
+      val lambda = lambdaArgument?.asSimpleLambda()
-        val stack = mutableListOf<IrTemporaryVariable>()
+      val title = when {
-
+        lambda != null -> lambda.inline()
-        fun push(expression: IrExpression): IrGetValue {
+        lambdaArgument != null -> {
-          val variable = irTemporary(expression)
+          val invoke = lambdaArgument.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
-          val source = fileSource.substring(expression.startOffset, expression.endOffset)
+          irCallOp(invoke.symbol, invoke.returnType, lambdaArgument)
-
-          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)
         }
-
+        else -> irString("Assertion failed")
-        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).appendln()
-//          val sorted = stack.sortedBy { it.indentation }
-//
-//          val indentations = sorted.map { it.indentation }
-//          var last = -1
-//          for (i in indentations) {
-//            if (i > last) {
-//              indent(i - last - 1).append("|")
-//            }
-//            last = i
-//          }
-//          appendln()
-//
-//          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).appendln()
-//          }
-//        })
-
-        val lambdaArgument = if (function.valueParameters.size == 2) expression.getValueArgument(1) else null
-        val lambda = lambdaArgument?.asSimpleLambda()
-        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 }
-
-            addArgument(irString(buildString {
-              appendln()
-              append(callSource).appendln()
-              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
-                appendln()
-                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(assertCondition), throwError)
       }
+
+      val tree = buildAssertTree(assertionArgument)
+      val root = tree.children.single()
+
+//      println(assertionArgument.dump())
+//      println(tree.dump())
+
+      return irBlock {
+        buildAssert(this@PowerAssertCallTransformer.context, file, fileSource, callSource, callIndent, title, root)
+      }
+    }
+  }
+}
+
+fun IrBlockBuilder.buildAssert(
+  context: JvmBackendContext,
+  file: IrFile,
+  fileSource: String,
+  callSource: String,
+  callIndent: Int,
+  title: IrExpression,
+  node: Node,
+  stack: MutableList<IrStackVariable> = mutableListOf(),
+  constructor: IrConstructorSymbol = context.ir.symbols.assertionErrorConstructor,
+  thenPart: IrBlockBuilder.(stack: MutableList<IrStackVariable>) -> IrExpression = { stack -> buildThrow(constructor, buildMessage(title, stack, callSource)) }
+) {
+  fun IrBlockBuilder.nest(children: List<Node>, index: Int, stack: MutableList<IrStackVariable>) {
+    val child = children[index]
+    buildAssert(context, file, fileSource, callSource, callIndent, title, child, stack, constructor) { stack ->
+      if (index + 1 == children.size) buildThrow(constructor, buildMessage(title, stack, callSource))
+      else irBlock { nest(children, index + 1, stack) }
+    }
+  }
+
+  when (node) {
+    is ExpressionNode -> {
+      +irIfNotThan(stack, file, fileSource, callIndent, node, thenPart)
+    }
+    is AndNode -> {
+      for (child in node.children) {
+        buildAssert(context, file, fileSource, callSource, callIndent, title, child, stack, constructor, thenPart)
+      }
+    }
+    is OrNode -> {
+      nest(node.children, 0, stack)
+    }
+  }
+}
+
+private inline fun IrBlockBuilder.irIfNotThan(
+  stack: MutableList<IrStackVariable>,
+  file: IrFile,
+  fileSource: String,
+  callIndent: Int,
+  node: ExpressionNode,
+  thenPart: IrBlockBuilder.(stack: MutableList<IrStackVariable>) -> IrExpression
+): IrWhen {
+  val stackTransformer = StackBuilder(this, stack, file, fileSource, callIndent, node.expressions)
+  val transformed = node.expressions.first().transform(stackTransformer, null)
+  return irIfThen(irNot(transformed), thenPart(stack))
+}
+
+class StackBuilder(
+  private val builder: IrBlockBuilder,
+  private val stack: MutableList<IrStackVariable>,
+  private val file: IrFile,
+  private val fileSource: String,
+  private val callIndent: Int,
+  private val transform: List<IrExpression>
+) : IrElementTransformerVoid() {
+  private fun push(expression: IrExpression): IrGetValue = with(builder) {
+    val variable = irTemporary(expression)
+    val source = fileSource.substring(expression)
+
+    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(IrStackVariable(variable, indentation, source))
+    irGet(variable)
+  }
+
+  override fun visitExpression(expression: IrExpression): IrExpression {
+    return if (expression in transform) {
+      push(super.visitExpression(expression))
+    } else {
+      super.visitExpression(expression)
     }
   }
 }
@@ -229,8 +208,5 @@ class PowerAssertCallTransformer(
 val IrFunction.isAssert: Boolean
   get() = name.asString() == "assert" && getPackageFragment()?.fqName == KotlinBuiltIns.BUILT_INS_PACKAGE_FQ_NAME
 
-fun IrFile.info(expression: IrElement): SourceRangeInfo {
-  return 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/test/kotlin/com/bnorm/power/test.kt

@@ -78,6 +78,124 @@ Not equal
 assert(1 == 2) { "Not equal" }
          |
          false
+""".trimIndent()
+    )
+  }
+
+  @Test
+  fun booleanExpressionsShortCircuit() {
+    assertMessage(
+      """
+fun main() {
+  val text: String? = null
+  assert(text != null && text.length == 1)
+}""",
+      """
+Assertion failed
+assert(text != null && text.length == 1)
+       |    |
+       |    false
+       null
+""".trimIndent()
+    )
+  }
+
+  @Test
+  fun booleanAnd() {
+    assertMessage(
+      """
+fun main() {
+  val text: String? = "Hello"
+  assert(text != null && text.length == 5 && text.toLowerCase() == text)
+}""",
+      """
+Assertion failed
+assert(text != null && text.length == 5 && text.toLowerCase() == text)
+       |    |          |    |      |       |    |             |  |
+       |    |          |    |      |       |    |             |  Hello
+       |    |          |    |      |       |    |             false
+       |    |          |    |      |       |    hello
+       |    |          |    |      |       Hello
+       |    |          |    |      true
+       |    |          |    5
+       |    |          Hello
+       |    true
+       Hello
+""".trimIndent()
+    )
+  }
+
+  @Test
+  fun booleanOr() {
+    assertMessage(
+      """
+fun main() {
+  val text: String? = "Hello"
+  assert(text == null || text.length == 1 || text.toLowerCase() == text)
+}""",
+      """
+Assertion failed
+assert(text == null || text.length == 1 || text.toLowerCase() == text)
+       |    |          |    |      |       |    |             |  |
+       |    |          |    |      |       |    |             |  Hello
+       |    |          |    |      |       |    |             false
+       |    |          |    |      |       |    hello
+       |    |          |    |      |       Hello
+       |    |          |    |      false
+       |    |          |    5
+       |    |          Hello
+       |    false
+       Hello
+""".trimIndent()
+    )
+  }
+
+  @Test
+  fun booleanMixAndFirst() {
+    assertMessage(
+      """
+fun main() {
+  val text: String? = "Hello"
+  assert(text != null && (text.length == 1 || text.toLowerCase() == text))
+}""",
+      """
+Assertion failed
+assert(text != null && (text.length == 1 || text.toLowerCase() == text))
+       |    |           |    |      |       |    |             |  |
+       |    |           |    |      |       |    |             |  Hello
+       |    |           |    |      |       |    |             false
+       |    |           |    |      |       |    hello
+       |    |           |    |      |       Hello
+       |    |           |    |      false
+       |    |           |    5
+       |    |           Hello
+       |    true
+       Hello
+""".trimIndent()
+    )
+  }
+
+  @Test
+  fun booleanMixOrFirst() {
+    assertMessage(
+      """
+fun main() {
+  val text: String? = "Hello"
+  assert(text == null || (text.length == 5 && text.toLowerCase() == text))
+}""",
+      """
+Assertion failed
+assert(text == null || (text.length == 5 && text.toLowerCase() == text))
+       |    |           |    |      |       |    |             |  |
+       |    |           |    |      |       |    |             |  Hello
+       |    |           |    |      |       |    |             false
+       |    |           |    |      |       |    hello
+       |    |           |    |      |       Hello
+       |    |           |    |      true
+       |    |           |    5
+       |    |           Hello
+       |    false
+       Hello
 """.trimIndent()
     )
   }