First attempt at refactoring Oracles

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

@@ -18,6 +18,7 @@ import at.forsyte.apalache.tla.types.tla
  *   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)`.
@@ -28,13 +29,13 @@ class IntOracle(val intCell: ArenaCell, nvalues: Int) extends Oracle {
    * 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 oracleValueIsEqualToIndexedValue(scope: RewriterScope, index: Int): TBuilderInstruction =
+  override def chosenValueIsEqualToIndexedValue(scope: RewriterScope, index: BigInt): TBuilderInstruction =
     tla.eql(intCell.toBuilder, tla.int(index))
 
-  override def getIndexOfOracleValueFromModel(solverContext: SolverContext): Int =
+  override def getIndexOfChosenValueFromModel(solverContext: SolverContext): BigInt =
     solverContext.evalGroundExpr(intCell.toBuilder) match {
-      case ValEx(TlaInt(i)) => i.toInt
-      case _                => throw new IllegalStateException(s"Invalid call to evalPosition, not an integer.")
+      case ValEx(TlaInt(i)) => i
+      case _ => throw new IllegalStateException(s"Invalid call to \"getIndexOfOracleValueFromModel\", not an integer.")
     }
 
 }

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

@@ -29,7 +29,7 @@ trait Oracle {
    * 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 oracleValueIsEqualToIndexedValue(scope: RewriterScope, index: Int): TBuilderInstruction
+  def chosenValueIsEqualToIndexedValue(scope: RewriterScope, index: BigInt): TBuilderInstruction
 
   /**
    * Produce a ground expression that contains assertions for the possible oracle values.
@@ -47,9 +47,8 @@ trait Oracle {
       scope: RewriterScope,
       assertions: Seq[TBuilderInstruction],
       elseAssertionsOpt: Option[Seq[TBuilderInstruction]] = None): TBuilderInstruction = {
-    if (assertions.size != this.size || elseAssertionsOpt.exists { _.size != this.size })
-      throw new IllegalStateException(s"Invalid call to Oracle, assertion sequences must have length $size.")
-
+    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
@@ -66,7 +65,7 @@ trait Oracle {
         // 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(oracleValueIsEqualToIndexedValue(scope, i), thenFormula, elseFormula)
+              tla.ite(chosenValueIsEqualToIndexedValue(scope, i), thenFormula, elseFormula)
             }: _*
         )
     }
@@ -76,5 +75,5 @@ trait Oracle {
    * Decode the value of the oracle variable into an integer index. This method assumes that the solver context has
    * produced an SMT model.
    */
-  def getIndexOfOracleValueFromModel(solverContext: SolverContext): Int
+  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))
+  }
+}