Merge pull request #2661 from informalsystems/th/qnt-parse-bigint

(De-)serialize BigInt from/to JSON number in Quint import

.unreleased/bug-fixes/2661-qnt-bigint-ids.md

@@ -0,0 +1 @@
+- Fixed deserialization of Quint `bigint`s, see #2661

tla-io/src/main/scala/at/forsyte/apalache/io/quint/Quint.scala

@@ -54,7 +54,7 @@ class Quint(quintOutput: QuintOutput) {
   private type NullaryOpReader[A] = Reader[Set[String], A]
 
   // Find the type for an id via the lookup table provided in the quint output
-  private def getTypeFromLookupTable(id: Int): QuintType = {
+  private def getTypeFromLookupTable(id: BigInt): QuintType = {
     table.get(id) match {
       case None => throw new QuintIRParseError(s"No entry found for id ${id} in lookup table")
       case Some(lookupEntry) =>
@@ -95,7 +95,7 @@ class Quint(quintOutput: QuintOutput) {
       } yield (tlaBody, typedParams)
   }
 
-  private def typeTagOfId(id: Int): TypeTag = {
+  private def typeTagOfId(id: BigInt): TypeTag = {
     Typed(typeConv.convert(types(id).typ))
   }
 
@@ -234,7 +234,7 @@ class Quint(quintOutput: QuintOutput) {
         case args  => tla.enumSet(args: _*)
       }
 
-    def listConstruction(id: Int): Converter =
+    def listConstruction(id: BigInt): Converter =
       variadicApp {
         // Empty lists must be handled specially since we cannot infer their type
         // from the given arguments
@@ -301,7 +301,7 @@ class Quint(quintOutput: QuintOutput) {
             tla.foldSeq(testLambda, tla.emptySeq(elemType), seq)
           })
 
-    def exceptWithUpdate(opName: String, id: Int): Converter =
+    def exceptWithUpdate(opName: String, id: BigInt): Converter =
       // f.setBy(x, op) ~~>
       //
       // LET f_cache = f IN

tla-io/src/main/scala/at/forsyte/apalache/io/quint/QuintIR.scala

@@ -32,6 +32,8 @@ package at.forsyte.apalache.io.quint
 // The `key` package allows customizing the JSON key for a class tag or attribute
 // See https://com-lihaoyi.github.io/upickle/#CustomKeys
 import upickle.implicits.key
+import upickle.core.{Abort, Visitor}
+
 import scala.util.Try
 
 // We use a slightly customized deserializer
@@ -42,6 +44,44 @@ private[quint] object QuintDeserializer extends upickle.AttributeTagged {
   // override forms a custom deserializer that is just like the default, except
   // the value of "kind" is used to differentiate
   override def tagName = "kind"
+
+  // Override the built-in BigInt{Reader,Writer} to also support reading/writing
+  // from/to JSON numbers.
+  implicit override val BigIntWriter: QuintDeserializer.Writer[BigInt] = new Writer[BigInt] {
+    override def write0[V](out: Visitor[_, V], v: BigInt): V = out.visitFloat64StringParts(v.toString(), -1, -1, 0)
+  }
+  implicit override val BigIntReader: QuintDeserializer.Reader[BigInt] = {
+    new NumericReader[BigInt] {
+      override def expectedMsg = "expected bigint"
+
+      // Below, `index` is the character offset of the lexical unit to be parsed inside the JSON string
+      override def visitString(s: CharSequence, index: Int) = visitFloat64String(s.toString, index)
+      override def visitInt32(d: Int, index: Int) = BigInt(d)
+      override def visitInt64(d: Long, index: Int) = BigInt(d)
+      override def visitUInt64(d: Long, index: Int) = BigInt(d)
+      override def visitFloat32(d: Float, index: Int) = BigDecimal(d).toBigIntExact match {
+        case Some(bigInt) => bigInt
+        case None         => throw Abort(expectedMsg + " got decimal (as float32)")
+      }
+      override def visitFloat64(d: Double, index: Int) = BigDecimal(d).toBigIntExact match {
+        case Some(bigInt) => bigInt
+        case None         => throw Abort(expectedMsg + " got decimal (as float64)")
+      }
+      override def visitFloat64StringParts(
+          s: CharSequence,
+          decIndex: Int,
+          expIndex: Int,
+          index: Int) = {
+        if (decIndex != -1) {
+          throw Abort(expectedMsg + " got decimal (as float64StringParts)")
+        }
+        if (expIndex != -1) {
+          throw Abort(expectedMsg + " got exp notation (as float64StringParts)")
+        }
+        BigInt(s.toString)
+      }
+    }
+  }
 }
 
 import QuintDeserializer.{macroRW, ReadWriter => RW}
@@ -55,7 +95,7 @@ class QuintUnsupportedError(errMsg: String) extends Exception("Unsupported quint
 
 private[quint] case class QuintLookupTableEntry(
     kind: String,
-    reference: Int)
+    reference: BigInt)
 private[quint] object QuintLookupTableEntry {
   implicit val rw: RW[QuintLookupTableEntry] = macroRW
 }
@@ -65,9 +105,9 @@ private[quint] case class QuintOutput(
     stage: String,
     modules: Seq[QuintModule],
     // Maps source IDs to types, see the `WithId` trait
-    types: Map[Int, QuintTypeScheme],
+    types: Map[BigInt, QuintTypeScheme],
     // Maps name IDs to declaration IDs
-    table: Map[Int, QuintLookupTableEntry])
+    table: Map[BigInt, QuintLookupTableEntry])
 
 object QuintOutput {
   implicit val rw: RW[QuintOutput] = macroRW
@@ -79,7 +119,7 @@ object QuintOutput {
 }
 
 private[quint] case class QuintModule(
-    id: Int,
+    id: BigInt,
     name: String,
     defs: Seq[QuintDef])
 private[quint] object QuintModule {
@@ -99,7 +139,7 @@ private[quint] object QuintTypeScheme {
 
 /** Source IDs, used to associate expressions with their inferred types */
 private[quint] trait WithID {
-  val id: Int
+  val id: BigInt
 }
 
 /** The representation of quint expressions */
@@ -119,29 +159,29 @@ private[quint] object QuintEx {
     )
 
   /** A name of: a variable, constant, parameter, user-defined operator */
-  @key("name") case class QuintName(id: Int, name: String) extends QuintEx {}
+  @key("name") case class QuintName(id: BigInt, name: String) extends QuintEx {}
   object QuintName {
     implicit val rw: RW[QuintName] = macroRW
   }
 
   /** The boolean literal value */
-  @key("bool") case class QuintBool(id: Int, value: Boolean) extends QuintEx {}
+  @key("bool") case class QuintBool(id: BigInt, value: Boolean) extends QuintEx {}
   object QuintBool {
     implicit val rw: RW[QuintBool] = macroRW
   }
 
-  @key("int") case class QuintInt(id: Int, value: BigInt) extends QuintEx {}
+  @key("int") case class QuintInt(id: BigInt, value: BigInt) extends QuintEx {}
   object QuintInt {
     implicit val rw: RW[QuintInt] = macroRW
   }
 
-  @key("str") case class QuintStr(id: Int, value: String) extends QuintEx {}
+  @key("str") case class QuintStr(id: BigInt, value: String) extends QuintEx {}
   object QuintStr {
     implicit val rw: RW[QuintStr] = macroRW
   }
 
   @key("app") case class QuintApp(
-      id: Int,
+      id: BigInt,
       /** The name of the operator being applied */
       opcode: String,
       /** A list of arguments to the operator */
@@ -152,14 +192,14 @@ private[quint] object QuintEx {
   }
 
   case class QuintLambdaParameter(
-      id: Int,
+      id: BigInt,
       name: String)
   object QuintLambdaParameter {
     implicit val rw: RW[QuintLambdaParameter] = macroRW
   }
 
   @key("lambda") case class QuintLambda(
-      id: Int,
+      id: BigInt,
       /** Identifiers for the formal parameters */
       params: Seq[QuintLambdaParameter],
       /** The qualifier for the defined operator */
@@ -173,7 +213,7 @@ private[quint] object QuintEx {
   }
 
   @key("let") case class QuintLet(
-      id: Int,
+      id: BigInt,
       /** The operator being defined for use in the body */
       opdef: QuintDef.QuintOpDef,
       /** The body */
@@ -231,7 +271,7 @@ private[quint] object QuintDef {
    * type of any operator declaration is given in the type map. So we simply omit it.
    */
   @key("def") case class QuintOpDef(
-      id: Int,
+      id: BigInt,
       /** definition name */
       name: String,
       /** qualifiers that identify definition kinds, like `def`, `val`, etc. */
@@ -244,7 +284,7 @@ private[quint] object QuintDef {
   }
 
   @key("const") case class QuintConst(
-      id: Int,
+      id: BigInt,
       /** name of the constant */
       name: String,
       typeAnnotation: QuintType)
@@ -254,7 +294,7 @@ private[quint] object QuintDef {
   }
 
   @key("var") case class QuintVar(
-      id: Int,
+      id: BigInt,
       /** name of the variable */
       name: String,
       typeAnnotation: QuintType)
@@ -264,7 +304,7 @@ private[quint] object QuintDef {
   }
 
   @key("assume") case class QuintAssume(
-      id: Int,
+      id: BigInt,
       /** name of the assumption, may be '_' */
       name: String,
       /** an expression to associate with the name */
@@ -281,7 +321,7 @@ private[quint] object QuintDef {
    *   - Type aliases always have an associated `type`
    */
   @key("typedef") case class QuintTypeDef(
-      id: Int,
+      id: BigInt,
       /** name of a type alias */
       name: String,
       /**
@@ -295,15 +335,19 @@ private[quint] object QuintDef {
     // We need custom ser/de here to cope with the optionality of the `type` field
     // see https://github.com/com-lihaoyi/upickle/issues/75
     private val toJson: QuintTypeDef => ujson.Value = {
-      case QuintTypeDef(id, name, None) => ujson.Obj("id" -> id, "name" -> name)
+      case QuintTypeDef(id, name, None) =>
+        ujson.Obj("id" -> QuintDeserializer.writeJs(id), "name" -> QuintDeserializer.writeJs(name))
       case QuintTypeDef(id, name, Some(t)) =>
-        ujson.Obj("id" -> id, "name" -> name, "type" -> QuintDeserializer.writeJs[QuintType](t))
+        ujson.Obj("id" -> QuintDeserializer.writeJs(id), "name" -> QuintDeserializer.writeJs(name),
+            "type" -> QuintDeserializer.writeJs(t))
     }
 
     private val ofJson: ujson.Value => QuintTypeDef = {
       case ujson.Obj(entries) if entries.get("id").isDefined && entries.get("name").isDefined =>
-        QuintTypeDef(entries.get("id").get.num.toInt, entries.get("name").get.str,
-            entries.get("type").map(t => QuintDeserializer.read[QuintType](t)))
+        val id = QuintDeserializer.read[BigInt](entries.get("id").get)
+        val name = QuintDeserializer.read[String](entries.get("name").get)
+        val tt = entries.get("type").map(t => QuintDeserializer.read[QuintType](t))
+        QuintTypeDef(id, name, tt)
       case invalidJson =>
         throw new QuintIRParseError(s"Unexpected JSON representation of Quint type definition: ${invalidJson}")
     }

tla-io/src/test/scala/at/forsyte/apalache/io/quint/TestQuintEx.scala

@@ -26,7 +26,7 @@ class TestQuintEx extends AnyFunSuite {
     // expression should use this thunk to produce a
     // unique ID (uid)
     private var nextId = 1
-    private def uid: Int = {
+    private def uid: BigInt = {
       val x = nextId
       nextId += 1
       x
@@ -38,8 +38,8 @@ class TestQuintEx extends AnyFunSuite {
     //
     // These two mutable map are used to construct those mappings, which are
     // then passed to the conversion class.
-    private val typeMap = collection.mutable.Map[Int, QuintType]()
-    private val lookupMap = collection.mutable.Map[Int, QuintLookupTableEntry]()
+    private val typeMap = collection.mutable.Map[BigInt, QuintType]()
+    private val lookupMap = collection.mutable.Map[BigInt, QuintLookupTableEntry]()
 
     // Register the type of an expression in the typeMap.
     // Think of this as a type annotation.
@@ -57,7 +57,7 @@ class TestQuintEx extends AnyFunSuite {
     // table. This mocks quint's practice of recording the type
     // of an applied operator in the lookup table for defined
     // operator, while omitting this data for bulitins.
-    def app(name: String, args: QuintEx*)(retType: QuintType, refId: Int = -1): QuintApp = {
+    def app(name: String, args: QuintEx*)(retType: QuintType, refId: BigInt = -1): QuintApp = {
       val id = uid
       if (refId != -1) {
         lookupMap += (id -> QuintLookupTableEntry("def", refId))

tla-io/src/test/scala/at/forsyte/apalache/io/quint/TestQuintIR.scala

@@ -3,12 +3,40 @@ package at.forsyte.apalache.io.quint
 import org.junit.runner.RunWith
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatestplus.junit.JUnitRunner
-import scala.util.Success
-import scala.util.Failure
+
+import scala.util.{Failure, Success}
 
 @RunWith(classOf[JUnitRunner])
 class TestQuintIR extends AnyFunSuite {
 
+  test("Deserializes BigInt") {
+    val someBigInt = "42" + Long.MaxValue.toString() + "42"
+
+    // unquoted; tests `visitFloat64StringParts`
+    assert(QuintDeserializer.read[BigInt]("0") == BigInt(0))
+    assert(QuintDeserializer.read[BigInt]("1") == BigInt(1))
+    assert(QuintDeserializer.read[BigInt]("-1") == BigInt(-1))
+    assert(QuintDeserializer.read[BigInt](someBigInt) == BigInt(someBigInt))
+    assert(QuintDeserializer.read[BigInt](s"-${someBigInt}") == BigInt("-" + someBigInt))
+
+    // quoted; tests `visitString`
+    assert(QuintDeserializer.read[BigInt]("\"0\"") == BigInt(0))
+    assert(QuintDeserializer.read[BigInt]("\"1\"") == BigInt(1))
+    assert(QuintDeserializer.read[BigInt]("\"-1\"") == BigInt(-1))
+    assert(QuintDeserializer.read[BigInt](s"\"${someBigInt}\"") == BigInt(someBigInt))
+    assert(QuintDeserializer.read[BigInt](s"\"-${someBigInt}\"") == BigInt("-" + someBigInt))
+  }
+
+  test("Serializes BigInt") {
+    val someBigIntStr = "42" + Long.MaxValue.toString() + "42"
+
+    assert(QuintDeserializer.write[BigInt](BigInt(0)) == "0")
+    assert(QuintDeserializer.write[BigInt](BigInt(1)) == "1")
+    assert(QuintDeserializer.write[BigInt](BigInt(-1)) == "-1")
+    assert(QuintDeserializer.write[BigInt](BigInt(someBigIntStr)) == someBigIntStr)
+    assert(QuintDeserializer.write[BigInt](BigInt(s"-${someBigIntStr}")) == s"-${someBigIntStr}")
+  }
+
   // tictactoe.json is located in tla-io/src/test/resources/tictactoe.json
   test("Can load tictactoe.json") {
     val tictactoeQuintJson = scala.io.Source.fromResource("tictactoe.json").mkString