Core: avoid explosive flattening

UnitTestBot · Nov 1, 2023 · a4dd431 · a4dd431
1 parent 856455d
commit a4dd431
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 12 deletions.
diff --git a/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KArrayExprSimplifier.kt b/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KArrayExprSimplifier.kt
@@ -953,7 +953,7 @@ interface KArrayExprSimplifier : KExprSimplifierBase {
         val original: KExpr<A>,
         val base: KExpr<A>,
         val values: List<KExpr<R>>
-    ) : KExpr<A>(ctx) {
+    ) : KExpr<A>(ctx), KSimplifierAuxExpr {
         override val sort: A
             get() = base.sort
 
@@ -1146,7 +1146,7 @@ interface KArrayExprSimplifier : KExprSimplifierBase {
     private sealed class SimplifierArraySelectBaseExpr<A : KArraySortBase<R>, R : KSort>(
         ctx: KContext,
         val array: KExpr<A>
-    ) : KExpr<R>(ctx) {
+    ) : KExpr<R>(ctx), KSimplifierAuxExpr {
         override val sort: R
             get() = array.sort.range
 
@@ -1217,7 +1217,7 @@ interface KArrayExprSimplifier : KExprSimplifierBase {
     private sealed class SelectFromStoreExprBase<A : KArraySortBase<R>, R : KSort>(
         ctx: KContext,
         array: KArrayStoreBase<A, R>
-    ) : KExpr<R>(ctx) {
+    ) : KExpr<R>(ctx), KSimplifierAuxExpr {
         val storeValue: KExpr<R> = array.value
 
         override val sort: R = storeValue.sort

diff --git a/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KBoolExprSimplifier.kt b/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KBoolExprSimplifier.kt
@@ -15,6 +15,7 @@ import io.ksmt.expr.KXorExpr
 import io.ksmt.expr.transformer.KTransformerBase
 import io.ksmt.sort.KBoolSort
 import io.ksmt.sort.KSort
+import java.util.IdentityHashMap
 
 interface KBoolExprSimplifier : KExprSimplifierBase {
 
@@ -213,7 +214,7 @@ interface KBoolExprSimplifier : KExprSimplifierBase {
         val condition: KExpr<KBoolSort>,
         val trueBranch: KExpr<T>,
         val falseBranch: KExpr<T>
-    ) : KApp<T, KBoolSort>(ctx) {
+    ) : KApp<T, KBoolSort>(ctx), KSimplifierAuxExpr {
 
         override val decl: KDecl<T>
             get() = ctx.mkIteDecl(trueBranch.sort)
@@ -240,7 +241,7 @@ interface KBoolExprSimplifier : KExprSimplifierBase {
         val simplifiedCondition: KExpr<KBoolSort>,
         val trueBranch: KExpr<T>,
         val falseBranch: KExpr<T>
-    ) : KApp<T, T>(ctx) {
+    ) : KApp<T, T>(ctx), KSimplifierAuxExpr {
 
         override val decl: KDecl<T>
             get() = ctx.mkIteDecl(trueBranch.sort)
@@ -272,7 +273,7 @@ interface KBoolExprSimplifier : KExprSimplifierBase {
         val neutralElement: KExpr<KBoolSort>,
         val zeroElement: KExpr<KBoolSort>,
         override val args: List<KExpr<KBoolSort>>
-    ) : KApp<KBoolSort, KBoolSort>(ctx) {
+    ) : KApp<KBoolSort, KBoolSort>(ctx), KSimplifierAuxExpr {
         override val sort: KBoolSort = ctx.boolSort
         override val decl: KDecl<KBoolSort>
             get() = when (operation) {
@@ -287,6 +288,7 @@ interface KBoolExprSimplifier : KExprSimplifierBase {
 
         val simplifiedArgs = arrayListOf<KExpr<KBoolSort>>()
 
+        private val visitedArgs = IdentityHashMap<KExpr<*>, Unit>()
         private val argsIteratorStack = arrayListOf(args.iterator())
         private var currentArgument: KExpr<KBoolSort>? = null
 
@@ -309,6 +311,9 @@ interface KBoolExprSimplifier : KExprSimplifierBase {
 
             while (hasUnprocessedArgument()) {
                 val argument = argsIteratorStack.last().next()
+
+                if (visitedArgs.put(argument, Unit) != null) continue
+
                 if (!tryFlatExpr(argument)) {
                     currentArgument = argument
                     return

diff --git a/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KBvExprSimplifier.kt b/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KBvExprSimplifier.kt
@@ -1083,7 +1083,7 @@ interface KBvExprSimplifier : KExprSimplifierBase {
     private class SimplifierFlatBvAddExpr<T : KBvSort>(
         ctx: KContext,
         override val args: List<KExpr<T>>
-    ) : KApp<T, T>(ctx) {
+    ) : KApp<T, T>(ctx), KSimplifierAuxExpr {
 
         override val decl: KDecl<T>
             get() = ctx.mkBvAddDecl(sort, sort)
@@ -1104,7 +1104,7 @@ interface KBvExprSimplifier : KExprSimplifierBase {
     private class SimplifierFlatBvMulExpr<T : KBvSort>(
         ctx: KContext,
         override val args: List<KExpr<T>>
-    ) : KApp<T, T>(ctx) {
+    ) : KApp<T, T>(ctx), KSimplifierAuxExpr {
 
         override val decl: KDecl<T>
             get() = ctx.mkBvMulDecl(sort, sort)
@@ -1125,7 +1125,7 @@ interface KBvExprSimplifier : KExprSimplifierBase {
     private class SimplifierFlatBvOrExpr<T : KBvSort>(
         ctx: KContext,
         override val args: List<KExpr<T>>
-    ) : KApp<T, T>(ctx) {
+    ) : KApp<T, T>(ctx), KSimplifierAuxExpr {
 
         override val decl: KDecl<T>
             get() = ctx.mkBvOrDecl(sort, sort)
@@ -1146,7 +1146,7 @@ interface KBvExprSimplifier : KExprSimplifierBase {
     private class SimplifierFlatBvAndExpr<T : KBvSort>(
         ctx: KContext,
         override val args: List<KExpr<T>>
-    ) : KApp<T, T>(ctx) {
+    ) : KApp<T, T>(ctx), KSimplifierAuxExpr {
 
         override val decl: KDecl<T>
             get() = ctx.mkBvAndDecl(sort, sort)
@@ -1167,7 +1167,7 @@ interface KBvExprSimplifier : KExprSimplifierBase {
     private class SimplifierFlatBvXorExpr<T : KBvSort>(
         ctx: KContext,
         override val args: List<KExpr<T>>
-    ) : KApp<T, T>(ctx) {
+    ) : KApp<T, T>(ctx), KSimplifierAuxExpr {
 
         override val decl: KDecl<T>
             get() = ctx.mkBvXorDecl(sort, sort)
@@ -1189,7 +1189,7 @@ interface KBvExprSimplifier : KExprSimplifierBase {
         ctx: KContext,
         override val sort: KBvSort,
         override val args: List<KExpr<KBvSort>>
-    ) : KApp<KBvSort, KBvSort>(ctx) {
+    ) : KApp<KBvSort, KBvSort>(ctx), KSimplifierAuxExpr {
 
         // We have no decl, but we don't care since decl is unused
         override val decl: KDecl<KBvSort>

diff --git a/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KExprSimplifier.kt b/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KExprSimplifier.kt
@@ -343,6 +343,10 @@ open class KExprSimplifier(override val ctx: KContext) :
         val result = transformedExpr(rewriteFrame.rewritten)
             ?: error("Nested rewrite failed")
 
+        if (rewriteFrame.rewritten is KSimplifierAuxExpr) {
+            eraseTransformationResult(rewriteFrame.rewritten)
+        }
+
         return result.uncheckedCast()
     }
 

diff --git a/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KExprSimplifierBase.kt b/ksmt-core/src/main/kotlin/io/ksmt/expr/rewrite/simplify/KExprSimplifierBase.kt
@@ -46,6 +46,11 @@ interface KExprSimplifierBase : KTransformer {
 @JvmInline
 internal value class SimplifierAuxExpression<T : KSort>(val expr: KExpr<T>)
 
+/**
+ * Mark simplifier auxiliary expressions.
+ * */
+internal interface KSimplifierAuxExpr
+
 internal inline fun <T : KSort> auxExpr(builder: () -> KExpr<T>): SimplifierAuxExpression<T> =
     SimplifierAuxExpression(builder())
 

diff --git a/ksmt-core/src/main/kotlin/io/ksmt/expr/transformer/KNonRecursiveTransformerBase.kt b/ksmt-core/src/main/kotlin/io/ksmt/expr/transformer/KNonRecursiveTransformerBase.kt
@@ -128,6 +128,14 @@ abstract class KNonRecursiveTransformerBase: KTransformer {
         return transformed[expr] as? KExpr<T>
     }
 
+    /**
+     * Erase [expr] transformation result.
+     * Useful for transformer auxiliary expressions.
+     * */
+    fun eraseTransformationResult(expr: KExpr<*>) {
+        transformed.remove(expr)
+    }
+
     /**
      * Allows to skip expression transformation and stop deepening.
      * */