IntConstStratifiedRule (#2706)

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

@@ -1,13 +1,16 @@
 package at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux
 
+import at.forsyte.apalache.tla.bmcmt.smt.SolverContext
 import at.forsyte.apalache.tla.bmcmt.stratifiedRules._
 import at.forsyte.apalache.tla.bmcmt.stratifiedRules.apalache.AssignmentStratifiedRule
-import at.forsyte.apalache.tla.bmcmt.stratifiedRules.base.{BuiltinConstStratifiedRule, SubstStratifiedRule}
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux.caches.IntValueCache
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.base.{
+  BuiltinConstStratifiedRule, IntConstStratifiedRule, SubstStratifiedRule,
+}
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.bool.{AndStratifiedRule, OrStratifiedRule}
 import at.forsyte.apalache.tla.bmcmt.stratifiedRules.set.{
   SetCtorStratifiedRule, SetCupStratifiedRule, SetFilterStratifiedRule,
 }
-import at.forsyte.apalache.tla.bmcmt.smt.SolverContext
-import at.forsyte.apalache.tla.bmcmt.stratifiedRules.bool.{AndStratifiedRule, OrStratifiedRule}
 import at.forsyte.apalache.tla.lir._
 import at.forsyte.apalache.tla.lir.oper.TlaOper
 import at.forsyte.apalache.tla.lir.values._
@@ -55,6 +58,8 @@ abstract class RewriterImpl(@unused private val ctx: SolverContext) extends Rewr
     case _                                    => "O@"
   }
 
+  val intValueCache: IntValueCache = new IntValueCache
+
   // A nice way to guess the candidate rules by looking at the expression key.
   // We use simple expressions to generate the keys.
   // For each key, there is a short list of rules that may be applicable.
@@ -79,8 +84,8 @@ abstract class RewriterImpl(@unused private val ctx: SolverContext) extends Rewr
           -> new BuiltinConstStratifiedRule,
         key(tla.natSet())
           -> new BuiltinConstStratifiedRule,
-//        key(ValEx(TlaInt(1)))
-//        -> List(new IntConstRule(this)),
+        key(tla.int(1))
+          -> new IntConstStratifiedRule(this, intValueCache),
 //      key(ValEx(TlaStr("red")))
 //        -> List(new StrConstRule(this)),
 //      // logic

tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/base/IntConstStratifiedRule.scala

@@ -0,0 +1,36 @@
+package at.forsyte.apalache.tla.bmcmt.stratifiedRules.base
+
+import at.forsyte.apalache.tla.bmcmt._
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.aux.caches.IntValueCache
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.{Rewriter, RewriterScope, StratifiedRule}
+import at.forsyte.apalache.tla.lir.values.TlaInt
+import at.forsyte.apalache.tla.lir.{TlaEx, ValEx}
+import at.forsyte.apalache.tla.types.tla
+
+/**
+ * Rewrites integer constants.
+ *
+ * @author
+ *   Jure Kukovec
+ */
+class IntConstStratifiedRule(rewriter: Rewriter, intValueCache: IntValueCache) extends StratifiedRule[BigInt] {
+
+  def isApplicable(ex: TlaEx, scope: RewriterScope): Boolean = ex match {
+    case ValEx(TlaInt(_)) => true
+    case _                => false
+  }
+
+  def buildArena(ex: TlaEx)(startingScope: RewriterScope): (RewriterScope, ArenaCell, BigInt) = ex match {
+    case ValEx(TlaInt(n)) =>
+      val (newArena, cell) = intValueCache.getOrCreate(startingScope.arena, n)
+      (startingScope.copy(arena = newArena), cell, n)
+
+    case _ =>
+      throw new RewriterException(getClass.getSimpleName + " is not applicable", ex)
+  }
+
+  def addConstraints(scope: RewriterScope, cell: ArenaCell, auxData: BigInt): Unit = {
+    rewriter.assert(tla.eql(cell.toBuilder, tla.int(auxData)))
+  }
+
+}

tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/stratifiedRules/base/SubstStratifiedRule.scala

@@ -1,7 +1,7 @@
 package at.forsyte.apalache.tla.bmcmt.stratifiedRules.base
 
 import at.forsyte.apalache.tla.bmcmt._
-import at.forsyte.apalache.tla.bmcmt.stratifiedRules.{RewriterScope, StratifiedRule}
+import at.forsyte.apalache.tla.bmcmt.stratifiedRules.{RewriterScope, StratifiedRuleInterface}
 import at.forsyte.apalache.tla.lir.{NameEx, TlaEx}
 import at.forsyte.apalache.tla.pp.TlaInputError
 import com.typesafe.scalalogging.LazyLogging
@@ -13,7 +13,7 @@ import com.typesafe.scalalogging.LazyLogging
  * @author
  *   Jure Kukovec
  */
-class SubstStratifiedRule extends StratifiedRule[Unit] with LazyLogging {
+class SubstStratifiedRule extends StratifiedRuleInterface with LazyLogging {
   override def isApplicable(ex: TlaEx, scop: RewriterScope): Boolean = ex match {
     case NameEx(x) =>
       // make sure that x is not an SMT constant, but a variable name
@@ -22,11 +22,11 @@ class SubstStratifiedRule extends StratifiedRule[Unit] with LazyLogging {
     case _ => false
   }
 
-  def buildArena(ex: TlaEx)(startingScope: RewriterScope): (RewriterScope, ArenaCell, Unit) = ex match {
+  def apply(ex: TlaEx)(startingScope: RewriterScope): (RewriterScope, ArenaCell) = ex match {
     case NameEx(x) =>
       val binding = startingScope.binding
       if (binding.contains(x)) {
-        (startingScope, binding(x), ())
+        (startingScope, binding(x))
       } else {
         logger.error("This error may show up when CONSTANTS are not initialized.")
         logger.error("Check the manual: https://apalache.informal.systems/docs/apalache/parameters.html")
@@ -36,6 +36,4 @@ class SubstStratifiedRule extends StratifiedRule[Unit] with LazyLogging {
     case _ =>
       throw new RewriterException("%s is not applicable".format(getClass.getSimpleName), ex)
   }
-
-  def addConstraints(scope: RewriterScope, cell: ArenaCell, auxData: Unit): Unit = ()
 }

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

@@ -28,6 +28,13 @@ class BoolRewriterTest extends AnyFunSuite with BeforeAndAfterEach {
     assert(trueCell.toNameEx.eqTyped(NameEx("$C$1")(Typed(BoolT1))))
   }
 
+  test("Boolean set rewriting rule: BOOLEAN ~~> $C$2") {
+    val boolSetEx = tla.booleanSet()
+
+    val cell = rewriter.rewrite(boolSetEx)(RewriterScope.initial())._2
+    assert(cell.toNameEx.eqTyped(NameEx("$C$2")(Typed(SetT1(BoolT1)))))
+  }
+
   test("Boolean operator rewriting rule: c_x /\\ c_y") {
     // Load up the arena with 2 var cells
     val (Seq(xCell, yCell), arena) = Seq(BoolT1, BoolT1).foldLeft((Seq.empty[ArenaCell], PureArena.initial)) {

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

@@ -0,0 +1,53 @@
+package at.forsyte.apalache.tla.bmcmt.stratifiedRules
+
+import at.forsyte.apalache.tla.lir.{IntT1, NameEx, SetT1, Typed}
+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 IntRewriterTest extends AnyFunSuite with BeforeAndAfterEach {
+
+  var rewriter: TestingRewriter = TestingRewriter(Map.empty)
+
+  override def beforeEach(): Unit = {
+    rewriter = TestingRewriter(Map.empty)
+  }
+
+  test("Integer set rewriting rule: Nat ~~> $C$3 / Int ~~> $C$4") {
+    val natSetEx = tla.natSet()
+    val intSetEx = tla.intSet()
+
+    val cellNat = rewriter.rewrite(natSetEx)(RewriterScope.initial())._2
+    assert(cellNat.toNameEx.eqTyped(NameEx("$C$3")(Typed(SetT1(IntT1)))))
+
+    val cellInt = rewriter.rewrite(intSetEx)(RewriterScope.initial())._2
+    assert(cellInt.toNameEx.eqTyped(NameEx("$C$4")(Typed(SetT1(IntT1)))))
+  }
+
+  test("Integer constant rewriting rule.") {
+    val i1 = tla.int(0)
+    val i2 = tla.int(42)
+
+    val initScope = RewriterScope.initial()
+
+    val (scope1, cell1) = rewriter.rewrite(i1)(initScope)
+
+    assert(
+        rewriter.assertSeq == Seq(tla.eql(cell1.toBuilder, i1).build)
+    )
+
+    val (_, cell2) = rewriter.rewrite(i2)(scope1)
+
+    assert(
+        rewriter.assertSeq == Seq(
+            tla.eql(cell1.toBuilder, i1),
+            tla.eql(cell2.toBuilder, i2),
+        ).map(_.build)
+    )
+
+  }
+
+}