Fix pretty printer: write annotations correctly and extend the standard modules

UNRELEASED.md

@@ -26,6 +26,8 @@
 * Parser: better parser for annotations, see #757
 * Parser: fixed two bugs in the declaration sorter, see #645 and #758
 * Printer: fixed the output for EXCEPT, see #746
+* Printer: fixed pretty printing of annotations, see #633
+* Printer: extending the standard modules, see #137
 * The command `config --enable-stats=true` creates `$HOME/.tlaplus` if needed, see #762
 
 ### Changes

tla-assignments/src/main/scala/at/forsyte/apalache/tla/assignments/passes/PrimingPassImpl.scala

@@ -4,7 +4,7 @@ import java.io.File
 import java.nio.file.Path
 import at.forsyte.apalache.infra.passes.{Pass, PassOptions, TlaModuleMixin}
 import at.forsyte.apalache.tla.lir._
-import at.forsyte.apalache.tla.lir.io.{PrettyWriter, TlaWriterFactory}
+import at.forsyte.apalache.tla.lir.io.{PrettyWriter, TlaWriter, TlaWriterFactory}
 import at.forsyte.apalache.tla.lir.storage.BodyMapFactory
 import at.forsyte.apalache.tla.lir.transformations.{TlaExTransformation, TransformationTracker}
 import at.forsyte.apalache.tla.lir.transformations.standard._
@@ -84,7 +84,7 @@ class PrimingPassImpl @Inject() (options: PassOptions, tracker: TransformationTr
     val newModule = new TlaModule(tlaModule.get.name, newDeclarations)
 
     val outdir = options.getOrError("io", "outdir").asInstanceOf[Path]
-    writerFactory.writeModuleToFile(newModule, new File(outdir.toFile, "out-priming.tla"))
+    writerFactory.writeModuleToFile(newModule, TlaWriter.STANDARD_MODULES, new File(outdir.toFile, "out-priming.tla"))
 
     setModule(newModule)
     true

tla-assignments/src/main/scala/at/forsyte/apalache/tla/assignments/passes/TransitionPassImpl.scala

@@ -5,7 +5,7 @@ import at.forsyte.apalache.tla.assignments._
 import at.forsyte.apalache.tla.imp.findBodyOf
 import at.forsyte.apalache.tla.imp.src.SourceStore
 import at.forsyte.apalache.tla.lir._
-import at.forsyte.apalache.tla.lir.io.TlaWriterFactory
+import at.forsyte.apalache.tla.lir.io.{TlaWriter, TlaWriterFactory}
 import at.forsyte.apalache.tla.lir.storage.{ChangeListener, SourceLocator}
 import at.forsyte.apalache.tla.lir.transformations.TransformationTracker
 import at.forsyte.apalache.tla.lir.transformations.standard.IncrementalRenaming
@@ -78,7 +78,8 @@ class TransitionPassImpl @Inject() (options: PassOptions, sourceStore: SourceSto
 
     // print the resulting module
     val outdir = options.getOrError("io", "outdir").asInstanceOf[Path]
-    writerFactory.writeModuleToFile(outModule, new File(outdir.toFile, "out-transition.tla"))
+    writerFactory.writeModuleToFile(outModule, TlaWriter.STANDARD_MODULES,
+        new File(outdir.toFile, "out-transition.tla"))
 
     setModule(outModule)
     true

tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/passes/AnalysisPassImpl.scala

@@ -3,7 +3,7 @@ package at.forsyte.apalache.tla.bmcmt.passes
 import at.forsyte.apalache.infra.passes.{Pass, PassOptions, TlaModuleMixin}
 import at.forsyte.apalache.tla.bmcmt.CheckerException
 import at.forsyte.apalache.tla.bmcmt.analyses._
-import at.forsyte.apalache.tla.lir.io.TlaWriterFactory
+import at.forsyte.apalache.tla.lir.io.{TlaWriter, TlaWriterFactory}
 import at.forsyte.apalache.tla.lir.transformations.TransformationTracker
 import at.forsyte.apalache.tla.lir.transformations.standard.ModuleByExTransformer
 import at.forsyte.apalache.tla.lir.{NullEx, TlaAssumeDecl, TlaEx, TlaOperDecl}
@@ -77,7 +77,7 @@ class AnalysisPassImpl @Inject() (val options: PassOptions, hintsStoreImpl: Form
     nextPass.setModule(marked)
 
     val outdir = options.getOrError("io", "outdir").asInstanceOf[Path]
-    writerFactory.writeModuleToFile(marked, new File(outdir.toFile, "out-analysis.tla"))
+    writerFactory.writeModuleToFile(marked, TlaWriter.STANDARD_MODULES, new File(outdir.toFile, "out-analysis.tla"))
 
     logger.info("  > Introduced expression grades")
     logger.info("  > Introduced %d formula hints".format(hintsStoreImpl.store.size))

tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/passes/VCGenPassImpl.scala

@@ -5,7 +5,7 @@ import java.nio.file.Path
 import at.forsyte.apalache.infra.passes.{Pass, PassOptions, TlaModuleMixin}
 import at.forsyte.apalache.tla.bmcmt.{CheckerException, VCGenerator}
 import at.forsyte.apalache.tla.lir.NullEx
-import at.forsyte.apalache.tla.lir.io.{PrettyWriter, TlaWriterFactory}
+import at.forsyte.apalache.tla.lir.io.{PrettyWriter, TlaWriter, TlaWriterFactory}
 import at.forsyte.apalache.tla.lir.transformations.TransformationTracker
 import at.forsyte.apalache.tla.lir.UntypedPredefs._
 import com.google.inject.Inject
@@ -51,7 +51,7 @@ class VCGenPassImpl @Inject() (options: PassOptions, tracker: TransformationTrac
       }
 
     val outdir = options.getOrError("io", "outdir").asInstanceOf[Path]
-    writerFactory.writeModuleToFile(newModule, new File(outdir.toFile, "out-vcgen.tla"))
+    writerFactory.writeModuleToFile(newModule, TlaWriter.STANDARD_MODULES, new File(outdir.toFile, "out-vcgen.tla"))
 
     nextPass.setModule(newModule)
     true

tla-import/src/main/scala/at/forsyte/apalache/io/annotations/PrettyWriterWithAnnotations.scala

@@ -1,8 +1,8 @@
 package at.forsyte.apalache.io.annotations
 
 import at.forsyte.apalache.io.annotations.store.AnnotationStore
-import at.forsyte.apalache.tla.lir.TlaDecl
-import at.forsyte.apalache.tla.lir.io.{PrettyWriter, TlaWriter}
+import at.forsyte.apalache.tla.lir.{TlaDecl, TlaEx, TlaModule}
+import at.forsyte.apalache.tla.lir.io.{PrettyWriter, TextLayout, TlaDeclAnnotator, TlaWriter}
 
 import java.io.PrintWriter
 
@@ -11,23 +11,48 @@ import java.io.PrintWriter
  *
  * @author Igor Konnov
  */
-class PrettyWriterWithAnnotations(annotationStore: AnnotationStore, writer: PrintWriter, textWidth: Int = 80,
-    indent: Int = 2)
-    extends PrettyWriter(writer, textWidth, indent) with TlaWriter {
-
-  // override the translation of a declaration
-  override def toDoc(decl: TlaDecl): Doc = {
-    val underlyingDoc = super.toDoc(decl)
-    annotationStore
-      .get(decl.ID)
-      .flatMap { annotations =>
-        Some(annotations.foldRight(underlyingDoc) { (anno, doc) =>
-          // A simple implementation that is using Annotation.toPrettyString.
-          // In the future, we should implement a full pretty printer both for annotations and types.
-          "(*" <> space <> anno.toPrettyString <> space <> "*)" <>
-            linebreak <> doc
-        })
+class PrettyWriterWithAnnotations(annotationStore: AnnotationStore, writer: PrintWriter,
+    layout: TextLayout = TextLayout())
+    extends TlaWriter {
+
+  private object annotator extends TlaDeclAnnotator {
+    override def apply(layout: TextLayout)(decl: TlaDecl): Option[List[String]] = {
+      annotationStore.get(decl.ID) match {
+        case None | Some(List()) =>
+          None
+
+        case Some(annotations) =>
+          Some(annotations.map(_.toPrettyString))
       }
-      .getOrElse(underlyingDoc)
+    }
+  }
+
+  private val prettyWriter: PrettyWriter = new PrettyWriter(writer, layout, annotator)
+
+  /**
+   * Write a module, including all declarations
+   *
+   * @param mod a module
+   */
+  override def write(mod: TlaModule, extendedModuleNames: List[String]): Unit = {
+    prettyWriter.write(mod, extendedModuleNames)
+  }
+
+  /**
+   * Write a declaration, including all expressions
+   *
+   * @param decl a declaration
+   */
+  override def write(decl: TlaDecl): Unit = {
+    prettyWriter.write(decl)
+  }
+
+  /**
+   * Write a TLA+ expression.
+   *
+   * @param expr an expression
+   */
+  override def write(expr: TlaEx): Unit = {
+    prettyWriter.write(expr)
   }
 }

tla-import/src/main/scala/at/forsyte/apalache/tla/imp/passes/SanyParserPassImpl.scala

@@ -1,14 +1,14 @@
 package at.forsyte.apalache.tla.imp.passes
 
 import at.forsyte.apalache.infra.passes.{Pass, PassOptions, TlaModuleMixin}
-import at.forsyte.apalache.tla.imp.{SanyImporter, SanyImporterException}
 import at.forsyte.apalache.io.annotations.store._
 import at.forsyte.apalache.tla.imp.src.SourceStore
-import at.forsyte.apalache.tla.lir.{CyclicDependencyError, TlaModule}
-import at.forsyte.apalache.tla.lir.io.{JsonReader, JsonWriter, PrettyWriter, TlaWriterFactory}
+import at.forsyte.apalache.tla.imp.{SanyImporter, SanyImporterException}
+import at.forsyte.apalache.tla.lir.UntypedPredefs._
+import at.forsyte.apalache.tla.lir.io.{JsonReader, JsonWriter, TlaWriter, TlaWriterFactory}
 import at.forsyte.apalache.tla.lir.storage.{ChangeListener, SourceLocator}
 import at.forsyte.apalache.tla.lir.transformations.standard.DeclarationSorter
-import at.forsyte.apalache.tla.lir.UntypedPredefs._
+import at.forsyte.apalache.tla.lir.{CyclicDependencyError, TlaModule}
 import com.google.inject.Inject
 import com.google.inject.name.Named
 import com.typesafe.scalalogging.LazyLogging
@@ -73,10 +73,8 @@ class SanyParserPassImpl @Inject() (
           }
         // save the output
         val outdir = options.getOrError("io", "outdir").asInstanceOf[Path]
-        writerFactory.writeModuleToFile(
-            rootModule.get,
-            new File(outdir.toFile, "out-parser.tla")
-        )
+        writerFactory.writeModuleToFile(rootModule.get, TlaWriter.STANDARD_MODULES,
+            new File(outdir.toFile, "out-parser.tla"))
         JsonWriter.write(
             rootModule.get,
             new File(outdir.toFile, "out-parser.json")
@@ -86,7 +84,7 @@ class SanyParserPassImpl @Inject() (
         val output = options.getOrElse("parser", "output", "")
         if (output.nonEmpty) {
           if (output.contains(".tla"))
-            writerFactory.writeModuleToFile(rootModule.get, new File(output))
+            writerFactory.writeModuleToFile(rootModule.get, TlaWriter.STANDARD_MODULES, new File(output))
           else if (output.contains(".json"))
             JsonWriter.write(rootModule.get, new File(output))
           else

tla-import/src/test/scala/at/forsyte/apalache/io/TestPrettyWriterWithAnnotations.scala

@@ -5,6 +5,7 @@ import at.forsyte.apalache.io.annotations.store.createAnnotationStore
 import at.forsyte.apalache.tla.lir._
 import UntypedPredefs._
 import at.forsyte.apalache.tla.lir.convenience.tla
+import at.forsyte.apalache.tla.lir.io.TextLayout
 import org.junit.runner.RunWith
 import org.scalatest.junit.JUnitRunner
 import org.scalatest.{BeforeAndAfterEach, FunSuite}
@@ -15,6 +16,7 @@ import java.io.{PrintWriter, StringWriter}
 class TestPrettyWriterWithAnnotations extends FunSuite with BeforeAndAfterEach {
   private var stringWriter: StringWriter = _
   private var printWriter: PrintWriter = _
+  private val layout80 = TextLayout().copy(textWidth = 80)
 
   override protected def beforeEach(): Unit = {
     stringWriter = new StringWriter()
@@ -26,12 +28,29 @@ class TestPrettyWriterWithAnnotations extends FunSuite with BeforeAndAfterEach {
     val store = createAnnotationStore()
     store += decl.ID -> List(Annotation("type", AnnotationStr("Int -> Bool")))
 
-    val writer = new PrettyWriterWithAnnotations(store, printWriter, 80)
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
     writer.write(decl)
     printWriter.flush()
     val expected =
-      """(* @type: Int -> Bool; *)
-        |VARIABLE myFun
+      """VARIABLE
+        |  (*
+        |    @type: Int -> Bool;
+        |  *)
+        |  myFun
+        |
+        |""".stripMargin
+    assert(expected == stringWriter.toString)
+  }
+
+  test("variable declaration without annotation") {
+    val decl = TlaVarDecl("myFun")
+    val store = createAnnotationStore()
+
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
+    writer.write(decl)
+    printWriter.flush()
+    val expected =
+      """VARIABLE myFun
         |
         |""".stripMargin
     assert(expected == stringWriter.toString)
@@ -42,13 +61,16 @@ class TestPrettyWriterWithAnnotations extends FunSuite with BeforeAndAfterEach {
     val store = createAnnotationStore()
     store += decl.ID -> List(Annotation("type", AnnotationStr("Int")), Annotation("sweet", AnnotationBool(true)))
 
-    val writer = new PrettyWriterWithAnnotations(store, printWriter, 80)
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
     writer.write(decl)
     printWriter.flush()
     val expected =
-      """(* @type: Int; *)
-        |(* @sweet(true) *)
-        |CONSTANT N
+      """CONSTANT
+        |  (*
+        |    @type: Int;
+        |    @sweet(true)
+        |  *)
+        |  N
         |
         |""".stripMargin
     assert(expected == stringWriter.toString)
@@ -59,14 +81,65 @@ class TestPrettyWriterWithAnnotations extends FunSuite with BeforeAndAfterEach {
     val store = createAnnotationStore()
     store += decl.ID -> List(Annotation("type", AnnotationStr("(Int, Str) -> Bool")))
 
-    val writer = new PrettyWriterWithAnnotations(store, printWriter, 80)
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
     writer.write(decl)
     printWriter.flush()
     val expected =
-      """(* @type: (Int, Str) -> Bool; *)
+      """(*
+        |  @type: (Int, Str) -> Bool;
+        |*)
         |MyOper(x, y) == TRUE
         |
         |""".stripMargin
     assert(expected == stringWriter.toString)
   }
+
+  test("recursive operator declaration") {
+    val decl = TlaOperDecl("RecOper", List(OperParam("x"), OperParam("y")), tla.bool(true))
+    decl.isRecursive = true
+    val store = createAnnotationStore()
+    store += decl.ID -> List(Annotation("type", AnnotationStr("(Int, Str) -> Bool")))
+
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
+    writer.write(decl)
+    printWriter.flush()
+    val expected =
+      """RECURSIVE RecOper(_, _)
+        |(*
+        |  @type: (Int, Str) -> Bool;
+        |*)
+        |RecOper(x, y) == TRUE
+        |
+        |""".stripMargin
+    assert(expected == stringWriter.toString)
+  }
+
+  test("operator declaration without annotation") {
+    val decl = TlaOperDecl("MyOper", List(OperParam("x"), OperParam("y")), tla.bool(true))
+    val store = createAnnotationStore()
+
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
+    writer.write(decl)
+    printWriter.flush()
+    val expected =
+      """MyOper(x, y) == TRUE
+        |
+        |""".stripMargin
+    assert(expected == stringWriter.toString)
+  }
+
+  test("operator declaration with an empty list") {
+    val decl = TlaOperDecl("MyOper", List(OperParam("x"), OperParam("y")), tla.bool(true))
+    val store = createAnnotationStore()
+    store += decl.ID -> List()
+
+    val writer = new PrettyWriterWithAnnotations(store, printWriter, layout80)
+    writer.write(decl)
+    printWriter.flush()
+    val expected =
+      """MyOper(x, y) == TRUE
+        |
+        |""".stripMargin
+    assert(expected == stringWriter.toString)
+  }
 }

tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/passes/ConfigurationPassImpl.scala

@@ -5,7 +5,7 @@ import at.forsyte.apalache.io.tlc.TlcConfigParserApalache
 import at.forsyte.apalache.io.tlc.config._
 import at.forsyte.apalache.tla.lir.UntypedPredefs._
 import at.forsyte.apalache.tla.lir._
-import at.forsyte.apalache.tla.lir.io.{PrettyWriter, TlaWriterFactory}
+import at.forsyte.apalache.tla.lir.io.{PrettyWriter, TlaWriter, TlaWriterFactory}
 import at.forsyte.apalache.tla.lir.oper.{TlaActionOper, TlaBoolOper, TlaOper, TlaTempOper}
 import at.forsyte.apalache.tla.lir.transformations.TransformationTracker
 import at.forsyte.apalache.tla.lir.transformations.impl.IdleTracker
@@ -70,10 +70,8 @@ class ConfigurationPassImpl @Inject() (
 
     // dump the configuration result
     val outdir = options.getOrError("io", "outdir").asInstanceOf[Path]
-    writerFactory.writeModuleToFile(
-        configuredModule,
-        new File(outdir.toFile, "out-config.tla")
-    )
+    writerFactory.writeModuleToFile(configuredModule, TlaWriter.STANDARD_MODULES,
+        new File(outdir.toFile, "out-config.tla"))
 
     outputTlaModule = Some(configuredModule)
     true

tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/passes/DesugarerPassImpl.scala

@@ -2,7 +2,7 @@ package at.forsyte.apalache.tla.pp.passes
 
 import at.forsyte.apalache.infra.passes.{Pass, PassOptions, TlaModuleMixin}
 import at.forsyte.apalache.tla.lir.TlaModule
-import at.forsyte.apalache.tla.lir.io.TlaWriterFactory
+import at.forsyte.apalache.tla.lir.io.{TlaWriter, TlaWriterFactory}
 import at.forsyte.apalache.tla.lir.transformations.TransformationTracker
 import at.forsyte.apalache.tla.lir.transformations.standard._
 import at.forsyte.apalache.tla.pp.{Desugarer, UniqueNameGenerator}
@@ -46,7 +46,7 @@ class DesugarerPassImpl @Inject() (
 
     // dump the result of preprocessing
     val outdir = options.getOrError("io", "outdir").asInstanceOf[Path]
-    writerFactory.writeModuleToFile(output, new File(outdir.toFile, "out-desugarer.tla"))
+    writerFactory.writeModuleToFile(output, TlaWriter.STANDARD_MODULES, new File(outdir.toFile, "out-desugarer.tla"))
     outputTlaModule = Some(output)
 
     true

tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/passes/InlinePassImpl.scala

@@ -2,7 +2,7 @@ package at.forsyte.apalache.tla.pp.passes
 
 import at.forsyte.apalache.infra.passes.{Pass, PassOptions, TlaModuleMixin}
 import at.forsyte.apalache.tla.lir.UntypedPredefs._
-import at.forsyte.apalache.tla.lir.io.TlaWriterFactory
+import at.forsyte.apalache.tla.lir.io.{TlaWriter, TlaWriterFactory}
 import at.forsyte.apalache.tla.lir.storage.BodyMapFactory
 import at.forsyte.apalache.tla.lir.transformations.TransformationTracker
 import at.forsyte.apalache.tla.lir.transformations.standard._
@@ -94,7 +94,7 @@ class InlinePassImpl @Inject() (val options: PassOptions, gen: UniqueNameGenerat
 
     // dump the result of preprocessing
     val outdir = options.getOrError("io", "outdir").asInstanceOf[Path]
-    writerFactory.writeModuleToFile(filtered, new File(outdir.toFile, "out-inline.tla"))
+    writerFactory.writeModuleToFile(filtered, TlaWriter.STANDARD_MODULES, new File(outdir.toFile, "out-inline.tla"))
 
     outputTlaModule = Some(filtered)
     true