Add shrinker for PBT-generated TLaType1's

apalache-mc · May 17, 2022 · 11db09e · 11db09e
1 parent 1c654a5
commit 11db09e
Showing 1 changed file with 45 additions and 1 deletion.
diff --git a/tla-bmcmt/src/test/scala/at/forsyte/apalache/tla/bmcmt/CrossTestEncodings.scala b/tla-bmcmt/src/test/scala/at/forsyte/apalache/tla/bmcmt/CrossTestEncodings.scala
@@ -12,13 +12,17 @@ import at.forsyte.apalache.tla.lir.transformations.impl.IdleTracker
 import at.forsyte.apalache.tla.lir.transformations.standard.IncrementalRenaming
 import ch.qos.logback.classic.{Level, Logger}
 import org.scalacheck.Arbitrary.arbitrary
-import org.scalacheck.Gen
 import org.scalacheck.Gen.{const, listOf, lzy, nonEmptyListOf, oneOf, sized}
 import org.scalacheck.Prop.{forAll, AnyOperators}
+import org.scalacheck.Shrink.shrink
+import org.scalacheck.{Gen, Shrink}
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatestplus.scalacheck.Checkers
 import org.slf4j.LoggerFactory
 
+import scala.annotation.nowarn
+import scala.collection.immutable.SortedMap
+
 /**
  * Test encodings against each other. Override [[AnyFunSuite.withFixture]] to set up SMT encodings used for oracle and
  * verifier.
@@ -98,6 +102,7 @@ trait CrossTestEncodings extends AnyFunSuite with Checkers {
           oneOf(genPrimitive, genSet, genSeq, genFun, /*genOper,*/ genTup, genRec /*, genRowRec, genVariant, genRow*/ )
       }
     }
+
     // Override to avoid reals, constants, and typevar-typed expressions.
     override def genPrimitive: Gen[TlaType1] =
       oneOf(const(BoolT1()), const(IntT1()), const(StrT1()) /*, const(RealT1()), genConst, genVar*/ )
@@ -271,6 +276,7 @@ trait CrossTestEncodings extends AnyFunSuite with Checkers {
 
   /**
    * Generate TLA+ expressions of a given [[TlaType1]].
+   *
    * @param witnessType
    *   [[TlaType1]] of the expression to generate.
    * @return
@@ -311,6 +317,7 @@ trait CrossTestEncodings extends AnyFunSuite with Checkers {
 
   /**
    * Generate TLA+ expressions of type `Set(witnessType)`.
+   *
    * @param witnessType
    *   Element [[TlaType1]] of the set expression to generate.
    * @return
@@ -335,6 +342,43 @@ trait CrossTestEncodings extends AnyFunSuite with Checkers {
     }
   }
 
+  def powerSet[A](xs: Seq[A]): Seq[Seq[A]] =
+    xs.foldLeft(Seq(Seq[A]())) { (sets, set) => sets ++ sets.map(_ :+ set) }
+
+  def powerSet[A, B](xs: SortedMap[A, B]): Seq[SortedMap[A, B]] = powerSet(xs)
+
+  @nowarn("cat=deprecation&msg=Stream") // ScalaCheck's Shrink still uses Stream, which is deprecated since Scala 2.13.0
+  implicit def shrinkType(implicit tp: Shrink[TlaType1]): Shrink[TlaType1] = Shrink {
+    case IntT1() | RealT1() | BoolT1() | StrT1() => Stream.empty
+    case FunT1(arg, res)                         => Stream(arg, res)
+    case SetT1(elem)                             => Stream(elem)
+    case SeqT1(elem)                             => Stream(elem)
+    case TupT1(elems @ _*) =>
+      Stream.concat(
+          // shrink to one of the element types
+          Stream(elems: _*),
+          // recursively shrink one tuple element type
+          elems.zipWithIndex.foldLeft(Stream.empty[TlaType1]) { case (acc, (tp, idx)) =>
+            acc ++ shrink(tp).map(shrunkType => TupT1(elems.updated(idx, shrunkType): _*))
+          },
+          // tuples over a subset of the current element types
+          Stream(powerSet(elems).map(elems => TupT1(elems: _*)): _*),
+      )
+    case RecT1(fieldTypes) =>
+      Stream.concat(
+          // shrink to one of the field types
+          Stream(fieldTypes.values.toSeq: _*),
+          // recursively shrink one record field type
+          fieldTypes.foldLeft(Stream.empty[TlaType1]) { case (acc, (fieldName, tp)) =>
+            acc ++ shrink(tp).map(shrunkType => RecT1(fieldTypes.updated(fieldName, shrunkType)))
+          },
+          // tuples over a subset of the current field types
+          Stream(powerSet(fieldTypes).map(elems => RecT1(elems)): _*),
+      )
+    case ConstT1(_) | OperT1(_, _) | RealT1() | RecRowT1(_) | RowT1(_, _) | SparseTupT1(_) | VarT1(_) | VariantT1(_) =>
+      throw new IllegalArgumentException(s"Shrinking unsupported type ${tp}. Should be disabled in type generator.")
+  }
+
   // Ignore this test until the bug described in https://github.com/informalsystems/apalache/issues/1660 is fixed.
   // TODO(#1668): Regularly run the encodings comparison.
   ignore("encodings are consistent") {