IntRangeCache rework (#2673)

* IntRangeCache

* uncomment

* Update tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/aux/caches/IntRangeCache.scala

Co-authored-by: Thomas Pani <[email protected]>

* Update tla-bmcmt/src/test/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/aux/FakeCtx.scala

Co-authored-by: Thomas Pani <[email protected]>

* PR comments

* PR comments 2

---------

Co-authored-by: Thomas Pani <[email protected]>

tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/aux/caches/IntRangeCache.scala

@@ -0,0 +1,58 @@
+package at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux.caches
+
+import at.forsyte.apalache.tla.bmcmt.smt.SolverContext
+import at.forsyte.apalache.tla.bmcmt.types.CellT
+import at.forsyte.apalache.tla.bmcmt.{ArenaCell, FixedElemPtr, PureArena}
+import at.forsyte.apalache.tla.lir.{IntT1, SetT1}
+import at.forsyte.apalache.tla.types.tla
+
+/**
+ * Cache ranges a..b and, as a special case, tuple domains.
+ *
+ * @author
+ *   Jure Kukovec
+ */
+class IntRangeCache(intValueCache: IntValueCache) extends Cache[PureArena, (Int, Int), (ArenaCell, Seq[ArenaCell])] {
+
+  /**
+   * Given a range as a boundary-pair `(a,b)`, returns a tuple `(rArena, (rCell, rCells))`, where:
+   *   - `rCell` is the cell representing the set `a..b`
+   *   - `rCells` is s sequence of cells `c_a, c_{a+1},..., c_b`, representing the range contents, i.e. `a, a+1, ...,
+   *     b`.
+   *   - `rArena` is an extension of `arena`, containing all of the above cells, and a relation
+   *     {{{rCell --(has)--> c_a, c_{a+1},..., c_b}}}
+   *
+   * Note that this method internally calls `intValueCache.getOrCreate`, which caches all integers in the range.
+   *
+   * If `a > b`, we treat the input as if `range` were `(1,0)`, regardless of the concrete values of `a` and `b`, which
+   * caches all empty ranges to the same value.
+   */
+  override protected def create(arena: PureArena, range: (Int, Int)): (PureArena, (ArenaCell, Seq[ArenaCell])) = {
+
+    // We normalize a..b, such that, if the initial a is greater than b (which would give an empty range),
+    // we always use 1..0. This lets us cache _all_ empty ranges into one cell.
+    val (iFrom, iTo) = if (range._1 > range._2) (1, 0) else range
+
+    val (arenaWithCachedInts, allCells) =
+      iFrom.to(iTo).foldLeft((arena, Seq.empty[ArenaCell])) { case ((partialArena, partialCells), key) =>
+        val (newArena, cell) = intValueCache.getOrCreate(partialArena, key)
+        (newArena, partialCells :+ cell)
+      }
+    // create the domain cell
+    val arenaWithDomCell = arenaWithCachedInts.appendCell(CellT.fromType1(SetT1(IntT1)))
+    val setCell = arenaWithDomCell.topCell
+    val arenaWithHas = arenaWithDomCell.appendHas(setCell, allCells.map(FixedElemPtr): _*)
+
+    (arenaWithHas, (setCell, allCells))
+  }
+
+  /** Return a function to add implementation-specific constraints for a single entry */
+  override def addConstraintsForElem(ctx: SolverContext): (((LevelT, LevelT), (ArenaCell, Seq[ArenaCell]))) => Unit = {
+    case (_, (setCell, elemCells)) =>
+      elemCells.foreach { elemCell =>
+        // We _know_ the pointer is fixed TRUE by construction, so instead of asserting X == true, we just assert X, where
+        // X = elemCell \in setCell
+        ctx.assertGroundExpr(tla.in(elemCell.toBuilder, setCell.toBuilder))
+      }
+  }
+}

tla-bmcmt/src/test/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/aux/IntRangeCacheTest.scala

@@ -0,0 +1,101 @@
+package at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux
+
+import at.forsyte.apalache.tla.bmcmt.PureArena
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux.caches.{IntRangeCache, IntValueCache}
+import at.forsyte.apalache.tla.types.tla
+import org.junit.runner.RunWith
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.funsuite.AnyFunSuite
+import org.scalatestplus.junit.JUnitRunner
+
+@RunWith(classOf[JUnitRunner])
+class IntRangeCacheTest extends AnyFunSuite with BeforeAndAfterEach {
+
+  var intValueCache: IntValueCache = new IntValueCache
+  var cache: IntRangeCache = new IntRangeCache(intValueCache)
+
+  override def beforeEach(): Unit = {
+    intValueCache = new IntValueCache
+    cache = new IntRangeCache(intValueCache)
+  }
+
+  test("Cache returns stored values after the first call to getOrCreate") {
+    val rangeTup = (1, 5)
+    val range = rangeTup._1.to(rangeTup._2)
+
+    val arena = PureArena.empty
+
+    // No cached values
+    assert(cache.get(rangeTup).isEmpty)
+    assert(range.forall { i => intValueCache.get(i).isEmpty })
+
+    val (newArena, retPair) = cache.getOrCreate(arena, rangeTup)
+
+    // range is now cached, arena has changed
+    assert(cache.get(rangeTup).nonEmpty && newArena != arena && range.forall { i => intValueCache.get(i).nonEmpty })
+
+    val (newArena2, retPair2) = cache.getOrCreate(newArena, rangeTup)
+
+    // 2nd call returns the _same_ arena and the previously computed cell
+    assert(newArena == newArena2 && retPair == retPair2)
+  }
+
+  test("Caching a..b ranges, where a > b, returns the same cell with no elements, for all such a,b") {
+    val rangeTup1 = (5, 1)
+    val arena = PureArena.empty
+    val (newArena, (rCell, cells)) = cache.getOrCreate(arena, rangeTup1)
+    assert(cells.isEmpty)
+
+    val rangeTup2 = (1203134, -221)
+    val (newArena2, (rCell2, cells2)) = cache.getOrCreate(arena, rangeTup2)
+    assert(newArena == newArena2 && rCell == rCell2 && cells == cells2)
+  }
+
+  test("Constraints are only added when addConstraintsForElem is explicitly called, and only once per value") {
+    val mockCtx: MockZ3SolverContext = new MockZ3SolverContext
+
+    val r1: (Int, Int) = (0, 10)
+    val r2: (Int, Int) = (3, 7)
+    val r3: (Int, Int) = (100, 0)
+
+    val a0 = PureArena.empty
+    val (a1, (cr1, csr1)) = cache.getOrCreate(a0, r1)
+    // Some extra calls, which shouldn't affect constraint generation
+    cache.getOrCreate(a0, r1)
+    cache.getOrCreate(a0, r1)
+    val (a2, (cr2, csr2)) = cache.getOrCreate(a1, r2)
+    // Some extra calls, which shouldn't affect constraint generation
+    cache.getOrCreate(a1, r2)
+    cache.getOrCreate(a1, r2)
+    val (_, (cr3, csr3)) = cache.getOrCreate(a2, r3)
+    // Some extra calls, which shouldn't affect constraint generation
+    cache.getOrCreate(a2, r3)
+    cache.getOrCreate(a2, r3)
+
+    assert(mockCtx.constraints.isEmpty)
+
+    cache.addAllConstraints(mockCtx)
+
+    // Dependent caches don't discharge constraints unless called, so we should have 0 constraints
+    // from IntValueCache, only the ones from the range cache
+    def nConstraints(r: (Int, Int)): Int = math.max(r._2 - r._1 + 1, 0)
+
+    assert(mockCtx.constraints.size == nConstraints(r1) + nConstraints(r2) + nConstraints(r3))
+    assert(
+        csr1.forall { c =>
+          mockCtx.constraints.contains(tla.in(c.toBuilder, cr1.toBuilder).build)
+        }
+    )
+    assert(
+        csr2.forall { c =>
+          mockCtx.constraints.contains(tla.in(c.toBuilder, cr2.toBuilder).build)
+        }
+    )
+    assert(
+        csr3.forall { c =>
+          mockCtx.constraints.contains(tla.in(c.toBuilder, cr3.toBuilder).build)
+        }
+    )
+  }
+
+}

tla-bmcmt/src/test/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/aux/IntValueCacheTest.scala

@@ -1,13 +1,14 @@
 package at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux
 
 import at.forsyte.apalache.tla.bmcmt.PureArena
-import at.forsyte.apalache.tla.bmcmt.smt.{SolverConfig, Z3SolverContext}
 import at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux.caches.IntValueCache
-import at.forsyte.apalache.tla.lir.TlaEx
 import at.forsyte.apalache.tla.types.tla
+import org.junit.runner.RunWith
 import org.scalatest.BeforeAndAfterEach
 import org.scalatest.funsuite.AnyFunSuite
+import org.scalatestplus.junit.JUnitRunner
 
+@RunWith(classOf[JUnitRunner])
 class IntValueCacheTest extends AnyFunSuite with BeforeAndAfterEach {
 
   var cache: IntValueCache = new IntValueCache
@@ -36,13 +37,7 @@ class IntValueCacheTest extends AnyFunSuite with BeforeAndAfterEach {
   }
 
   test("Constraints are only added when addConstraintsForElem is explicitly called, and only once per value") {
-    // we make a fake context, which just intercepts any expression added via assertGroundExpr and stores it in a Seq.
-    class FakeCtx extends Z3SolverContext(SolverConfig.default) {
-      var constraints: Seq[TlaEx] = Seq.empty
-      override def assertGroundExpr(ex: TlaEx): Unit = constraints = constraints :+ ex
-    }
-
-    val fakeCtx: FakeCtx = new FakeCtx
+    val mockCtx: MockZ3SolverContext = new MockZ3SolverContext
 
     val i1: BigInt = 42
     val i2: BigInt = 0
@@ -57,16 +52,16 @@ class IntValueCacheTest extends AnyFunSuite with BeforeAndAfterEach {
     cache.getOrCreate(a1, i2)
     cache.getOrCreate(a1, i2)
 
-    assert(fakeCtx.constraints.isEmpty)
+    assert(mockCtx.constraints.isEmpty)
 
-    cache.addAllConstraints(fakeCtx)
+    cache.addAllConstraints(mockCtx)
 
-    assert(fakeCtx.constraints.size == 2)
+    assert(mockCtx.constraints.size == 2)
     assert(
-        fakeCtx.constraints.contains(tla.eql(ci1.toBuilder, tla.int(i1)).build)
+        mockCtx.constraints.contains(tla.eql(ci1.toBuilder, tla.int(i1)).build)
     )
     assert(
-        fakeCtx.constraints.contains(tla.eql(ci2.toBuilder, tla.int(i2)).build)
+        mockCtx.constraints.contains(tla.eql(ci2.toBuilder, tla.int(i2)).build)
     )
 
   }

tla-bmcmt/src/test/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/aux/MockZ3SolverContext.scala

@@ -0,0 +1,14 @@
+package at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux
+
+import at.forsyte.apalache.tla.bmcmt.smt.{SolverConfig, Z3SolverContext}
+import at.forsyte.apalache.tla.lir.TlaEx
+
+/**
+ * A mock context, which just intercepts any expression added via assertGroundExpr and stores it in a Seq.
+ *
+ * We can use it in tests, to observe constraint generation without needing to instrument SMT.
+ */
+class MockZ3SolverContext extends Z3SolverContext(SolverConfig.default) {
+  var constraints: Seq[TlaEx] = Seq.empty
+  override def assertGroundExpr(ex: TlaEx): Unit = constraints = constraints :+ ex
+}