From 6a5c01bea4647b212864f3ea2c3c4b85cb535236 Mon Sep 17 00:00:00 2001
From: Brian Norman <brian@bnorm.com>
Date: Sat, 8 Feb 2020 14:12:23 -0600
Subject: [PATCH] Make sure expression are only called once when building stack

---
 .../main/kotlin/com/bnorm/power/AssertTree.kt  | 18 +++++++++++++++---
 .../bnorm/power/PowerAssertCallTransformer.kt  | 12 +++++++++---
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/kotlin-power-assert/src/main/kotlin/com/bnorm/power/AssertTree.kt b/kotlin-power-assert/src/main/kotlin/com/bnorm/power/AssertTree.kt
index d69698f93b1..923706e9a30 100644
--- a/kotlin-power-assert/src/main/kotlin/com/bnorm/power/AssertTree.kt
+++ b/kotlin-power-assert/src/main/kotlin/com/bnorm/power/AssertTree.kt
@@ -16,6 +16,7 @@
 
 package com.bnorm.power
 
+import org.jetbrains.kotlin.backend.common.deepCopyWithVariables
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrConst
@@ -60,9 +61,20 @@ class ExpressionNode(
     parent.mutableChildren.add(this)
   }
 
-  val expressions: MutableList<IrExpression> = mutableListOf()
+  private val _expressions: MutableList<IrExpression> = mutableListOf()
 
-  override fun toString() = "ExpressionNode($expressions)"
+  fun add(expression: IrExpression) {
+    _expressions.add(expression)
+  }
+
+  fun getExpressionsCopy(): 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().deepCopyWithVariables()
+    return (buildAssertTree(headCopy).children.single() as ExpressionNode)._expressions
+  }
+
+  override fun toString() = "ExpressionNode($_expressions)"
 }
 
 class RootNode : Node() {
@@ -84,7 +96,7 @@ fun buildAssertTree(expression: IrExpression): RootNode {
 
     override fun visitExpression(expression: IrExpression, data: Node) {
       val node = data as? ExpressionNode ?: ExpressionNode(data)
-      node.expressions += expression
+      node.add(expression)
       expression.acceptChildren(this, node)
     }
 
diff --git a/kotlin-power-assert/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt b/kotlin-power-assert/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt
index 9f826a53fb2..df36b55e386 100644
--- a/kotlin-power-assert/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt
+++ b/kotlin-power-assert/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt
@@ -114,6 +114,7 @@ class PowerAssertCallTransformer(
       return irBlock {
         buildAssert(this@PowerAssertCallTransformer.context, file, fileSource, callSource, callIndent, title, root)
       }
+//        .also { println(it.dump())}
     }
   }
 }
@@ -161,8 +162,9 @@ private inline fun IrBlockBuilder.irIfNotThan(
   node: ExpressionNode,
   thenPart: IrBlockBuilder.(subStack: MutableList<IrStackVariable>) -> IrExpression
 ): IrWhen {
-  val stackTransformer = StackBuilder(this, stack, file, fileSource, callIndent, node.expressions)
-  val transformed = node.expressions.first().transform(stackTransformer, null)
+  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()))
 }
 
@@ -197,7 +199,11 @@ class StackBuilder(
   }
 
   override fun visitExpression(expression: IrExpression): IrExpression {
-    return super.visitExpression(expression).also { if (expression in transform) push(it) }
+    return if (expression in transform) {
+      push(super.visitExpression(expression))
+    } else {
+      super.visitExpression(expression)
+    }
   }
 }