Added transformations for quantification over special integer sets

tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/IntegerQuantificationOptimizer.scala

@@ -0,0 +1,78 @@
+package at.forsyte.apalache.tla.pp
+
+import at.forsyte.apalache.tla.lir._
+import at.forsyte.apalache.tla.lir.oper._
+import at.forsyte.apalache.tla.lir.transformations.{TlaExTransformation, TransformationTracker}
+import at.forsyte.apalache.tla.lir.values.{TlaInt, TlaIntSet, TlaNatSet}
+
+/**
+ * This transformation rewrites expressions of the form `\E x \in S: P`, where S is one of `Int`, `Nat`, `a..b`, into
+ * `\E x: Q` (unbounded quantification) and `\A x \in S: P` into `\A x: R`.
+ *
+ * `Q` has the shape `T /\ P`, and `R` has the shape `T => P`, where `T` depends on `S`:
+ *   - If `S` is `Int`, `T` is just `TRUE` (and we can say `Q = R = P`)
+ *   - If `S` is `Nat`, `T` is `x >= 0`
+ *   - If `S` is `a..b`, `T` is `a <= x <= b`
+ *
+ * This is the only place we support unbounded quantification, as we can leverage the type annotation of `x` (which
+ * should be IntT1()). We can then introduce a special rewriting rule for this quantification form.
+ *
+ * @author
+ *   Jure Kukovec
+ */
+
+class IntegerQuantificationOptimizer(tracker: TransformationTracker) extends TlaExTransformation {
+
+  private val intTag = Typed(IntT1())
+  private val boolTag = Typed(BoolT1())
+
+  override def apply(expr: TlaEx): TlaEx = {
+    transform(expr)
+  }
+
+  // This optimization only applies to Int, Nat and a..b sets (of integers)
+  private def isAllowedSet(s: TlaEx): Boolean = s match {
+    case ValEx(TlaIntSet | TlaNatSet)      => true
+    case OperEx(TlaArithOper.dotdot, _, _) => true
+    case _                                 => false
+  }
+
+  // We pass name as String in the next two methods, because re-creating
+  // NameEx(name)(intTag) constructs an expression with a fresh UID, thus preserving
+  // the property of UID uniqueness. If we had just copied nameEx: NameEx, all the instances
+  // of nameEx would have had the same UID.
+
+  // name >= 0
+  private def ge0(name: String): TlaEx =
+    OperEx(TlaArithOper.ge, NameEx(name)(intTag), ValEx(TlaInt(0))(intTag))(Typed(BoolT1()))
+
+  // lowerBound <= name <= lowerBound
+  private def inRange(name: String, lowerBound: TlaEx, upperBound: TlaEx): TlaEx = {
+    val lConstraint = OperEx(TlaArithOper.le, lowerBound, NameEx(name)(intTag))(boolTag)
+    val uConstraint = OperEx(TlaArithOper.le, NameEx(name)(intTag), upperBound)(boolTag)
+    OperEx(TlaBoolOper.and, lConstraint, uConstraint)(boolTag)
+  }
+
+  def transform: TlaExTransformation = tracker.trackEx {
+    // Assume typechecking succeeded -> name has type Int
+    case OperEx(oper @ (TlaBoolOper.forall | TlaBoolOper.exists), nameEx @ NameEx(name), setEx, pred)
+        if isAllowedSet(setEx) =>
+      // If we have \A, then T => P, otherwise T /\ P
+      val (unboundedOper, outerPred) = oper match {
+        case TlaBoolOper.forall => (TlaBoolOper.forallUnbounded, TlaBoolOper.implies)
+        case TlaBoolOper.exists => (TlaBoolOper.existsUnbounded, TlaBoolOper.and)
+      }
+
+      // T depends on the shape of S
+      val newPred = setEx match {
+        case ValEx(TlaIntSet) => pred
+        case ValEx(TlaNatSet) => OperEx(outerPred, ge0(name), pred)(boolTag)
+        case OperEx(TlaArithOper.dotdot, lowerBound, upperBound) =>
+          OperEx(outerPred, inRange(name, lowerBound, upperBound), pred)(boolTag)
+      }
+      // nameEx can be copied, and because this is the only place we do so, it will still have a unique ID
+      OperEx(unboundedOper, nameEx, newPred)(boolTag)
+
+    case ex => ex
+  }
+}

tla-pp/src/test/scala/at/forsyte/apalache/tla/pp/TestQuantificationOptimizer.scala

@@ -0,0 +1,101 @@
+package at.forsyte.apalache.tla.pp
+
+import at.forsyte.apalache.tla.lir.{Typed, _}
+import at.forsyte.apalache.tla.lir.convenience.tla
+import at.forsyte.apalache.tla.lir.transformations.impl.IdleTracker
+import org.junit.runner.RunWith
+import at.forsyte.apalache.tla.lir.TypedPredefs._
+import at.forsyte.apalache.tla.lir.oper.{TlaArithOper, TlaBoolOper, TlaOper}
+import at.forsyte.apalache.tla.lir.values.{TlaIntSet, TlaNatSet}
+import org.scalacheck.Gen
+import org.scalacheck.Prop.{forAll, AnyOperators}
+import org.scalatest.funsuite.AnyFunSuite
+import org.scalatestplus.junit.JUnitRunner
+import org.scalatestplus.scalacheck.Checkers
+
+@RunWith(classOf[JUnitRunner])
+class TestQuantificationOptimizer extends AnyFunSuite with Checkers {
+
+  val optimizer = new IntegerQuantificationOptimizer(new IdleTracker)
+
+  val gens: IrGenerators = new IrGenerators {
+    override val maxArgs: Int = 3
+  }
+
+  val ops =
+    gens.simpleOperators ++ gens.logicOperators ++ gens.arithOperators ++ gens.setOperators ++ gens.functionOperators
+
+  val rangeGen: Gen[TlaEx] = for {
+    a <- gens.genTlaEx(ops)(gens.emptyContext)
+    b <- gens.genTlaEx(ops)(gens.emptyContext)
+  } yield tla.dotdot(a.as(IntT1()), b.as(IntT1())).as(SetT1(IntT1()))
+
+  val boolTag = Typed(BoolT1())
+  val varname = tla.name("x").as(IntT1())
+
+  def varnameInBounds(lowerBound: TlaEx, upperBound: TlaEx): TlaEx = {
+    val lConstraint = tla.le(lowerBound, varname).as(BoolT1())
+    val uConstraint = tla.le(varname, upperBound).as(BoolT1())
+    tla.and(lConstraint, uConstraint).as(BoolT1())
+  }
+
+  val varnameGe0 = tla.ge(varname, tla.int(0).as(IntT1())).as(BoolT1())
+
+  def newQuantifierPredicate(oper: TlaOper, setEx: TlaEx, ex: TlaEx): TlaEx = {
+    // T /\ P for \E, T => P for \A
+    val outerOperatorCtor: (TlaEx, TlaEx) => BuilderEx =
+      if (oper == TlaBoolOper.exists) tla.and(_, _) else tla.impl(_, _)
+
+    setEx match {
+      case ValEx(TlaIntSet) => ex
+      case ValEx(TlaNatSet) => outerOperatorCtor(varnameGe0, ex).as(BoolT1())
+      case OperEx(TlaArithOper.dotdot, lowerBound, upperBound) =>
+        outerOperatorCtor(varnameInBounds(lowerBound, upperBound), ex).as(BoolT1())
+      case _ => ex
+    }
+  }
+
+  // Checks that the actual translation matches expectations, based on the shape of the set
+  def withSetGen(predefSetOrRangeGen: Gen[TlaEx]): Unit = {
+    val pairGen = for {
+      predefSetOrRangeEx <- predefSetOrRangeGen
+      exRandomTag <- gens.genTlaEx(ops)(gens.emptyContext)
+    } yield (predefSetOrRangeEx, exRandomTag)
+    val prop = forAll(pairGen) { case (setEx, exWithRandomTag) =>
+      // gens produces garbage w.r.t. type tags, we just force the toplevel ex to be boolean.
+      val ex = exWithRandomTag.withTag(boolTag)
+      val existsEx = tla.exists(varname, setEx, ex).as(BoolT1())
+      val forallEx = tla.forall(varname, setEx, ex).as(BoolT1())
+
+      val expected = {
+        val expectedPredicateExists = newQuantifierPredicate(TlaBoolOper.exists, setEx, ex)
+        val expectedPredicateForall = newQuantifierPredicate(TlaBoolOper.forall, setEx, ex)
+
+        val existsExpected: TlaEx = OperEx(TlaBoolOper.existsUnbounded, varname, expectedPredicateExists)(boolTag)
+        val forallExpected: TlaEx = OperEx(TlaBoolOper.forallUnbounded, varname, expectedPredicateForall)(boolTag)
+        (existsExpected, forallExpected)
+      }
+
+      val actual = (optimizer(existsEx), optimizer(forallEx))
+
+      expected =? actual
+    }
+    check(prop, minSuccessful(500), sizeRange(7))
+  }
+
+  test("Translation of Nat") {
+    val natSet = tla.natSet().as(SetT1(IntT1()))
+
+    withSetGen(Gen.const(natSet))
+  }
+
+  test("Translation of Int") {
+    val intSet = tla.intSet().as(SetT1(IntT1()))
+
+    withSetGen(Gen.const(intSet))
+  }
+
+  test("Translation of a..b") {
+    withSetGen(rangeGen)
+  }
+}