Fixes for compilation to TLA+

.unreleased/bug-fixes/fold-let-extraction.md

@@ -0,0 +1 @@
+Fixed a problem where folds produced by the Shai command `TLA` had an invalid form and could not be parsed back (#2891)

.unreleased/bug-fixes/sanitize-bindings.md

@@ -0,0 +1 @@
+Fixed a problem where bindings from exists and forall operators where not properly sanitized before printing (#2891)

.unreleased/bug-fixes/slice-to-replaceAt.md

@@ -0,0 +1 @@
+Fixed a problem where translation from `slice` to `replaceAt` added an incorrect increment to the last argument (#2891)

.unreleased/features/compilation-type-annotations.md

@@ -0,0 +1 @@
+TLA+ modules produced by the Shai command `TLA` now include type annotations (#2891)

shai/src/main/scala/at/forsyte/apalache/shai/v1/CmdExecutorService.scala

@@ -1,5 +1,6 @@
 package at.forsyte.apalache.shai.v1
 
+import java.io.{PrintWriter, StringWriter}
 import scala.util.Try
 import com.typesafe.scalalogging.Logger
 import io.grpc.Status
@@ -16,7 +17,8 @@ import at.forsyte.apalache.tla.bmcmt.config.CheckerModule
 import at.forsyte.apalache.tla.passes.imp.ParserModule
 import at.forsyte.apalache.tla.passes.typecheck.TypeCheckerModule
 import at.forsyte.apalache.tla.lir.TlaModule
-import at.forsyte.apalache.io.lir.PrettyWriter
+import at.forsyte.apalache.io.annotations.PrettyWriterWithAnnotations
+import at.forsyte.apalache.io.annotations.store._
 
 /**
  * Provides the [[CmdExecutorService]]
@@ -99,7 +101,32 @@ class CmdExecutorService(logger: Logger) extends ZioCmdExecutor.ZCmdExecutor[ZEn
   }
 
   private def tlaModuleToJsonString(module: TlaModule): ujson.Value = {
-    ujson.Str(PrettyWriter.writeAsString(module))
+    val annotationStore = createAnnotationStore()
+
+    val buf = new StringWriter()
+    val prettyWriter = new PrettyWriterWithAnnotations(annotationStore, new PrintWriter(buf))
+    val modules_to_extend = List("Integers", "Sequences", "FiniteSets", "TLC", "Apalache", "Variants")
+    prettyWriter.write(module, modules_to_extend)
+    val moduleString = buf.toString()
+
+    val modifiedModule = extractLetFromFolds(moduleString)
+    ujson.Str(modifiedModule)
+  }
+
+  // Apalache inlines fold operator arguments as LET .. IN expressions, but this
+  // is not valid for SANY. In order to produce a valid TLA+ module from Quint
+  // files, we transform expressions like:
+  // ```
+  // ApaFoldSet(LET __QUINT_LAMBDAn(a, b) == c IN __QUINT_LAMBDAn, init, set)
+  // ```
+  //
+  // into:
+  // ```
+  // LET __QUINT_LAMBDAn(a, b) == c IN ApaFoldSet(__QUINT_LAMBDAn, init, set)
+  // ```
+  private def extractLetFromFolds(module: String): String = {
+    val regex = """(?s)(ApaFold[\w]*\()\s*(LET\s.*?\sIN\s+)(__QUINT_LAMBDA)"""
+    return module.replaceAll(regex, "$2 $1$3")
   }
 
   // Allows us to handle invalid protobuf messages on the ZIO level, before

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

@@ -89,7 +89,7 @@ class PrettyWriter(
     prettyWriteDoc(declToDoc(decl) <> line <> line)
   }
 
-  def write(expr: TlaEx): Unit = prettyWriteDoc(exToDoc((0, 0), expr, sanatizeID))
+  def write(expr: TlaEx): Unit = prettyWriteDoc(exToDoc((0, 0), expr, sanitizeID))
 
   def writeComment(commentStr: String): Unit = {
     prettyWriteDoc(wrapWithComment(commentStr) <> line)
@@ -103,14 +103,14 @@ class PrettyWriter(
   def writeFooter(): Unit = prettyWriteDoc(moduleTerminalDoc)
 
   private def moduleNameDoc(name: String): Doc = {
-    val middle = s" MODULE ${sanatizeID(name)} "
+    val middle = s" MODULE ${sanitizeID(name)} "
     val nDashes = math.max(5, (layout.textWidth - middle.length) / 2) // int div rounds down
     s"${List.fill(nDashes)("-").mkString}$middle${List.fill(nDashes)("-").mkString}" <> line
   }
 
   private def moduleExtendsDoc(moduleNames: List[String]): Doc =
     if (moduleNames.isEmpty) emptyDoc
-    else line <> text("EXTENDS") <> space <> hsep(moduleNames.map(n => text(sanatizeID(n))), comma) <> line
+    else line <> text("EXTENDS") <> space <> hsep(moduleNames.map(n => text(sanitizeID(n))), comma) <> line
 
   private def moduleTerminalDoc: Doc =
     s"${List.fill(layout.textWidth)("=").mkString}" <> line
@@ -209,7 +209,7 @@ class PrettyWriter(
         val sign = PrettyWriter.bindingOps(op)
         val doc =
           group(
-              group(text(sign) <> space <> text(x.toString) <> space <>
+              group(text(sign) <> space <> text(sanitizeID(x.toString)) <> space <>
                 text(PrettyWriter.binaryOps(TlaSetOper.in)) <> softline <>
                 exToDoc(op.precedence, set, nameResolver) <> text(":")) <>
                 nest(line <> exToDoc(op.precedence, pred, nameResolver))
@@ -538,7 +538,7 @@ class PrettyWriter(
             }))
   }
 
-  def declToDoc(decl: TlaDecl, nameResolver: String => String = sanatizeID): Doc = {
+  def declToDoc(decl: TlaDecl, nameResolver: String => String = sanitizeID): Doc = {
     val annotations = declAnnotator(layout)(decl)
 
     decl match {
@@ -680,15 +680,15 @@ class PrettyWriter(
   private val validIdentifierPrefix = """_*[a-zA-Z]""".r
 
   // Sanitize an identifier to ensure it can be read by TLC
-  private def sanatizeID(s: String): String = {
+  private def sanitizeID(s: String): String = {
     val s0 = if (validIdentifierPrefix.findPrefixOf(s).isDefined) s else "id" + s
     invalidIdentifierParts.replaceAllIn(s0, "_")
   }
 
   private def parseableName(name: String): String = {
     // An operator name may contain '!' if it comes from an instance. Replace '!' with "_i_".
     // Additionally, the name may contain '$', which is produced during preprocessing. Replace '$' with "_si_".
-    sanatizeID(name.replaceAll("!", "_i_").replaceAll("\\$", "_si_"))
+    sanitizeID(name.replaceAll("!", "_i_").replaceAll("\\$", "_si_"))
   }
 
   def close(): Unit = writer.close()

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

@@ -558,7 +558,7 @@ class Quint(quintOutput: QuintOutput) {
         case "foldl"     => ternaryApp(opName, (seq, init, op) => tla.foldSeq(op, init, seq))
         case "nth"       => binaryApp(opName, (seq, idx) => tla.app(seq, incrTla(idx)))
         case "replaceAt" => ternaryApp(opName, (seq, idx, x) => tla.except(seq, incrTla(idx), x))
-        case "slice"     => ternaryApp(opName, (seq, from, to) => tla.subseq(seq, incrTla(from), incrTla(to)))
+        case "slice"     => ternaryApp(opName, (seq, from, to) => tla.subseq(seq, incrTla(from), to))
         case "select"    => MkTla.selectSeq(opName, typeConv.convert(types(id).typ))
         case "range" =>
           binaryApp(opName,

tla-io/src/test/scala/at/forsyte/apalache/io/lir/TestPrettyWriter.scala

@@ -379,6 +379,16 @@ class TestPrettyWriter extends AnyFunSuite with BeforeAndAfterEach {
     assert(expected == stringWriter.toString)
   }
 
+  test("an exists with a binding with invalid characters") {
+    val writer = new PrettyWriter(printWriter, layout40)
+    val expr = exists(name("a::x"), name("a::y"), name("a::z"))
+    writer.write(expr)
+    printWriter.flush()
+    // a multi-line vertical box always breaks from the previous line, as otherwise it is incredibly hard to indent
+    val expected = "\\E a_x \\in a_y: a_z"
+    assert(expected == stringWriter.toString)
+  }
+
   test("nested \\EE and \\AA") {
     val writer = new PrettyWriter(printWriter, layout10)
     val ex1 =

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

@@ -513,7 +513,7 @@ class TestQuintEx extends AnyFunSuite {
 
   test("can convert builtin slice operator application") {
     assert(convert(Q.app("slice", Q.intList, Q._0, Q._1)(
-            QuintSeqT(QuintIntT()))) == "Sequences!SubSeq(<<1, 2, 3>>, 0 + 1, 1 + 1)")
+            QuintSeqT(QuintIntT()))) == "Sequences!SubSeq(<<1, 2, 3>>, 0 + 1, 1)")
   }
 
   test("can convert builtin select operator application") {