From 830caa71596caf488b7717b39ba2a4e535ddb779 Mon Sep 17 00:00:00 2001
From: Markus Alexander Kuppe <github.com@lemmster.de>
Date: Wed, 28 Aug 2024 13:28:50 -0700
Subject: [PATCH] ExprOptimizer: Eliminate `SUBSET` operator when it occurs in
 subexpressions:
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

```tla
----- MODULE ljasf -----
EXTENDS Integers

VARIABLE
    \* @type: { v: Set( { p: Set( { v: Set(Int) } ) } ) };
    m

TypeOK ==
        m \in [ v : SUBSET [ p : SUBSET [ v: SUBSET {1,2,3,4,5}] ] ]

Init ==
    m = [ v |-> {[p |-> {[v |-> {1}]} ]} ]

Next ==
    UNCHANGED m
====
```

```
input:
m["v"] ∈ SUBSET ({["p" ↦ t_5$1] : t_5$1 ∈ SUBSET ({["v" ↦ t_4$1] : t_4$1 ∈ SUBSET {1, 2, 3, 4, 5}})})

old:
∀t_a ∈ m["v"]: ((DOMAIN t_a = {"p"}) ∧ (∀t_b ∈ t_a["p"]: ((DOMAIN t_b = {"v"}) ∧ t_b["v"] ∈ SUBSET {1, 2, 3, 4, 5})))

new:
∀t_a ∈ m["v"]: ((DOMAIN t_a = {"p"}) ∧ (∀t_b ∈ t_a["p"]: ((DOMAIN t_b = {"v"}) ∧ (∀t_c ∈ t_b["v"]: t_c ∈ {1, 2, 3, 4, 5}))))
```
---
 .unreleased/features/rewrite.md                     |  1 +
 .../at/forsyte/apalache/tla/pp/ExprOptimizer.scala  |  9 +++++++--
 .../forsyte/apalache/tla/pp/TestExprOptimizer.scala | 13 +++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100644 .unreleased/features/rewrite.md

diff --git a/.unreleased/features/rewrite.md b/.unreleased/features/rewrite.md
new file mode 100644
index 0000000000..42c7985991
--- /dev/null
+++ b/.unreleased/features/rewrite.md
@@ -0,0 +1 @@
+Handle expressions such as  S \in SUBSET T  in ExprOptimizer by rewriting the expression into  \A r \in S: r \in T
\ No newline at end of file
diff --git a/tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/ExprOptimizer.scala b/tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/ExprOptimizer.scala
index 9faa5d1b68..f296d447eb 100644
--- a/tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/ExprOptimizer.scala
+++ b/tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/ExprOptimizer.scala
@@ -93,7 +93,7 @@ class ExprOptimizer(nameGen: UniqueNameGenerator, tracker: TransformationTracker
         val b = BoolT1
         val domEq = tla.eql(tla.dom(rec).as(strSetT), tla.enumSet(fields: _*).as(strSetT)).as(b)
         val fieldsEq = fields.zip(values.zip(sets)).map { case (key, (value, set)) =>
-          tla.in(tla.appFun(rec, key).as(value.typeTag.asTlaType1()), set).as(b)
+          apply(tla.in(tla.appFun(rec, key).as(value.typeTag.asTlaType1()), set).as(b))
         }
         apply(tla.and(domEq +: fieldsEq: _*).as(b))
       }
@@ -121,10 +121,15 @@ class ExprOptimizer(nameGen: UniqueNameGenerator, tracker: TransformationTracker
         val domEq = tla.eql(tla.dom(r).as(SetT1(domType)), tla.enumSet(fields: _*).as(strSetT)).as(b)
 
         val fieldsEq = fields.zip(values.zip(sets)).map { case (key, (value, set)) =>
-          tla.in(tla.appFun(r, key).as(value.typeTag.asTlaType1()), set).as(b)
+          apply(tla.in(tla.appFun(r, key).as(value.typeTag.asTlaType1()), set).as(b))
         }
         apply(tla.forall(r, setRec, tla.and(domEq +: fieldsEq: _*).as(b)).as(b))
       }
+
+    // S ∈ SUBSET T ~~> ∀ s ∈ S: s ∈ T
+    case OperEx(TlaSetOper.in, lSet, OperEx(TlaSetOper.powerset, rSet)) =>
+      val lMem = tla.name(nameGen.newName()).as(getElemType(lSet))
+      apply(tla.forall(lMem, lSet, tla.in(lMem, apply(rSet)).as(BoolT1)).as(BoolT1))
   }
 
   /**
diff --git a/tla-pp/src/test/scala/at/forsyte/apalache/tla/pp/TestExprOptimizer.scala b/tla-pp/src/test/scala/at/forsyte/apalache/tla/pp/TestExprOptimizer.scala
index a38e77fefc..bd673a5313 100644
--- a/tla-pp/src/test/scala/at/forsyte/apalache/tla/pp/TestExprOptimizer.scala
+++ b/tla-pp/src/test/scala/at/forsyte/apalache/tla/pp/TestExprOptimizer.scala
@@ -184,6 +184,19 @@ class TestExprOptimizer extends AnyFunSuite with BeforeAndAfterEach {
     assert(expected == output)
   }
 
+  test("""S \in SUBSET T ~~> \A s \in S: s \in T""") {
+    val T = name("T").as(intSetT)
+    val S = name("S").as(intSetT)
+    val powSetT = powSet(T).as(intSetT)
+    val input = in(S, powSetT).as(boolT)
+    val output = optimizer.apply(input)
+
+    val s = name("t_1").as(intT)
+    val expected = forall(s, S, in(s, T).as(boolT)).as(boolT)
+
+    assert(expected == output)
+  }
+
   // optimizations for set cardinalities
 
   test("""Cardinality(S) = 0 becomes S = {}""") {