From ad51a39b770aa6b72dfb5815e40349643db6f278 Mon Sep 17 00:00:00 2001
From: mario-bucev <mario.bucev@epfl.ch>
Date: Thu, 29 Feb 2024 16:50:45 +0100
Subject: [PATCH] Fix ImperativeCleanup cleaning too much (#1500)

---
 .../imperative/ImperativeCleanup.scala        | 34 +++++++++++--------
 .../verification/invalid/TupleCleanup.scala   | 21 ++++++++++++
 2 files changed, 40 insertions(+), 15 deletions(-)
 create mode 100644 frontends/benchmarks/verification/invalid/TupleCleanup.scala

diff --git a/core/src/main/scala/stainless/extraction/imperative/ImperativeCleanup.scala b/core/src/main/scala/stainless/extraction/imperative/ImperativeCleanup.scala
index 39ce60a29..908718e03 100644
--- a/core/src/main/scala/stainless/extraction/imperative/ImperativeCleanup.scala
+++ b/core/src/main/scala/stainless/extraction/imperative/ImperativeCleanup.scala
@@ -54,21 +54,25 @@ class ImperativeCleanup(override val s: Trees, override val t: oo.Trees)
     object ReconstructTuple {
       def unapply(e: s.Expr): Option[s.Expr] = e match {
         case s.Let(vd, tuple, Lets(lets, s.Tuple(es))) =>
-          val letsMap = lets.map { case (vd, e) => (vd.id, e) }.toMap
-          if (
-            vd.getType.isInstanceOf[s.TupleType] &&
-            es.length == vd.getType.asInstanceOf[s.TupleType].bases.length &&
-            es.zipWithIndex.forall {
-              case (e0 : s.Variable, i) =>
-                letsMap.contains(e0.id) &&
-                letsMap(e0.id) == s.TupleSelect(vd.toVariable, i + 1)
-              case (e0, i) =>
-                e0 == s.TupleSelect(vd.toVariable, i + 1)
-            }
-          )
-            Some(tuple)
-          else
-            None
+          vd.getType match {
+            case s.TupleType(bases) =>
+              val letsMap = lets.map { case (vd, e) => (vd.id, e) }.toMap
+              // All let-bindings are used in the "returned" value `es`
+              lazy val returnedLets = letsMap.keySet.filter(id => es.exists { case v: s.Variable => v.id == id; case _ => false })
+              // All values in `es` are reconstruction of selections of `vd`
+              lazy val esIsTupleSel = es.zipWithIndex.forall {
+                case (e0: s.Variable, i) =>
+                  letsMap.contains(e0.id) &&
+                  letsMap(e0.id) == s.TupleSelect(vd.toVariable, i + 1)
+                case (e0, i) =>
+                  e0 == s.TupleSelect(vd.toVariable, i + 1)
+              }
+              if (es.length == bases.length && returnedLets.size == letsMap.size && esIsTupleSel)
+                Some(tuple)
+              else None
+            case _ =>
+              None
+          }
 
         case s.Let(vd, e, Lets(Seq(), v)) if v == vd.toVariable =>
           Some(e)
diff --git a/frontends/benchmarks/verification/invalid/TupleCleanup.scala b/frontends/benchmarks/verification/invalid/TupleCleanup.scala
new file mode 100644
index 000000000..448b5bf81
--- /dev/null
+++ b/frontends/benchmarks/verification/invalid/TupleCleanup.scala
@@ -0,0 +1,21 @@
+import stainless.lang._
+import stainless.collection._
+
+object TupleCleanup {
+  def test1(x: BigInt, y: BigInt): (BigInt, BigInt) = {
+    val (xx, yy) = (x, y)
+    val boom = Nil[BigInt]().head
+    (xx, yy)
+  }
+
+  def test2(xs: List[BigInt]): (List[BigInt], List[BigInt]) = {
+    decreases(xs)
+    xs match {
+      case Nil() => (Nil(), Nil())
+      case Cons(i, t) =>
+        val (s2, g2) = test2(t)
+        val boom = Nil[BigInt]().head
+        (s2, g2)
+    }
+  }
+}
\ No newline at end of file