From 6576959a263bafbb8c6a11345729fde51d8bd915 Mon Sep 17 00:00:00 2001
From: Kukovec <jure.kukovec@gmail.com>
Date: Tue, 14 Nov 2023 15:45:09 +0100
Subject: [PATCH] Throw error when attempting to construct a UNION of function
 sets (#2778)

* Expansion marker fix

* missing file

* test elaboration
---
 test/tla/Bug2772.tla                          | 11 ++++++++++
 test/tla/cli-integration-tests.md             | 21 +++++++++++++++++++
 .../tla/bmcmt/analyses/ExpansionMarker.scala  |  8 +++----
 3 files changed, 36 insertions(+), 4 deletions(-)
 create mode 100644 test/tla/Bug2772.tla

diff --git a/test/tla/Bug2772.tla b/test/tla/Bug2772.tla
new file mode 100644
index 0000000000..70ae8eb38f
--- /dev/null
+++ b/test/tla/Bug2772.tla
@@ -0,0 +1,11 @@
+---- MODULE Bug2772 ----
+EXTENDS Integers, Apalache
+
+ErrInv == [x \in {0} |-> 1] \in UNION { [d -> {1}] : d \in {{0}} }
+
+OkInv == \E d \in {{0}}: [x \in {0} |-> 1] \in [d -> {1}]
+
+Init == TRUE
+
+Next == TRUE
+====
diff --git a/test/tla/cli-integration-tests.md b/test/tla/cli-integration-tests.md
index 52789aecde..958054cb4a 100644
--- a/test/tla/cli-integration-tests.md
+++ b/test/tla/cli-integration-tests.md
@@ -2332,6 +2332,27 @@ $ apalache-mc check --config=Test2750.cfg Test2750.tla | sed 's/[IEW]@.*//'
 EXITCODE: OK
 ```
 
+### check Bug2772.tla errors on unsupported syntax
+
+It should not be possible to pass input, which would require function set expansion without triggering an exception.
+
+```sh
+$ apalache-mc check --inv=ErrInv --length=1 Bug2772.tla | sed 's/[IEW]@.*//'
+...
+EXITCODE: ERROR (255)
+```
+
+### check Bug2772.tla succeeds on supported syntax
+
+However, if one uses a semantically equivalent (syntactically different) expression, where the function set is not forced to expand, it should pass.
+
+```sh
+$ apalache-mc check --inv=OkInv --length=1 Bug2772.tla | sed 's/[IEW]@.*//'
+...
+EXITCODE: OK
+```
+
+
 ### check profiling
 
 Check that the profiler output is produced as explained in
diff --git a/tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/analyses/ExpansionMarker.scala b/tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/analyses/ExpansionMarker.scala
index 2ef3c09b6e..fb7f3edd5e 100644
--- a/tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/analyses/ExpansionMarker.scala
+++ b/tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/analyses/ExpansionMarker.scala
@@ -63,7 +63,7 @@ class ExpansionMarker @Inject() (tracker: TransformationTracker) extends TlaExTr
     case ex @ OperEx(op @ TlaOper.chooseBounded, name, set, pred) =>
       // CHOOSE does not require the set to be expanded
       val tag = ex.typeTag
-      OperEx(op, name, transform(false)(set), transform(false)(pred))(tag)
+      OperEx(op, name, transform(false)(set), transform(shallExpand)(pred))(tag)
 
     case ex @ OperEx(op @ ApalacheOper.guess, set) =>
       // Guess does not require the set to be expanded
@@ -73,7 +73,7 @@ class ExpansionMarker @Inject() (tracker: TransformationTracker) extends TlaExTr
     case ex @ OperEx(op, name, set, pred) if op == TlaBoolOper.exists || op == TlaBoolOper.forall =>
       // non-skolemizable quantifiers require their sets to be expanded
       val tag = ex.typeTag
-      OperEx(op, name, transform(true)(set), transform(false)(pred))(tag)
+      OperEx(op, name, transform(true)(set), transform(shallExpand)(pred))(tag)
 
     case ex @ OperEx(op, elem, set)
         if op == TlaSetOper.in || op == ApalacheInternalOper.selectInSet || op == ApalacheInternalOper.selectInFun ||
@@ -91,11 +91,11 @@ class ExpansionMarker @Inject() (tracker: TransformationTracker) extends TlaExTr
     case ex @ OperEx(op @ TlaSetOper.filter, name, set, pred) =>
       // For the moment, we require the set to be expanded. However, we could think of collecting constraints on the way.
       val tag = ex.typeTag
-      OperEx(op, name, transform(true)(set), transform(false)(pred))(tag)
+      OperEx(op, name, transform(true)(set), transform(shallExpand)(pred))(tag)
 
     case ex @ OperEx(op, body, args @ _*)
         if op == TlaSetOper.map || op == TlaFunOper.funDef || op == TlaFunOper.recFunDef =>
-      val tbody: TlaEx = transform(false)(body)
+      val tbody: TlaEx = transform(shallExpand)(body)
       val targs = args.map(transform(true))
       val tag = ex.typeTag
       OperEx(op, tbody +: targs: _*)(tag)