Clean up temporary variable nesting logic

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

@@ -16,11 +16,10 @@
 
 package com.bnorm.power
 
-import com.bnorm.power.diagram.IrTemporaryVariable
+import com.bnorm.power.diagram.buildDiagramNesting
-import com.bnorm.power.diagram.DiagramGenerator
 import com.bnorm.power.diagram.buildTree
 import com.bnorm.power.diagram.info
-import com.bnorm.power.diagram.irDiagram
+import com.bnorm.power.diagram.irDiagramString
 import com.bnorm.power.diagram.substring
 import com.bnorm.power.internal.ReturnableBlockTransformer
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
@@ -121,6 +120,7 @@ class PowerAssertCallTransformer(
       return super.visitCall(expression)
     }
 
+    // TODO - support more arguments by currying?
     val assertionArgument = expression.getValueArgument(0)!!
     val messageArgument = if (function.valueParameters.size == 2) expression.getValueArgument(1) else null
 
@@ -129,45 +129,44 @@ class PowerAssertCallTransformer(
       messageCollector.info(expression, "Expression is constant and will not be power-assert transformed")
       return super.visitCall(expression)
     }
+//      println(root.dump())
 
     val symbol = currentScope!!.scope.scopeOwnerSymbol
-    DeclarationIrBuilder(context, symbol).run {
+    val builder = DeclarationIrBuilder(context, symbol, expression.startOffset, expression.endOffset)
-      at(expression)
+    return builder.buildDiagramNesting(root) { argument, variables ->
-
+      val lambda = messageArgument?.asSimpleLambda()
-      val generator = object : DiagramGenerator() {
+      val title = when {
-        override fun IrBuilderWithScope.buildCall(argument: IrExpression, variables: List<IrTemporaryVariable>): IrExpression {
+        messageArgument is IrConst<*> -> messageArgument
-
+        messageArgument is IrStringConcatenation -> messageArgument
-          val lambda = messageArgument?.asSimpleLambda()
+        lambda != null -> lambda.deepCopyWithSymbols(parent).inline(parent)
-          val title = when {
+          .transform(ReturnableBlockTransformer(context, symbol), null)
-            messageArgument is IrConst<*> -> messageArgument
+        messageArgument != null -> {
-            messageArgument is IrStringConcatenation -> messageArgument
+          val invoke =
-            lambda != null -> lambda.deepCopyWithSymbols(parent).inline(parent).transform(ReturnableBlockTransformer(context, symbol), null)
+            messageArgument.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
-            messageArgument != null -> {
+          irCallOp(invoke.symbol, invoke.returnType, messageArgument)
-              val invoke = messageArgument.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
-              irCallOp(invoke.symbol, invoke.returnType, messageArgument)
-            }
-            // TODO what should the default message be?
-            assertionArgument.type.isBoolean() -> irString("Assertion failed")
-            else -> null
-          }
-
-          val prefix = title?.deepCopyWithSymbols(parent)
-          val diagram = irDiagram(file, fileSource, prefix, expression, variables)
-          return delegate.buildCall(this, expression, argument, diagram)
         }
+        // TODO what should the default message be?
+        assertionArgument.type.isBoolean() -> irString("Assertion failed")
+        else -> null
       }
 
-//      println(root.dump())
+      val prefix = title?.deepCopyWithSymbols(parent)
-      return generator.buildExpression(this, root)
+      val diagram = irDiagramString(file, fileSource, prefix, expression, variables)
+      delegate.buildCall(this, expression, argument, diagram)
+    }
 //        .also { println(expression.dump()) }
 //        .also { println(it.dump()) }
 //        .also { println(expression.dumpKotlinLike()) }
 //        .also { println(it.dumpKotlinLike()) }
-    }
   }
 
   private interface FunctionDelegate {
-    fun buildCall(builder: IrBuilderWithScope, original: IrCall, argument: IrExpression, message: IrExpression): IrExpression
+    fun buildCall(
+      builder: IrBuilderWithScope,
+      original: IrCall,
+      argument: IrExpression,
+      message: IrExpression
+    ): IrExpression
   }
 
   private fun findDelegate(function: IrFunction): FunctionDelegate? {
@@ -184,7 +183,12 @@ class PowerAssertCallTransformer(
         return@mapNotNull when {
           isStringSupertype(messageParameter.type) -> {
             object : FunctionDelegate {
-              override fun buildCall(builder: IrBuilderWithScope, original: IrCall, argument: IrExpression, message: IrExpression): IrExpression = with(builder) {
+              override fun buildCall(
+                builder: IrBuilderWithScope,
+                original: IrCall,
+                argument: IrExpression,
+                message: IrExpression
+              ): IrExpression = with(builder) {
                 irCall(overload, type = original.type).apply {
                   dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
                   extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
@@ -199,7 +203,12 @@ class PowerAssertCallTransformer(
           }
           isStringFunction(messageParameter.type) -> {
             object : FunctionDelegate {
-              override fun buildCall(builder: IrBuilderWithScope, original: IrCall, argument: IrExpression, message: IrExpression): IrExpression = with(builder) {
+              override fun buildCall(
+                builder: IrBuilderWithScope,
+                original: IrCall,
+                argument: IrExpression,
+                message: IrExpression
+              ): IrExpression = with(builder) {
                 val scope = this
                 val lambda = builder.context.irFactory.buildFun {
                   name = Name.special("<anonymous>")
@@ -213,7 +222,13 @@ class PowerAssertCallTransformer(
                   }
                   parent = scope.parent
                 }
-                val expression = IrFunctionExpressionImpl(original.startOffset, original.endOffset, messageParameter.type, lambda, IrStatementOrigin.LAMBDA)
+                val expression = IrFunctionExpressionImpl(
+                  original.startOffset,
+                  original.endOffset,
+                  messageParameter.type,
+                  lambda,
+                  IrStatementOrigin.LAMBDA
+                )
                 irCall(overload, type = original.type).apply {
                   dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
                   extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)

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

@@ -0,0 +1,142 @@
+/*
+ * 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.diagram
+
+import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
+import org.jetbrains.kotlin.ir.builders.irBlock
+import org.jetbrains.kotlin.ir.builders.irFalse
+import org.jetbrains.kotlin.ir.builders.irIfThenElse
+import org.jetbrains.kotlin.ir.builders.irTrue
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
+
+fun IrBuilderWithScope.buildDiagramNesting(
+  root: Node,
+  call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
+): IrExpression {
+  return buildExpression(root, listOf()) { argument, subStack ->
+    call(argument, subStack)
+  }
+}
+
+private fun IrBuilderWithScope.buildExpression(
+  node: Node,
+  variables: List<IrTemporaryVariable>,
+  call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
+): IrExpression = when (node) {
+  is ExpressionNode -> add(node, variables, call)
+  is AndNode -> nest(node, 0, variables, call)
+  is OrNode -> nest(node, 0, variables, call)
+  else -> TODO("Unknown node type=$node")
+}
+
+/**
+ * ```
+ * val result = call(1 + 2 + 3)
+ * ```
+ * Transforms to
+ * ```
+ * val result = run {
+ *   val tmp0 = 1 + 2
+ *   val tmp1 = tmp0 + 3
+ *   call(tmp1, <diagram>)
+ * }
+ * ```
+ */
+private fun IrBuilderWithScope.add(
+  node: ExpressionNode,
+  variables: List<IrTemporaryVariable>,
+  call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
+): IrExpression {
+  return irBlock {
+    val head = node.expressions.first().deepCopyWithSymbols(scope.getLocalDeclarationParent())
+    val expressions = (buildTree(head) as ExpressionNode).expressions
+    val transformer = IrTemporaryExtractionTransformer(this@irBlock, expressions.toSet())
+    val transformed = expressions.first().transform(transformer, null)
+    +call(transformed, variables + transformer.variables)
+  }
+}
+
+/**
+ * ```
+ * val result = call(1 == 1 && 2 == 2)
+ * ```
+ * Transforms to
+ * ```
+ * val result = run {
+ *   val tmp0 = 1 == 1
+ *   if (tmp0) {
+ *     val tmp1 = 2 == 2
+ *     call(tmp1, <diagram>)
+ *   }
+ *   else call(false, <diagram>)
+ * }
+ * ```
+ */
+private fun IrBuilderWithScope.nest(
+  node: AndNode,
+  index: Int,
+  variables: List<IrTemporaryVariable>,
+  call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
+): IrExpression {
+  val children = node.children
+  val child = children[index]
+  return buildExpression(child, variables) { argument, newVariables ->
+    if (index + 1 == children.size) call(argument, newVariables) // last expression, result is false
+    else irIfThenElse(
+      context.irBuiltIns.anyType,
+      argument,
+      nest(node, index + 1, newVariables, call), // more expressions, continue nesting
+      call(irFalse(), newVariables), // short-circuit result to false
+    )
+  }
+}
+
+/**
+ * ```
+ * val result = call(1 == 1 || 2 == 2)
+ * ```
+ * Transforms to
+ * ```
+ * val result = run {
+ *   val tmp0 = 1 == 1
+ *   if (tmp0) call(true, <diagram>)
+ *   else {
+ *     val tmp1 = 2 == 2
+ *     call(tmp1, <diagram>)
+ *   }
+ * }
+ * ```
+ */
+private fun IrBuilderWithScope.nest(
+  node: OrNode,
+  index: Int,
+  variables: List<IrTemporaryVariable>,
+  call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
+): IrExpression {
+  val children = node.children
+  val child = children[index]
+  return buildExpression(child, variables) { argument, newVariables ->
+    if (index + 1 == children.size) call(argument, newVariables) // last expression, result is false
+    else irIfThenElse(
+      context.irBuiltIns.anyType,
+      argument,
+      call(irTrue(), newVariables), // short-circuit result to true
+      nest(node, index + 1, newVariables, call), // more expressions, continue nesting
+    )
+  }
+}

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

@@ -1,153 +0,0 @@
-/*
- * 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.diagram
-
-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.irFalse
-import org.jetbrains.kotlin.ir.builders.irIfThenElse
-import org.jetbrains.kotlin.ir.builders.irTrue
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
-
-abstract class DiagramGenerator {
-  abstract fun IrBuilderWithScope.buildCall(
-    argument: IrExpression,
-    variables: List<IrTemporaryVariable>
-  ): IrExpression
-
-  fun buildExpression(
-    builder: IrBuilderWithScope,
-    root: Node
-  ): IrExpression {
-    return builder.irBlock {
-      buildExpression(root, listOf()) { argument, subStack ->
-        buildCall(argument, subStack)
-      }
-    }
-  }
-
-  private fun IrStatementsBuilder<*>.buildExpression(
-    node: Node,
-    variables: List<IrTemporaryVariable>,
-    call: IrStatementsBuilder<*>.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
-  ) {
-    when (node) {
-      is ExpressionNode -> add(node, variables, call)
-      is AndNode -> nest(node, 0, variables, call)
-      is OrNode -> nest(node, 0, variables, call)
-      else -> TODO("Unknown node type=$node")
-    }
-  }
-
-  /**
-   * ```
-   * val result = call(1 + 2 + 3)
-   * ```
-   * Transforms to
-   * ```
-   * val result = run {
-   *   val tmp0 = 1 + 2
-   *   val tmp1 = tmp0 + 3
-   *   call(tmp1, <diagram>)
-   * }
-   * ```
-   */
-  private fun IrStatementsBuilder<*>.add(
-    node: ExpressionNode,
-    variables: List<IrTemporaryVariable>,
-    call: IrStatementsBuilder<*>.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
-  ) {
-    val head = node.expressions.first().deepCopyWithSymbols(scope.getLocalDeclarationParent())
-    val expressions = (buildTree(head) as ExpressionNode).expressions
-    val transformer = IrTemporaryExtractionTransformer(this, expressions.toSet())
-    val transformed = expressions.first().transform(transformer, null)
-
-    +call(transformed, variables + transformer.variables)
-  }
-
-  /**
-   * ```
-   * val result = call(1 == 1 && 2 == 2)
-   * ```
-   * Transforms to
-   * ```
-   * val result = run {
-   *   val tmp0 = 1 == 1
-   *   if (tmp0) {
-   *     val tmp1 = 2 == 2
-   *     call(tmp1, <diagram>)
-   *   }
-   *   else call(false, <diagram>)
-   * }
-   * ```
-   */
-  private fun IrStatementsBuilder<*>.nest(
-    node: AndNode,
-    index: Int,
-    variables: List<IrTemporaryVariable>,
-    call: IrStatementsBuilder<*>.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
-  ) {
-    val children = node.children
-    val child = children[index]
-    buildExpression(child, variables) { argument, newVariables ->
-      if (index + 1 == children.size) call(argument, newVariables) // last expression, result is false
-      else irIfThenElse(
-        context.irBuiltIns.anyType,
-        argument,
-        irBlock { nest(node, index + 1, newVariables, call) }, // more expressions, continue nesting
-        call(irFalse(), newVariables), // short-circuit result to false
-      )
-    }
-  }
-
-  /**
-   * ```
-   * val result = call(1 == 1 || 2 == 2)
-   * ```
-   * Transforms to
-   * ```
-   * val result = run {
-   *   val tmp0 = 1 == 1
-   *   if (tmp0) call(true, <diagram>)
-   *   else {
-   *     val tmp1 = 2 == 2
-   *     call(tmp1, <diagram>)
-   *   }
-   * }
-   * ```
-   */
-  private fun IrStatementsBuilder<*>.nest(
-    node: OrNode,
-    index: Int,
-    variables: List<IrTemporaryVariable>,
-    call: IrStatementsBuilder<*>.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
-  ) {
-    val children = node.children
-    val child = children[index]
-    buildExpression(child, variables) { argument, newVariables ->
-      if (index + 1 == children.size) call(argument, newVariables) // last expression, result is false
-      else irIfThenElse(
-        context.irBuiltIns.anyType,
-        argument,
-        call(irTrue(), newVariables), // short-circuit result to true
-        irBlock { nest(node, index + 1, newVariables, call) }, // more expressions, continue nesting
-      )
-    }
-  }
-}

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

@@ -32,7 +32,7 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 
-fun IrBuilderWithScope.irDiagram(
+fun IrBuilderWithScope.irDiagramString(
   file: IrFile,
   fileSource: String,
   prefix: IrExpression? = null,