First attempt at refactoring Oracles

tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/aux/oracles/IntOracle.scala

@@ -0,0 +1,49 @@
+package at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux.oracles
+import at.forsyte.apalache.tla.bmcmt.ArenaCell
+import at.forsyte.apalache.tla.bmcmt.smt.SolverContext
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.RewriterScope
+import at.forsyte.apalache.tla.bmcmt.types.CellT
+import at.forsyte.apalache.tla.lir.{IntT1, ValEx}
+import at.forsyte.apalache.tla.lir.values.TlaInt
+import at.forsyte.apalache.tla.typecomp.TBuilderInstruction
+import at.forsyte.apalache.tla.types.tla
+
+/**
+ * An oracle that uses an integer variable. Although using integers as an oracle is the most straightforward decision,
+ * whenever a specialized oracle is available, it should be used instead, for performance reasons.
+ *
+ * @author
+ *   Jure Kukovec
+ */
+class IntOracle(val intCell: ArenaCell, nvalues: Int) extends Oracle {
+  require(nvalues >= 0, "IntOracle must have a non-negative number of candidate values.")
+
+  /**
+   * The number of values that this oracle is defined over: `0..(size - 1)`.
+   */
+  override def size: Int = nvalues
+
+  /**
+   * Produce an expression that states that the chosen value equals to the value `v_{index}`. The actual implementation
+   * may be different from an integer comparison.
+   */
+  override def chosenValueIsEqualToIndexedValue(scope: RewriterScope, index: BigInt): TBuilderInstruction =
+    tla.eql(intCell.toBuilder, tla.int(index))
+
+  override def getIndexOfChosenValueFromModel(solverContext: SolverContext): BigInt =
+    solverContext.evalGroundExpr(intCell.toBuilder) match {
+      case ValEx(TlaInt(i)) => i
+      case _ => throw new IllegalStateException(s"Invalid call to \"getIndexOfOracleValueFromModel\", not an integer.")
+    }
+
+}
+
+object IntOracle {
+  def create(scope: RewriterScope, nvalues: Int): (RewriterScope, IntOracle) = {
+    val newArena = scope.arena.appendCell(CellT.fromType1(IntT1))
+    val oracleCell = newArena.topCell
+    val oracle = new IntOracle(oracleCell, nvalues)
+    // the oracle value must be equal to one of the value cells
+    (scope.copy(arena = newArena), oracle)
+  }
+}

tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/aux/oracles/Oracle.scala

@@ -0,0 +1,79 @@
+package at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux.oracles
+
+import at.forsyte.apalache.tla.bmcmt.PureArena
+import at.forsyte.apalache.tla.bmcmt.smt.SolverContext
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.RewriterScope
+import at.forsyte.apalache.tla.typecomp.TBuilderInstruction
+import at.forsyte.apalache.tla.types.tla
+
+/**
+ * An oracle provides a means of choosing a single value, out of a finite collection of candidate values. It can
+ * generate SMT constraints, which capture the conditions under which the chosen value may/must be equal to the
+ * candidate values.
+ *
+ * For the purposes of describing the oracle's behavior, we can assume that there are `N` distinct possible values,
+ * which are indexed with `0,1,2,...,N-1` (though they need not be represented as such in the implementation). We refer
+ * to the `i`-th candidate value as `vi`.
+ *
+ * @author
+ *   Jure Kukovec
+ */
+trait Oracle {
+
+  /**
+   * The number of values that this oracle is defined over: `0..(size - 1)`.
+   */
+  def size: Int
+
+  /**
+   * Produce an expression that states that the chosen value equals to the value `v_{index}`. The actual implementation
+   * may be different from an integer comparison.
+   */
+  def chosenValueIsEqualToIndexedValue(scope: RewriterScope, index: BigInt): TBuilderInstruction
+
+  /**
+   * Produce a ground expression that contains assertions for the possible oracle values.
+   *
+   * The `assertions` sequence must contain `N` elements, where `N` equals the number of oracle value candidates. The
+   * `i`-th element of `assertions` describes a formula, which holds if the oracle value is equal to the `i`-th
+   * candidate value `vi`.
+   *
+   * Optionally, a sequence option `elseAssertionsOpt`, containing a sequence of the same length as `assertions`, may be
+   * provided. If it is, the `i`-th element of the sequence describes a formula, which holds if the oracle value is
+   * _not_ equal to the `i`-th candidate value `vi`. If this value is not provided, a default sequence, containing `N`
+   * copies of `TRUE` is taken in place of the aforementioned formulas.
+   */
+  def caseAssertions(
+      scope: RewriterScope,
+      assertions: Seq[TBuilderInstruction],
+      elseAssertionsOpt: Option[Seq[TBuilderInstruction]] = None): TBuilderInstruction = {
+    require(assertions.size == this.size && elseAssertionsOpt.forall { _.size == this.size },
+        s"Invalid call to Oracle, assertion sequences must have length $size.")
+    size match {
+      // If the oracle has no candidate values, then caseAssertions should return a no-op for SMT, i.e. TRUE.
+      case 0 => PureArena.cellTrue(scope.arena).toBuilder
+
+      // If the oracle has a single candidate value, then the chosen value will always equal to the (sole) candidate value.
+      // In other words, for such an oracle, whenEqualTo(..., 0) will always hold.
+      // Therefore, we don't need an ITE (and don't need to consider elseAssertions), we just take the first assertion.
+      case 1 => assertions.head
+
+      case _ =>
+        val elseAssertions = elseAssertionsOpt.getOrElse(Seq.fill(size)(tla.bool(true)))
+        // We zip a sequence of tuples, with the 1st and 2nd elements of each tuple being the "then" and "else" cases of an ite.
+        // If elseAssertionsOpt is not empty, each tuple has its 1st element from assertions and its 2nd form elseAssertionsOpt.
+        // If elseAssertionsOpt is empty, each tuple has its 1st element from assertions and its 2nd defaults to TRUE.
+        tla.and(
+            assertions.zip(elseAssertions).zipWithIndex.map { case ((thenFormula, elseFormula), i) =>
+              tla.ite(chosenValueIsEqualToIndexedValue(scope, i), thenFormula, elseFormula)
+            }: _*
+        )
+    }
+  }
+
+  /**
+   * Decode the value of the oracle variable into an integer index. This method assumes that the solver context has
+   * produced an SMT model.
+   */
+  def getIndexOfChosenValueFromModel(solverContext: SolverContext): BigInt
+}

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

@@ -0,0 +1,146 @@
+package at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux.oracles
+
+import at.forsyte.apalache.tla.bmcmt.PureArena
+import at.forsyte.apalache.tla.bmcmt.arena.PureArenaAdapter
+import at.forsyte.apalache.tla.bmcmt.smt.{SolverConfig, Z3SolverContext}
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.RewriterScope
+import at.forsyte.apalache.tla.lir.{BoolT1, NameEx, OperEx, TlaEx, ValEx}
+import at.forsyte.apalache.tla.lir.oper.TlaOper
+import at.forsyte.apalache.tla.lir.values.TlaInt
+import at.forsyte.apalache.tla.typecomp.TBuilderInstruction
+import at.forsyte.apalache.tla.types.tla
+import org.junit.runner.RunWith
+import org.scalacheck.{Gen, Prop}
+import org.scalacheck.Prop.forAll
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.funsuite.AnyFunSuite
+import org.scalatestplus.junit.JUnitRunner
+import org.scalatestplus.scalacheck.Checkers
+
+@RunWith(classOf[JUnitRunner])
+class TestIntOracle extends AnyFunSuite with BeforeAndAfterEach with Checkers {
+
+  var initScope: RewriterScope = RewriterScope.initial()
+
+  override def beforeEach(): Unit = {
+    initScope = RewriterScope.initial()
+  }
+
+  val intGen: Gen[Int] = Gen.choose(-10, 10)
+  val nonNegIntGen: Gen[Int] = Gen.choose(0, 10)
+
+  val maxSizeAndIndexGen: Gen[(Int, Int)] = for {
+    max <- nonNegIntGen
+    idx <- Gen.choose(0, max)
+  } yield (max, idx)
+
+  test("Oracle cannot be constructed with negative size") {
+    val prop =
+      forAll(intGen) {
+        case i if i < 0 =>
+          Prop.throws(classOf[IllegalArgumentException]) {
+            IntOracle.create(initScope, i)
+          }
+        case i => IntOracle.create(initScope, i)._2.size == i
+      }
+
+    check(prop, minSuccessful(1000), sizeRange(4))
+  }
+
+  test("oracleValueIsEqualToIndexedValue returns an integer comparison") {
+    val prop =
+      forAll(maxSizeAndIndexGen) { case (size, index) =>
+        val (scope, oracle) = IntOracle.create(initScope, size)
+        val cmp: TlaEx = oracle.chosenValueIsEqualToIndexedValue(scope, index)
+        cmp match {
+          case OperEx(TlaOper.eq, NameEx(name), ValEx(TlaInt(i))) => name == oracle.intCell.toString && i == index
+          case _                                                  => false
+        }
+      }
+
+    check(prop, minSuccessful(1000), sizeRange(4))
+
+  }
+
+  val (assertionsA, assertionsB): (Seq[TBuilderInstruction], Seq[TBuilderInstruction]) = 0
+    .to(10)
+    .map { i =>
+      (tla.name(s"A$i", BoolT1), tla.name(s"B$i", BoolT1))
+    }
+    .unzip
+
+  test("caseAssertions requires assertion sequences of equal length") {
+    val assertionsGen: Gen[(Seq[TBuilderInstruction], Option[Seq[TBuilderInstruction]])] = for {
+      i <- Gen.choose(0, assertionsA.size)
+      j <- Gen.choose(0, assertionsB.size)
+      opt <- Gen.option(Gen.const(assertionsB.take(j)))
+    } yield (assertionsA.take(i), opt)
+
+    val prop =
+      forAll(Gen.zip(nonNegIntGen, assertionsGen)) { case (size, (assertionsIfTrue, assertionsIfFalseOpt)) =>
+        val (scope, oracle) = IntOracle.create(initScope, size)
+        if (assertionsIfTrue.size != oracle.size || assertionsIfFalseOpt.exists { _.size != oracle.size })
+          Prop.throws(classOf[IllegalArgumentException]) {
+            oracle.caseAssertions(scope, assertionsIfTrue, assertionsIfFalseOpt)
+          }
+        else true
+      }
+
+    check(prop, minSuccessful(1000), sizeRange(4))
+
+  }
+
+  test("caseAssertions constructs a collection of ITEs, or shorthands") {
+    val gen: Gen[(Int, Seq[TBuilderInstruction], Option[Seq[TBuilderInstruction]])] = for {
+      size <- nonNegIntGen
+      opt <- Gen.option(Gen.const(assertionsB.take(size)))
+    } yield (size, assertionsA.take(size), opt)
+
+    val prop =
+      forAll(gen) { case (size, assertionsIfTrue, assertionsIfFalseOpt) =>
+        val (scope, oracle) = IntOracle.create(initScope, size)
+        val caseEx: TlaEx = oracle.caseAssertions(scope, assertionsIfTrue, assertionsIfFalseOpt)
+        size match {
+          case 0 =>
+            caseEx == PureArena.cellTrue(scope.arena).toBuilder.build
+          case 1 =>
+            caseEx == assertionsA.head.build
+          case _ =>
+            assertionsIfFalseOpt match {
+              case None =>
+                val ites = assertionsIfTrue.zipWithIndex.map { case (a, i) =>
+                  tla.ite(tla.eql(oracle.intCell.toBuilder, tla.int(i)), a, tla.bool(true))
+                }
+                caseEx == tla.and(ites: _*).build
+              case Some(assertionsIfFalse) =>
+                val ites = assertionsIfTrue.zip(assertionsIfFalse).zipWithIndex.map { case ((at, af), i) =>
+                  tla.ite(tla.eql(oracle.intCell.toBuilder, tla.int(i)), at, af)
+                }
+                caseEx == tla.and(ites: _*).build
+            }
+        }
+      }
+
+    check(prop, minSuccessful(1000), sizeRange(4))
+
+  }
+
+  // We cannot test getIndexOfChosenValueFromModel without running the solver
+  test("getIndexOfChosenValueFromModel recovers the index correctly") {
+    val prop =
+      forAll(Gen.zip(maxSizeAndIndexGen)) { case (size, index) =>
+        val ctx = new Z3SolverContext(SolverConfig.default)
+        val paa = PureArenaAdapter.create(ctx) // We use PAA, since it performs the basic context initialization
+        val (scope, oracle) = IntOracle.create(initScope.copy(arena = paa.arena), size)
+        ctx.declareCell(oracle.intCell)
+
+        val eql = oracle.chosenValueIsEqualToIndexedValue(scope, index)
+        ctx.assertGroundExpr(eql)
+        ctx.sat()
+        oracle.getIndexOfChosenValueFromModel(ctx) == index
+      }
+
+    // 1000 is too many, since each run invokes the solver
+    check(prop, minSuccessful(200), sizeRange(4))
+  }
+}