Prepare for diagramming of generic expressions

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

@@ -16,15 +16,23 @@
 
 package com.bnorm.power
 
+import com.bnorm.power.diagram.IrTemporaryVariable
+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.substring
 import com.bnorm.power.internal.ReturnableBlockTransformer
-import java.io.File
-import org.jetbrains.kotlin.backend.common.*
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.backend.common.ir.asSimpleLambda
 import org.jetbrains.kotlin.backend.common.ir.inline
 import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
 import org.jetbrains.kotlin.backend.common.lower.at
-import org.jetbrains.kotlin.cli.common.messages.*
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
@@ -39,16 +47,30 @@ import org.jetbrains.kotlin.ir.builders.parent
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.path
-import org.jetbrains.kotlin.ir.expressions.*
+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.IrStringConcatenation
 import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionExpressionImpl
-import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.types.IrSimpleType
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.types.IrTypeArgument
+import org.jetbrains.kotlin.ir.types.IrTypeProjection
+import org.jetbrains.kotlin.ir.types.getClass
+import org.jetbrains.kotlin.ir.types.isBoolean
+import org.jetbrains.kotlin.ir.types.isSubtypeOf
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
+import org.jetbrains.kotlin.ir.util.functions
+import org.jetbrains.kotlin.ir.util.isFunctionOrKFunction
+import org.jetbrains.kotlin.ir.util.kotlinFqName
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import java.io.File
 
 fun FileLoweringPass.runOnFileInOrder(irFile: IrFile) {
   irFile.acceptVoid(object : IrElementVisitorVoid {
@@ -98,8 +120,7 @@ class PowerAssertCallTransformer(
     val messageArgument = if (function.valueParameters.size == 2) expression.getValueArgument(1) else null
 
     // If the tree does not contain any children, the expression is not transformable
-    val tree = buildAssertTree(assertionArgument)
-    val root = tree.children.singleOrNull() ?: run {
+    val root = buildTree(assertionArgument) ?: run {
       messageCollector.info(expression, "Expression is constant and will not be power-assert transformed")
       return super.visitCall(expression)
     }
@@ -108,8 +129,8 @@ class PowerAssertCallTransformer(
     DeclarationIrBuilder(context, symbol).run {
       at(expression)
 
-      val generator = object : PowerAssertGenerator() {
-        override fun IrBuilderWithScope.buildAssertThrow(subStack: List<IrStackVariable>): IrExpression {
+      val generator = object : DiagramGenerator() {
+        override fun IrBuilderWithScope.buildAssertThrow(variables: List<IrTemporaryVariable>): IrExpression {
 
           val lambda = messageArgument?.asSimpleLambda()
           val title = when {
@@ -124,20 +145,24 @@ class PowerAssertCallTransformer(
             else -> irString("Assertion failed")
           }
 
-          return delegate.buildCall(this, expression, buildMessage(file, fileSource, title.deepCopyWithSymbols(parent), expression, subStack))
+          val prefix = title.deepCopyWithSymbols(parent)
+          val diagram = irDiagram(file, fileSource, prefix, expression, variables)
+          return delegate.buildCall(this, expression, irFalse(), diagram)
         }
       }
 
-//      println(expression.dump())
-//      println(tree.dump())
+//      println(root.dump())
 
       return generator.buildAssert(this, root)
+//        .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, message: IrExpression): IrExpression
+    fun buildCall(builder: IrBuilderWithScope, original: IrCall, argument: IrExpression, message: IrExpression): IrExpression
   }
 
   private fun findDelegate(fqName: FqName): FunctionDelegate? {
@@ -148,25 +173,26 @@ class PowerAssertCallTransformer(
         if (parameters.size != 2) return@mapNotNull null
         if (!parameters[0].type.isBoolean()) return@mapNotNull null
 
+        val messageParameter = parameters.last()
         return@mapNotNull when {
-          isStringSupertype(parameters[1].type) -> {
+          isStringSupertype(messageParameter.type) -> {
             object : FunctionDelegate {
-              override fun buildCall(builder: IrBuilderWithScope, original: IrCall, message: IrExpression): IrExpression = with(builder) {
+              override fun buildCall(builder: IrBuilderWithScope, original: IrCall, argument: IrExpression, message: IrExpression): IrExpression = with(builder) {
                 irCall(overload, type = overload.owner.returnType).apply {
                   dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
                   extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
                   for (i in 0 until original.typeArgumentsCount) {
                     putTypeArgument(i, original.getTypeArgument(i))
                   }
-                  putValueArgument(0, irFalse())
+                  putValueArgument(0, argument)
                   putValueArgument(1, message)
                 }
               }
             }
           }
-          isStringFunction(parameters[1].type) -> {
+          isStringFunction(messageParameter.type) -> {
             object : FunctionDelegate {
-              override fun buildCall(builder: IrBuilderWithScope, original: IrCall, 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>")
@@ -180,14 +206,14 @@ class PowerAssertCallTransformer(
                   }
                   parent = scope.parent
                 }
-                val expression = IrFunctionExpressionImpl(original.startOffset, original.endOffset, parameters[1].type, lambda, IrStatementOrigin.LAMBDA)
+                val expression = IrFunctionExpressionImpl(original.startOffset, original.endOffset, messageParameter.type, lambda, IrStatementOrigin.LAMBDA)
                 irCall(overload, type = overload.owner.returnType).apply {
                   dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
                   extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
                   for (i in 0 until original.typeArgumentsCount) {
                     putTypeArgument(i, original.getTypeArgument(i))
                   }
-                  putValueArgument(0, irFalse())
+                  putValueArgument(0, argument)
                   putValueArgument(1, expression)
                 }
               }

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

@@ -1,104 +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
-
-import org.jetbrains.kotlin.backend.common.lower.irIfThen
-import org.jetbrains.kotlin.backend.common.lower.irNot
-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.builders.parent
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrWhen
-import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
-import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
-
-abstract class PowerAssertGenerator {
-  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(this.parent)
-    val stackTransformer = StackBuilder(this, stack, expressions)
-    val transformed = expressions.first().transform(stackTransformer, null)
-    return irIfThen(irNot(transformed), thenPart(stack.toMutableList()))
-  }
-
-  class StackBuilder(
-    private val builder: IrStatementsBuilder<*>,
-    private val stack: MutableList<IrStackVariable>,
-    private val transform: List<IrExpression>
-  ) : IrElementTransformerVoid() {
-
-    override fun visitExpression(expression: IrExpression): IrExpression {
-      return if (expression in transform) {
-        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/diagram/DiagramGenerator.kt

@@ -0,0 +1,147 @@
+/*
+ * 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.backend.common.lower.irIfThen
+import org.jetbrains.kotlin.backend.common.lower.irNot
+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.expressions.IrExpression
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
+
+abstract class DiagramGenerator {
+  abstract fun IrBuilderWithScope.buildAssertThrow(variables: List<IrTemporaryVariable>): IrExpression
+
+  fun buildAssert(
+    builder: IrBuilderWithScope,
+    root: Node
+  ): IrExpression {
+    return builder.irBlock {
+      buildAssert(root, listOf()) { subStack ->
+        buildAssertThrow(subStack)
+      }
+    }
+  }
+
+  private fun IrStatementsBuilder<*>.buildAssert(
+    node: Node,
+    variables: List<IrTemporaryVariable>,
+    thenPart: IrStatementsBuilder<*>.(List<IrTemporaryVariable>) -> IrExpression
+  ): List<IrTemporaryVariable> = when (node) {
+    is ExpressionNode -> add(node, variables, thenPart)
+    is AndNode -> chain(node, variables, thenPart)
+    is OrNode -> nest(node, 0, variables)
+    else -> TODO("Unknown node type=$node")
+  }
+
+  /**
+   * TODO - this is how the following function should behave
+   * ```
+   * val result = call(1 + 2 + 3)
+   * ```
+   * Transformed to
+   * ```
+   * val result = run {
+   *   val tmp0 = 1 + 2
+   *   val tmp1 = tmp0 + 3
+   *   call(tmp1, <diagram>)
+   * }
+   * ```
+   */
+  private fun IrStatementsBuilder<*>.add(
+    node: ExpressionNode,
+    variables: List<IrTemporaryVariable>,
+    thenPart: IrStatementsBuilder<*>.(List<IrTemporaryVariable>) -> IrExpression
+  ): List<IrTemporaryVariable> {
+    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)
+
+    if (transformed.type == context.irBuiltIns.booleanType) {
+      // TODO irIfThenElse to support property nesting of ANDs & ORs
+      +irIfThen(irNot(transformed), thenPart(variables + transformer.variables))
+    } else {
+      +thenPart(variables + transformer.variables)
+    }
+
+    return transformer.variables
+  }
+
+  /**
+   * TODO - this is how the following function should behave
+   * ```
+   * val result = call(1 == 1 && 2 == 2)
+   * ```
+   * Transformed to
+   * ```
+   * val result = run {
+   *   val tmp0 = 1 == 1
+   *   if (tmp0) {
+   *     val tmp1 = 2 == 2
+   *     if (tmp1) call(true, <diagram>)
+   *     else call(false, <diagram>)
+   *   }
+   *   else call(false, <diagram>)
+   * }
+   * ```
+   */
+  private fun IrStatementsBuilder<*>.chain(
+    node: AndNode,
+    variables: List<IrTemporaryVariable>,
+    thenPart: IrStatementsBuilder<*>.(List<IrTemporaryVariable>) -> IrExpression
+  ): List<IrTemporaryVariable> {
+    val newVariables = mutableListOf<IrTemporaryVariable>()
+    for (child in node.children) {
+      newVariables.addAll(buildAssert(child, variables + newVariables, thenPart))
+    }
+    return newVariables
+  }
+
+  /**
+   * TODO - this is how the following function should behave
+   * ```
+   * val result = call(1 == 1 || 2 == 2)
+   * ```
+   * Transformed to
+   * ```
+   * val result = run {
+   *   val tmp0 = 1 == 1
+   *   if (tmp0) call(true, <diagram>)
+   *   else {
+   *     val tmp1 = 2 == 2
+   *     if (tmp1) call(true, <diagram>)
+   *     else call(false, <diagram>)
+   *   }
+   * }
+   * ```
+   */
+  private fun IrStatementsBuilder<*>.nest(
+    node: OrNode,
+    index: Int,
+    variables: List<IrTemporaryVariable>
+  ): List<IrTemporaryVariable> {
+    val children = node.children
+    val child = children[index]
+    buildAssert(child, variables) { newVariables ->
+      if (index + 1 == children.size) buildAssertThrow(newVariables)
+      else irBlock { nest(node, index + 1, newVariables) }
+    }
+    return emptyList()
+  }
+}

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

@@ -14,25 +14,31 @@
  * limitations under the License.
  */
 
-package com.bnorm.power
+package com.bnorm.power.diagram
 
 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
-import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
+import org.jetbrains.kotlin.ir.util.dumpKotlinLike
 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
+abstract class Node {
+  private val _children = mutableListOf<Node>()
+  val children: List<Node> get() = _children
 
-  protected fun dump(builder: StringBuilder, indent: Int) {
+  fun addChild(node: Node) {
+    _children.add(node)
+  }
+
+  fun dump(): String = buildString {
+    dump(this, 0)
+  }
+
+  private fun dump(builder: StringBuilder, indent: Int) {
     builder.append("  ".repeat(indent)).append(this).appendLine()
     for (child in children) {
       child.dump(builder, indent + 1)
@@ -40,56 +46,30 @@ sealed class Node {
   }
 }
 
-class AndNode(override val parent: Node) : Node() {
-  init {
-    parent.mutableChildren.add(this)
-  }
-
+class AndNode : Node() {
   override fun toString() = "AndNode"
 }
 
-class OrNode(override val parent: Node) : Node() {
-  init {
-    parent.mutableChildren.add(this)
-  }
-
+class OrNode : Node() {
   override fun toString() = "OrNode"
 }
 
-class ExpressionNode(
-  override val parent: Node
-) : Node() {
-  init {
-    parent.mutableChildren.add(this)
-  }
-
-  private val _expressions: MutableList<IrExpression> = mutableListOf()
+class ExpressionNode : Node() {
+  private val _expressions = mutableListOf<IrExpression>()
+  val expressions: List<IrExpression> = _expressions
 
   fun add(expression: IrExpression) {
     _expressions.add(expression)
   }
 
-  fun getExpressionsCopy(initialParent: IrDeclarationParent?): List<IrExpression> {
-    // Return a copy of all the expression by creating a deep copy of the head
-    // expression and running back through the assertion tree builder
-    val headCopy = _expressions.first().deepCopyWithSymbols(initialParent)
-    return (buildAssertTree(headCopy).children.single() as ExpressionNode)._expressions
-  }
-
-  override fun toString() = "ExpressionNode($_expressions)"
+  override fun toString() = "ExpressionNode(${_expressions.map { it.dumpKotlinLike() }})"
 }
 
-class RootNode : Node() {
-  override val parent: Node? = null
-
-  fun dump(): String = buildString {
-    dump(this, 0)
+fun buildTree(expression: IrExpression): Node? {
+  class RootNode : Node() {
+    override fun toString() = "RootNode"
   }
 
-  override fun toString() = "RootNode"
-}
-
-fun buildAssertTree(expression: IrExpression): RootNode {
   val tree = RootNode()
   expression.accept(object : IrElementVisitor<Unit, Node> {
     val INCREMENT_DECREMENT_OPERATORS = setOf(
@@ -104,14 +84,14 @@ fun buildAssertTree(expression: IrExpression): RootNode {
     }
 
     override fun visitExpression(expression: IrExpression, data: Node) {
-      val node = data as? ExpressionNode ?: ExpressionNode(data)
+      val node = data as? ExpressionNode ?: ExpressionNode().also { data.addChild(it) }
       node.add(expression)
       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)
+        val node = data as? ExpressionNode ?: ExpressionNode().also { data.addChild(it) }
         node.add(expression)
         return // Skip the internals of increment/decrement operations
       }
@@ -136,7 +116,7 @@ fun buildAssertTree(expression: IrExpression): RootNode {
       when (expression.origin) {
         IrStatementOrigin.ANDAND -> {
           // flatten `&&` expressions to be at the same level
-          val node = data as? AndNode ?: AndNode(data)
+          val node = data as? AndNode ?: AndNode().also { data.addChild(it) }
 
           require(expression.branches.size == 2)
           val thenBranch = expression.branches[0]
@@ -157,7 +137,7 @@ fun buildAssertTree(expression: IrExpression): RootNode {
         }
         IrStatementOrigin.OROR -> {
           // flatten `||` expressions to be at the same level
-          val node = data as? OrNode ?: OrNode(data)
+          val node = data as? OrNode ?: OrNode().also { data.addChild(it) }
 
           require(expression.branches.size == 2)
           val thenBranchCondition = expression.branches[0].condition
@@ -182,11 +162,11 @@ fun buildAssertTree(expression: IrExpression): RootNode {
         else -> {
           // Add as basic expression and terminate
           // TODO this has to be broken and not work in all cases...
-          ExpressionNode(data)
+          ExpressionNode().also { data.addChild(it) }
         }
       }
     }
   }, tree)
 
-  return tree
+  return tree.children.singleOrNull()
 }

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

@@ -14,8 +14,9 @@
  * limitations under the License.
  */
 
-package com.bnorm.power
+package com.bnorm.power.diagram
 
+import com.bnorm.power.irString
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.SourceRangeInfo
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
@@ -31,38 +32,26 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 
-data class IrStackVariable(
-  val temporary: IrVariable,
-  val original: IrExpression
-)
-
-data class ValueDisplay(
-  val value: IrVariable,
-  val indent: Int,
-  val row: Int,
-  val source: String
-)
-
-fun IrBuilderWithScope.buildMessage(
+fun IrBuilderWithScope.irDiagram(
   file: IrFile,
   fileSource: String,
-  title: IrExpression,
-  expression: IrExpression,
-  stack: List<IrStackVariable>
+  prefix: IrExpression? = null,
+  original: IrExpression,
+  variables: List<IrTemporaryVariable>
 ): IrExpression {
 
-  val originalInfo = file.info(expression)
+  val originalInfo = file.info(original)
   val callIndent = originalInfo.startColumnNumber
 
-  val stackValues = stack.map { it.toValueDisplay(fileSource, callIndent, file, originalInfo) }
+  val stackValues = variables.map { it.toValueDisplay(fileSource, callIndent, file, originalInfo) }
 
   val valuesByRow = stackValues.groupBy { it.row }
-  val rows = fileSource.substring(expression)
+  val rows = fileSource.substring(original)
     .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
     .split("\n")
 
   return irConcat().apply {
-    addArgument(title)
+    if (prefix != null) addArgument(prefix)
 
     for ((row, rowSource) in rows.withIndex()) {
       val rowValues = valuesByRow[row]?.let { values -> values.sortedBy { it.indent } } ?: emptyList()
@@ -102,7 +91,14 @@ fun IrBuilderWithScope.buildMessage(
   }
 }
 
-private fun IrStackVariable.toValueDisplay(
+private data class ValueDisplay(
+  val value: IrVariable,
+  val indent: Int,
+  val row: Int,
+  val source: String
+)
+
+private fun IrTemporaryVariable.toValueDisplay(
   fileSource: String,
   callIndent: Int,
   file: IrFile,

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

@@ -0,0 +1,33 @@
+package com.bnorm.power.diagram
+
+import org.jetbrains.kotlin.ir.builders.IrStatementsBuilder
+import org.jetbrains.kotlin.ir.builders.irGet
+import org.jetbrains.kotlin.ir.builders.irTemporary
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+
+data class IrTemporaryVariable(
+  val temporary: IrVariable,
+  val original: IrExpression
+)
+
+class IrTemporaryExtractionTransformer(
+  private val builder: IrStatementsBuilder<*>,
+  private val transform: Set<IrExpression>
+) : IrElementTransformerVoid() {
+  private val _variables = mutableListOf<IrTemporaryVariable>()
+  val variables: List<IrTemporaryVariable> = _variables
+
+  override fun visitExpression(expression: IrExpression): IrExpression {
+    return if (expression in transform) {
+      val copy = expression.deepCopyWithSymbols(builder.scope.getLocalDeclarationParent())
+      val variable = builder.irTemporary(super.visitExpression(expression))
+      _variables.add(IrTemporaryVariable(variable, copy))
+      builder.irGet(variable)
+    } else {
+      super.visitExpression(expression)
+    }
+  }
+}