Extract a more general power assert generator

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

@@ -17,7 +17,7 @@
 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.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irCall
 import org.jetbrains.kotlin.ir.builders.irConcat
 import org.jetbrains.kotlin.ir.builders.irGet
@@ -29,27 +29,29 @@ import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
 
 data class IrStackVariable(
   val variable: IrVariable,
-  val indentation: Int,
+  val startColumnNumber: Int,
   val source: String
 )
 
-fun IrBlockBuilder.buildThrow(
+fun IrBuilderWithScope.buildThrow(
   constructor: IrConstructorSymbol,
   message: IrExpression
 ): IrThrow = irThrow(irCall(constructor).apply {
   putValueArgument(0, message)
 })
 
-fun IrBlockBuilder.buildMessage(
+fun IrBuilderWithScope.buildMessage(
   title: IrExpression,
   stack: List<IrStackVariable>,
-  callSource: String
+  callSource: String,
+  callIndent: Int = 0
 ): IrExpression {
+  val stack = stack.map { it.copy(startColumnNumber = it.startColumnNumber - callIndent) }
   return irConcat().apply {
     addArgument(title)
 
-    val sorted = stack.sortedBy { it.indentation }
-    val indentations = sorted.map { it.indentation }
+    val sorted = stack.sortedBy { it.startColumnNumber }
+    val indentations = sorted.map { it.startColumnNumber }
 
     addArgument(irString(buildString {
       newline()
@@ -69,13 +71,13 @@ fun IrBlockBuilder.buildMessage(
         var last = -1
         newline()
         for (i in indentations) {
-          if (i == tmp.indentation) break
+          if (i == tmp.startColumnNumber) break
           if (i > last) {
             indent(i - last - 1).append("|")
           }
           last = i
         }
-        indent(tmp.indentation - last - 1)
+        indent(tmp.startColumnNumber - last - 1)
       }))
       addArgument(irGet(tmp.variable))
     }

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

@@ -21,26 +21,17 @@ 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
-import org.jetbrains.kotlin.backend.common.lower.irNot
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.ir.IrElement
-import org.jetbrains.kotlin.ir.builders.IrBlockBuilder
-import org.jetbrains.kotlin.ir.builders.irBlock
+import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irCallOp
-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.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
@@ -65,14 +56,12 @@ 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
+  private val constructor: IrConstructorSymbol = context.ir.symbols.assertionErrorConstructor
 
   override fun lower(irFile: IrFile) {
     file = irFile
@@ -105,110 +94,28 @@ class PowerAssertCallTransformer(
         else -> irString("Assertion failed")
       }
 
-      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)
+      val generator = object : PowerAssertGenerator(file, fileSource) {
+        override fun IrBuilderWithScope.buildAssertThrow(subStack: List<IrStackVariable>): IrExpression {
+          return buildThrow(constructor, buildMessage(title, subStack, callSource, callIndent))
+        }
       }
+
+      val tree = buildAssertTree(assertionArgument)
+      val root = tree.children.single()
+      return generator.buildAssert(this, root)
 //        .also { println(it.dump())}
     }
   }
 }
 
-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 = { subStack -> buildThrow(constructor, buildMessage(title, subStack, 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) { subStack ->
-      if (index + 1 == children.size) buildThrow(constructor, buildMessage(title, subStack, callSource))
-      else irBlock { nest(children, index + 1, subStack) }
-    }
-  }
-
-  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.(subStack: MutableList<IrStackVariable>) -> IrExpression
-): IrWhen {
-  val expressions = node.getExpressionsCopy()
-  val stackTransformer = StackBuilder(this, stack, file, fileSource, callIndent, expressions)
-  val transformed = expressions.first().transform(stackTransformer, null)
-  return irIfThen(irNot(transformed), thenPart(stack.toMutableList()))
-}
-
-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)
-    }
-  }
-}
-
 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

@@ -0,0 +1,134 @@
+/*
+ * 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.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(
+  private val file: IrFile,
+  private val fileSource: String
+) {
+
+//  private fun IrStatementsBuilder<*>.buildAssertThrow(
+//    callSource: String,
+//    title: IrExpression,
+//    subStack: MutableList<IrStackVariable>,
+//    callIndent: Int = 0
+//  ) = buildThrow(constructor, buildMessage(title, subStack, callSource, callIndent))
+
+  abstract fun IrBuilderWithScope.buildAssertThrow(subStack: List<IrStackVariable>): IrExpression
+
+  fun buildAssert(
+    builder: IrBuilderWithScope,
+    root: Node
+  ): IrExpression {
+    return builder.irBlock {
+      buildAssert(root, mutableListOf()) { subStack ->
+        buildAssertThrow(subStack)
+      }
+    }
+  }
+
+  private fun IrStatementsBuilder<*>.buildAssert(
+    node: Node,
+    stack: MutableList<IrStackVariable>,
+    thenPart: IrStatementsBuilder<*>.(stack: MutableList<IrStackVariable>) -> IrExpression
+  ) {
+    fun IrStatementsBuilder<*>.nest(children: List<Node>, index: Int, stack: MutableList<IrStackVariable>) {
+      val child = children[index]
+      buildAssert(child, stack) { subStack ->
+        if (index + 1 == children.size) buildAssertThrow(subStack)
+        else irBlock { nest(children, index + 1, subStack) }
+      }
+    }
+
+    when (node) {
+      is ExpressionNode -> {
+        +irIfNotThan(stack, node, thenPart)
+      }
+      is AndNode -> {
+        for (child in node.children) {
+          buildAssert(child, stack, thenPart)
+        }
+      }
+      is OrNode -> {
+        nest(node.children, 0, stack)
+      }
+    }
+  }
+
+  private inline fun IrStatementsBuilder<*>.irIfNotThan(
+    stack: MutableList<IrStackVariable>,
+    node: ExpressionNode,
+    thenPart: IrStatementsBuilder<*>.(subStack: MutableList<IrStackVariable>) -> IrExpression
+  ): IrWhen {
+    val expressions = node.getExpressionsCopy()
+    val stackTransformer = StackBuilder(this, stack, expressions)
+    val transformed = expressions.first().transform(stackTransformer, null)
+    return irIfThen(irNot(transformed), thenPart(stack.toMutableList()))
+  }
+
+  inner class StackBuilder(
+    private val builder: IrStatementsBuilder<*>,
+    private val stack: MutableList<IrStackVariable>,
+    private val transform: List<IrExpression>
+  ) : IrElementTransformerVoid() {
+    private fun push(expression: IrExpression): IrGetValue = with(builder) {
+      val variable = irTemporary(expression)
+      val source = fileSource.substring(expression)
+
+      var startColumnNumber = file.info(expression).startColumnNumber
+      if (expression is IrMemberAccessExpression) {
+        // TODO Is this the best way to fix indentation of infix operators?
+        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)
+    }
+
+    override fun visitExpression(expression: IrExpression): IrExpression {
+      return if (expression in transform) {
+        push(super.visitExpression(expression))
+      } else {
+        super.visitExpression(expression)
+      }
+    }
+  }
+}