Merge pull request #2808 from fan-tom/parse-assume-name

Parse ASSUME declarations names

.unreleased/features/01.md

@@ -0,0 +1 @@
+Parse name of `ASSUME` declarations into IR and preserve them during serialization to JSON, see #2808

tla-io/src/main/scala/at/forsyte/apalache/io/json/JsonToTla.scala

@@ -157,9 +157,10 @@ class JsonToTla[T <: JsonRepresentation](
         opDecl
 
       case "TlaAssumeDecl" =>
+        val definedName = declJson.getFieldOpt("name").map(scalaFactory.asStr)
         val bodyField = getOrThrow(declJson, "body")
         val body = asTlaEx(bodyField)
-        TlaAssumeDecl(body)(typeTag)
+        TlaAssumeDecl(definedName, body)(typeTag)
       case _ => throw new JsonDeserializationError(s"$kind is not a valid TlaDecl kind")
     }
     setLoc(decl, getSourceLocationOpt(declJson))

tla-io/src/main/scala/at/forsyte/apalache/io/json/JsonToTlaViaBuilder.scala

@@ -190,9 +190,10 @@ class JsonToTlaViaBuilder[T <: JsonRepresentation](
         opDecl
 
       case "TlaAssumeDecl" =>
+        val definedName = declJson.getFieldOpt("name").map(scalaFactory.asStr)
         val bodyField = getOrThrow(declJson, "body")
         val body = asTBuilderInstruction(bodyField)
-        TlaAssumeDecl(body)(typeTag)
+        TlaAssumeDecl(definedName, body)(typeTag)
       case _ => throw new JsonDeserializationError(s"$kind is not a valid TlaDecl kind")
     }
     setLoc(decl, getSourceLocationOpt(declJson))

tla-io/src/main/scala/at/forsyte/apalache/io/json/TlaToJson.scala

@@ -188,13 +188,17 @@ class TlaToJson[T <: JsonRepresentation](
             "body" -> bodyJson,
         )
 
-      case TlaAssumeDecl(body) =>
+      case TlaAssumeDecl(definedName, body) =>
         val bodyJson = apply(body)
-        withLoc(
+
+        val optionalName: Option[(String, T)] = definedName.map("name" -> _)
+        val fields = Seq[(String, T)](
             typeFieldName -> typeTagPrinter(decl.typeTag),
             kindFieldName -> "TlaAssumeDecl",
             "body" -> bodyJson,
-        )
+        ) ++ optionalName
+
+        withLoc(fields: _*)
     }
   }
 

tla-io/src/main/scala/at/forsyte/apalache/io/lir/PrettyWriter.scala

@@ -556,8 +556,12 @@ class PrettyWriter(
           "VARIABLE" <> nest(line <> wrapWithComment(annotations.get) <> line <> parseableName(name))
         }
 
-      case TlaAssumeDecl(body) =>
-        val doc = group("ASSUME" <> parens(exToDoc((0, 0), body, nameResolver)))
+      case TlaAssumeDecl(definedName, body) =>
+        val doc = definedName match {
+          case None       => group("ASSUME" <> parens(exToDoc((0, 0), body, nameResolver)))
+          case Some(name) => group("ASSUME" <+> name <+> "==" <+> parens(exToDoc((0, 0), body, nameResolver)))
+        }
+
         if (annotations.isEmpty) {
           doc
         } else {

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

@@ -737,10 +737,11 @@ class Quint(quintOutput: QuintOutput) {
         // no methods for them are provided by the ScopedBuilder.
         case QuintConst(id, name, _) => Some(None, TlaConstDecl(name)(typeTagOfId(id)))
         case QuintVar(id, name, _)   => Some(None, TlaVarDecl(name)(typeTagOfId(id)))
-        case QuintAssume(_, _, quintEx) =>
+        case d @ QuintAssume(_, name, quintEx) =>
           val tlaEx = build(tlaExpression(quintEx).run(nullaryOps))
+          val definedName = Option.unless(d.isUnnamed)(name)
           // assume declarations have no entry in the type map, and are always typed bool
-          Some(None, TlaAssumeDecl(tlaEx)(Typed(BoolT1)))
+          Some(None, TlaAssumeDecl(definedName, tlaEx)(Typed(BoolT1)))
         case op: QuintOpDef if op.qualifier == "run" =>
           // We don't currently support run definitions
           None

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

@@ -310,7 +310,16 @@ private[quint] object QuintDef {
       name: String,
       /** an expression to associate with the name */
       assumption: QuintEx)
-      extends QuintDef {}
+      extends QuintDef {
+
+    /**
+     * @return
+     *   `true` if this assume declaration has no user-defined name, `false` otherwise
+     *
+     * anonymous assume declarations are named with the "hole", `_`
+     */
+    def isUnnamed: Boolean = name == "_"
+  }
   object QuintAssume {
     implicit val rw: RW[QuintAssume] = macroRW
   }

tla-io/src/main/scala/at/forsyte/apalache/tla/imp/AssumeTranslator.scala

@@ -17,7 +17,8 @@ class AssumeTranslator(
           context,
           OutsideRecursion(),
       ).translate(node.getAssume)
-    TlaAssumeDecl(body)(Untyped)
+    val name = Option(node.getDef).map(_.getName.toString)
+    TlaAssumeDecl(name, body)(Untyped)
   }
 }
 

tla-io/src/test/scala/at/forsyte/apalache/io/TestPrettyWriterWithTypes.scala

@@ -97,4 +97,38 @@ class TestPrettyWriterWithTypes extends AnyFunSuite with BeforeAndAfterEach {
         |""".stripMargin
     assert(expected == stringWriter.toString)
   }
+
+  test("unnamed assume declaration") {
+    val decl = TlaAssumeDecl(None, tla.eql(tla.name("x"), tla.bool(true)))(Typed(BoolT1))
+    val store = createAnnotationStore()
+
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
+    writer.write(decl)
+    printWriter.flush()
+    val expected =
+      """(*
+        |  @type: Bool;
+        |*)
+        |ASSUME(x = TRUE)
+        |
+        |""".stripMargin
+    assert(expected == stringWriter.toString)
+  }
+
+  test("named assume declaration") {
+    val decl = TlaAssumeDecl(Some("myAssume"), tla.eql(tla.name("x"), tla.bool(true)))(Typed(BoolT1))
+    val store = createAnnotationStore()
+
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
+    writer.write(decl)
+    printWriter.flush()
+    val expected =
+      """(*
+        |  @type: Bool;
+        |*)
+        |ASSUME myAssume == (x = TRUE)
+        |
+        |""".stripMargin
+    assert(expected == stringWriter.toString)
+  }
 }

tla-io/src/test/scala/at/forsyte/apalache/io/json/TestTlaToUJson.scala

@@ -1,7 +1,7 @@
 package at.forsyte.apalache.io.json
 
 import at.forsyte.apalache.io.json.impl.TlaToUJson
-import at.forsyte.apalache.tla.lir.{TestingPredefs, TlaConstDecl, TlaDecl, TlaEx, TlaVarDecl, TypeTag}
+import at.forsyte.apalache.tla.lir.{TestingPredefs, TlaAssumeDecl, TlaConstDecl, TlaDecl, TlaEx, TlaVarDecl, TypeTag}
 import at.forsyte.apalache.tla.lir.convenience.tla
 import at.forsyte.apalache.tla.lir.oper.{TlaFunOper, TlaSetOper}
 import org.junit.runner.RunWith
@@ -89,15 +89,25 @@ class TestTlaToUJson extends AnyFunSuite with BeforeAndAfterEach with TestingPre
   test("Non-operator declarations") {
     val constDecl = TlaConstDecl("C")
     val varDecl = TlaVarDecl("x")
+    val namedAssumeDecl = TlaAssumeDecl(Some("myAssume"), tla.eql(tla.name("x"), tla.bool(true)))
+    val unnamedAssumeDecl = TlaAssumeDecl(None, tla.eql(tla.name("x"), tla.bool(true)))
 
     val constJson = getEncVal(constDecl)
     val varJson = getEncVal(varDecl)
+    val namedAssumeJson = getEncVal(namedAssumeDecl)
+    val unnamedAssumeJson = getEncVal(unnamedAssumeDecl)
 
     assert(constJson(kindField).str == "TlaConstDecl")
     assert(constJson("name").str == constDecl.name)
 
     assert(varJson(kindField).str == "TlaVarDecl")
     assert(varJson("name").str == varDecl.name)
+
+    assert(namedAssumeJson(kindField).str == "TlaAssumeDecl")
+    assert(namedAssumeJson("name").str == namedAssumeDecl.name)
+
+    assert(unnamedAssumeJson(kindField).str == "TlaAssumeDecl")
+    assert(!unnamedAssumeJson.obj.contains("name"))
   }
 
   test("Operator declarations") {

tla-io/src/test/scala/at/forsyte/apalache/io/json/TestUJsonToTla.scala

@@ -45,7 +45,7 @@ class TestUJsonToTla extends AnyFunSuite with Checkers {
 
     val decls: Seq[TlaDecl] = Seq(
         tla.declOp("X", tla.eql(tla.name("a"), tla.int(1)), OperParam("a")),
-        TlaAssumeDecl(tla.eql(tla.int(1), tla.int(0))),
+        TlaAssumeDecl(None, tla.eql(tla.int(1), tla.int(0))),
         TlaConstDecl("c"),
         TlaVarDecl("v"),
     )

tla-io/src/test/scala/at/forsyte/apalache/io/json/TestUJsonToTlaViaBuilder.scala

@@ -57,7 +57,7 @@ class TestUJsonToTlaViaBuilder extends AnyFunSuite with Checkers {
 
     val decls: Seq[TlaDecl] = Seq(
         tla.decl("X", tla.eql(tla.name("a", IntT1), tla.int(1)), tla.param("a", IntT1)),
-        TlaAssumeDecl(tla.eql(tla.int(1), tla.int(0))),
+        TlaAssumeDecl(None, tla.eql(tla.int(1), tla.int(0))),
         TlaConstDecl("c"),
         TlaVarDecl("v"),
     )

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

@@ -1,11 +1,15 @@
 package at.forsyte.apalache.io.quint
 
-import at.forsyte.apalache.tla.lir.{IntT1, OperT1, RecRowT1, RowT1, SetT1, TlaEx, TlaType1, Typed, VarT1}
+import at.forsyte.apalache.tla.lir.{
+  BoolT1, IntT1, OperT1, RecRowT1, RowT1, SetT1, TlaAssumeDecl, TlaEx, TlaType1, Typed, ValEx, VarT1,
+}
 import org.junit.runner.RunWith
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatestplus.junit.JUnitRunner
 import QuintType._
 import QuintEx._
+import at.forsyte.apalache.io.quint.QuintDef.QuintAssume
+import at.forsyte.apalache.tla.lir.values.TlaBool
 import at.forsyte.apalache.tla.typecomp.build
 
 // You can run all these tests in watch mode in the
@@ -789,4 +793,18 @@ class TestQuintEx extends AnyFunSuite {
     val err = intercept[QuintIRParseError](convert(Q.oneOfSet))
     assert(err.getMessage().contains("`oneOf` can only occur as the principle operator of a `nondet` declaration"))
   }
+
+  test("can convert ASSUME declaration") {
+    val translator = new Quint(Q.quintOutput)
+    val nullaryOps = Set[String]()
+
+    val name = "myAssume"
+    val namedAssume = QuintAssume(1, name, QuintBool(2, true))
+    val tlaNamedAssumeDef = translator.tlaDef(namedAssume).run(nullaryOps).get._2
+    assert(tlaNamedAssumeDef == TlaAssumeDecl(Some(name), ValEx(TlaBool(true))(Typed(BoolT1)))(Typed(BoolT1)))
+
+    val unnamedAssume = QuintAssume(1, "_", QuintBool(2, true))
+    val tlaUnnamedAssumeDef = translator.tlaDef(unnamedAssume).run(nullaryOps).get._2
+    assert(tlaUnnamedAssumeDef == TlaAssumeDecl(None, ValEx(TlaBool(true))(Typed(BoolT1)))(Typed(BoolT1)))
+  }
 }

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

@@ -1,5 +1,7 @@
 package at.forsyte.apalache.io.quint
 
+import at.forsyte.apalache.io.quint.QuintDef.QuintAssume
+import at.forsyte.apalache.io.quint.QuintEx.QuintBool
 import org.junit.runner.RunWith
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatestplus.junit.JUnitRunner
@@ -37,6 +39,14 @@ class TestQuintIR extends AnyFunSuite {
     assert(QuintDeserializer.write[BigInt](BigInt(s"-${someBigIntStr}")) == s"-${someBigIntStr}")
   }
 
+  test("QuintAssume#isUnnamed works correctly") {
+    val namedAssume = QuintAssume(1, "myAssume", QuintBool(2, true))
+    assert(!namedAssume.isUnnamed)
+
+    val unnamedAssume = QuintAssume(1, "_", QuintBool(2, true))
+    assert(unnamedAssume.isUnnamed)
+  }
+
   // 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

tla-io/src/test/scala/at/forsyte/apalache/tla/imp/TestSanyImporter.scala

@@ -1808,12 +1808,14 @@ class TestSanyImporter extends SanyImporterTestBase {
 
   test("assumptions") {
     // checking that the assumptions are imported properly
+    val assumptionName = "nonZero"
     val text =
-      """
+      s"""
         |---- MODULE assumptions ----
         |CONSTANT N
         |ASSUME N = 4
         |ASSUME N /= 10
+        |ASSUME $assumptionName == N /= 0
         |================================
         |""".stripMargin
 
@@ -1825,20 +1827,28 @@ class TestSanyImporter extends SanyImporterTestBase {
     expectSourceInfoInDefs(root)
 
     modules(rootName).declarations(1) match {
-      case TlaAssumeDecl(e) => assert(eql(name("N"), int(4)).untyped() == e)
+      case TlaAssumeDecl(_, e) => assert(eql(name("N"), int(4)).untyped() == e)
 
       case e @ _ => fail("expected an assumption, found: " + e)
     }
 
     modules(rootName).declarations(2) match {
-      case TlaAssumeDecl(e) => assert(neql(name("N"), int(10)).untyped() == e)
+      case TlaAssumeDecl(_, e) => assert(neql(name("N"), int(10)).untyped() == e)
+
+      case e @ _ => fail("expected an assumption, found: " + e)
+    }
+
+    modules(rootName).declarations(3) match {
+      case TlaAssumeDecl(definedName, e) =>
+        assert(neql(name("N"), int(0)).untyped() == e)
+        assert(definedName contains assumptionName)
 
       case e @ _ => fail("expected an assumption, found: " + e)
     }
 
     // regression test for issue #25
     val names = HashSet(modules(rootName).assumeDeclarations.map(_.name): _*)
-    assert(2 == names.size) // all assumptions must have unique names
+    assert(3 == names.size) // all assumptions must have unique names
   }
 
   test("ignore theorems") {

tla-typechecker/src/main/scala/at/forsyte/apalache/tla/typecheck/etc/ToEtcExpr.scala

@@ -45,7 +45,12 @@ class ToEtcExpr(
         val operType = OperT1(Seq(BoolT1), BoolT1)
         val application = mkUniqApp(Seq(operType), this(d.body))
         // We have to introduce a lambda abstraction, as the type checker is expecting this form.
-        mkLet(BlameRef(d.ID), "__Assume_" + d.ID, mkAbs(ExactRef(d.ID), application), inScopeEx)
+        mkLet(
+            BlameRef(d.ID),
+            d.name,
+            mkAbs(ExactRef(d.ID), application),
+            inScopeEx,
+        )
 
       case d: TlaOperDecl =>
         // Foo(x) == ...
@@ -131,7 +136,7 @@ class ToEtcExpr(
     OperT1(nBools, BoolT1)
   }
 
-  // Valid when the input seq has two items, the first of which is a VlaEx(TlaStr(_))
+  // Valid when the input seq has two items, the first of which is a ValEx(TlaStr(_))
   private val validateRecordPair: Seq[TlaEx] => (String, TlaEx) = {
     // Only pairs coordinating pairs and sets are valid. See TlaSetOper.recSet
     case Seq(ValEx(TlaStr(name)), set) =>

tla-typechecker/src/main/scala/at/forsyte/apalache/tla/typecheck/integration/TypeRewriter.scala

@@ -74,8 +74,8 @@ class TypeRewriter(tracker: TransformationTracker, defaultTag: UID => TypeTag)(t
       case d @ TlaVarDecl(_) =>
         decl.withTag(getOrDefault(d.ID))
 
-      case d @ TlaAssumeDecl(body) =>
-        TlaAssumeDecl(this(body))(getOrDefault(d.ID))
+      case d @ TlaAssumeDecl(definedName, body) =>
+        TlaAssumeDecl(definedName, this(body))(getOrDefault(d.ID))
 
       case d @ TlaTheoremDecl(name, body) =>
         TlaTheoremDecl(name, this(body))(getOrDefault(d.ID))

tla-typechecker/src/test/scala/at/forsyte/apalache/tla/typecheck/etc/TestToEtcExprDecls.scala

@@ -107,7 +107,7 @@ class TestToEtcExprDecls extends AnyFunSuite with ToEtcExprBase with BeforeAndAf
   }
 
   test("assumes") {
-    val assume = TlaAssumeDecl(tla.name("x"))
+    val assume = TlaAssumeDecl(None, tla.name("x"))
     val terminal = mkUniqConst(BoolT1)
     // becomes:
     // let Assume1 == ((Bool => Bool) "x") in
@@ -116,7 +116,8 @@ class TestToEtcExprDecls extends AnyFunSuite with ToEtcExprBase with BeforeAndAf
     // The body is wrapped with the application of an operator that has the signature Bool => Bool.
     // This allows us to check that the assumption has Boolean type.
     val application = mkUniqApp(Seq(parser("Bool => Bool")), assumption)
-    val expected = mkUniqLet("__Assume_" + assume.ID, mkUniqAbs(application), terminal)
+    val expected =
+      mkUniqLet(assume.name, mkUniqAbs(application), terminal)
     // Translate the declaration of positive.
     // We have to pass the next expression in scope, which is just TRUE in this case.
     assert(expected == mkToEtcExpr(Map())(assume, terminal))

tlair/src/main/scala/at/forsyte/apalache/tla/lir/TlaDecl.scala

@@ -60,13 +60,16 @@ case class TlaVarDecl(name: String)(implicit typeTag: TypeTag) extends TlaDecl w
 /**
  * An assumption defined by ASSUME(...)
  *
+ * @param definedName
+ *   optional assumption name, like name in `ASSUME name == x = 4`, or none, like in `ASSUME x = 4`
  * @param body
  *   the assumption body
  */
-case class TlaAssumeDecl(body: TlaEx)(implicit typeTag: TypeTag) extends TlaDecl with Serializable {
-  val name: String = "ASSUME" + body.ID
+case class TlaAssumeDecl(definedName: Option[String], body: TlaEx)(implicit typeTag: TypeTag)
+    extends TlaDecl with Serializable {
+  override val name: String = definedName.getOrElse("ASSUME" + body.ID)
 
-  override def withTag(newTypeTag: TypeTag): TlaAssumeDecl = TlaAssumeDecl(body)(newTypeTag)
+  override def withTag(newTypeTag: TypeTag): TlaAssumeDecl = TlaAssumeDecl(definedName, body)(newTypeTag)
 }
 
 /**

tlair/src/main/scala/at/forsyte/apalache/tla/lir/TlaLevelFinder.scala

@@ -74,7 +74,7 @@ class TlaLevelFinder(module: TlaModule) {
       case TlaVarDecl(_) =>
         TlaLevelState
 
-      case TlaAssumeDecl(_) =>
+      case TlaAssumeDecl(_, _) =>
         TlaLevelConst
 
       case TlaOperDecl(name, _, body) =>

tlair/src/main/scala/at/forsyte/apalache/tla/lir/io/Printer.scala

@@ -243,8 +243,11 @@ object UTFPrinter extends Printer {
       case TlaVarDecl(name) =>
         "VARIABLE " + name
 
-      case TlaAssumeDecl(body) =>
-        apply(body)
+      case TlaAssumeDecl(definedName, body) =>
+        definedName match {
+          case None       => s"ASSUME " + apply(body)
+          case Some(name) => s"ASSUME $name ${m_defeq} " + apply(body)
+        }
 
       case TlaOperDecl(name, formalParams, body) =>
         val ps =