<commit_before>package at.forsyte.apalache.tla.imp

import java.io.{File, PrintWriter}
import java.nio.file.Files

import at.forsyte.apalache.tla.imp.src.{SourcePosition, SourceRegion, SourceStore}
import at.forsyte.apalache.tla.lir.actions.TlaActionOper
import at.forsyte.apalache.tla.lir.control.{LetInOper, TlaControlOper}
import at.forsyte.apalache.tla.lir.convenience.tla
import at.forsyte.apalache.tla.lir.oper._
import at.forsyte.apalache.tla.lir.predef._
import at.forsyte.apalache.tla.lir.temporal.TlaTempOper
import at.forsyte.apalache.tla.lir.values._
import at.forsyte.apalache.tla.lir.{OperEx, ValEx, _}
import org.junit.runner.RunWith
import org.scalatest.FunSuite
import org.scalatest.junit.JUnitRunner

import scala.io.Source

/**
  * Tests for the SANY importer.
  *
  * @author konnov
  */
@RunWith(classOf[JUnitRunner])
class TestSanyImporter extends FunSuite {

  def expectTlaDecl(sourceStore: SourceStore, name: String, params: List[FormalParam],
                    body: TlaEx): (TlaDecl => Unit) = {
    case d: TlaOperDecl =>
      assert(name == d.name)
      assert(params == d.formalParams)
      assert(body == d.body)
      assert(d.body.ID.valid) // an identifier has been assigned
      assert(sourceStore.contains(d.body.safeId)) // and source file information has been saved

    case d@_ =>
      fail("Expected a TlaOperDecl, found: " + d)
  }

  def findAndExpectTlaDecl(sourceStore: SourceStore,
                           mod: TlaModule,
                           name: String,
                           params: List[FormalParam],
                           body: TlaEx): Unit = {
    mod.declarations.find {
      _.name == name
    } match {
      case Some(d: TlaOperDecl) =>
        expectTlaDecl(sourceStore, name, List(), body)

      case _ =>
        fail("Expected a TlaDecl")
    }
  }

  test("empty module") {
    val text =
      """---- MODULE justASimpleTest ----
        |================================
      """.stripMargin

    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, new SourceStore)
      .loadFromSource("justASimpleTest", Source.fromString(text))
    assert("justASimpleTest" == rootName)
  }

  test("empty module from a file") {
    val text =
      """---- MODULE justASimpleTest ----
        |================================
      """.stripMargin

    val tempDir = Files.createTempDirectory("sanyimp").toFile
    val temp = new File(tempDir, "justASimpleTest.tla")
    val source = Source.fromString(text)
    // write the contents to a temporary file
    val pw = new PrintWriter(temp)
    try {
      source.getLines().foreach(line => pw.println(line))
    } finally {
      pw.close()
    }

    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, new SourceStore)
      .loadFromFile(temp)
    assert("justASimpleTest" == rootName)
  }

  test("one variable") {
    val text =
      """---- MODULE onevar ----
        |VARIABLE x
        |================================
      """.stripMargin

    val (rootName, modules) =
      new SanyImporter(EnvironmentHandlerGenerator.makeEH, new SourceStore)
        .loadFromSource("onevar", Source.fromString(text))
    val mod = expectSingleModule("onevar", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case vd: TlaVarDecl =>
        assert("x" == vd.name)

      case _ =>
        fail("Expected a declaration of variable x")
    }
  }

  test("one constant") {
    val text =
      """---- MODULE oneconst ----
        |CONSTANT n
        |================================
      """.stripMargin

    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, new SourceStore)
      .loadFromSource("oneconst", Source.fromString(text))
    val mod = expectSingleModule("oneconst", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case constDecl: TlaConstDecl =>
        assert("n" == constDecl.name)
    }
  }

  test("constant operator (int)") {
    val text =
      """---- MODULE constop ----
        |MyOp == 1
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaInt(1)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 9)) == loc.region)
    }
  }

  test("constant operator (decimal)") {
    val text =
      """---- MODULE constop ----
        |MyOp == 123.456
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaDecimal(BigDecimal("123.456"))) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 15)) == loc.region)
    }
  }

  test("constant operator (string)") {
    val text =
      """---- MODULE constop ----
        |MyOp == "Hello"
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaStr("Hello")) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 15)) == loc.region)
    }
  }

  test("constant operator (FALSE)") {
    val text =
      """---- MODULE constop ----
        |MyOp == FALSE
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaBool(false)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert("constop" == loc.filename)
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 13)) == loc.region)
    }
  }

  test("constant operator (TRUE)") {
    val text =
      """---- MODULE constop ----
        |MyOp == TRUE
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaBool(true)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
    }
  }

  test("constant operator (a single variable)") {
    val text =
      """---- MODULE constop ----
        |VARIABLE x
        |
        |MyOp == x
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(2 == mod.declarations.size)

    mod.declarations(1) match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(NameEx("x") == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(4, 9), SourcePosition(4, 9)) == loc.region)
    }
  }

  test("constant operator (a single constant)") {
    val text =
      """---- MODULE constop ----
        |CONSTANT n
        |
        |MyOp == n
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(2 == mod.declarations.size)

    mod.declarations(1) match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(NameEx("n") == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
    }
  }

  test("constant operator (a builtin operator)") {
    val text =
      """---- MODULE builtinop ----
        |MyOp == FALSE \/ TRUE
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("builtinop", Source.fromString(text))
    val mod = expectSingleModule("builtinop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(OperEx(TlaBoolOper.or, ValEx(TlaFalse), ValEx(TlaTrue)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 21)) == loc.region)
    }
  }

  test("empty set") {
    val text =
      """---- MODULE emptyset ----
        |MyOp == {}
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("emptyset", Source.fromString(text))
    val mod = expectSingleModule("emptyset", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(OperEx(TlaSetOper.enumSet) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
    }
  }


  test("constant operator (a builtin op and variables)") {
    val text =
      """---- MODULE builtinop ----
        |VARIABLE x
        |MyOp == x /\ TRUE
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("builtinop", Source.fromString(text))
    val mod = expectSingleModule("builtinop", rootName, modules)
    assert(2 == mod.declarations.size)

    mod.declarations(1) match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(OperEx(TlaBoolOper.and, NameEx("x"), ValEx(TlaTrue)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
    }
  }

  test("built-in operators") {
    // Following the operator table from tla2sany.semantic.BuiltInLevel.LevelData.
    // Here we test the very basic application of every operator.
    // For more advanced expressions, see the individual tests for each operator.
    val text =
    """---- MODULE builtins ----
      |VARIABLES x
      |True == TRUE
      |
      |Eq == FALSE = TRUE
      |Ne == FALSE /= TRUE
      |Prime == x'
      |Not == ~x
      |Or == FALSE \/ TRUE
      |And == FALSE /\ TRUE
      |Equiv == FALSE <=> TRUE
      |Implies == FALSE => TRUE
      |Subset == SUBSET x
      |Union == UNION x
      |Domain == DOMAIN x
      |Subseteq == x \subseteq x
      |In == x \in x
      |Notin == x \notin x
      |Setminus == x \ x
      |Cap == x \cap x
      |Cup == x \cup x
      |Times == x \X x
      |LeadsTo == TRUE ~> TRUE
      |Box == []TRUE
      |Diamond == <>TRUE
      |Enabled == ENABLED x
      |Unchanged == UNCHANGED x
      |Cdot == (True \cdot True)
      |Guarantees == True -+-> True
      |Angleact == <<True>>_x
      |BoundedChoose == CHOOSE y \in x: TRUE
      |BoundedExists == \E y \in x: TRUE
      |BoundedForall == \A y \in x: TRUE
      |CartesianProd == x \X x \X x
      |Pair == <<1, 2>>
      |Tuple == <<1, 2, 3>>
      |Case == CASE 1 -> 2 [] 3 -> 4 [] 5 -> 6
      |CaseOther == CASE 1 -> 2 [] 3 -> 4 [] 5 -> 6 [] OTHER -> 7
      |ConjList == /\ FALSE
      |            /\ TRUE
      |            /\ FALSE
      |DisjList == \/ FALSE
      |            \/ TRUE
      |            \/ FALSE
      |Except == [ x EXCEPT ![0] = 1 ]
      |ExceptAt == [ x EXCEPT ![0] = @ /\ TRUE]
      |FcnApply == x[1]
      |FcnCtor == [ y \in x |-> y \cup y ]
      |FcnCtor2 == [ a \in x, b \in BOOLEAN |-> <<a, b>> ]
      |FcnCtor3 == [ <<a, b>> \in x \X BOOLEAN |-> <<a, b>> ]
      |IfThenElse == IF TRUE THEN FALSE ELSE TRUE
      |RcdCtor == [ a |-> 1, b |-> 2 ]
      |RcdSelect == x.foo
      |SetEnumerate == { 1, 2, 3 }
      |SetOfFcns == [ x -> x ]
      |SetOfRcds == [ a: x, b: x ]
      |StrongFairness == SF_x(True)
      |WeakFairness == WF_x(True)
      |SquareAct == [True]_x
      |TemporalExists == \EE y : True
      |TemporalForall == \AA y : True
      |UnboundedChoose == CHOOSE y: TRUE
      |UnboundedExists == \E y: TRUE
      |UnboundedForall == \A y: TRUE
      |SetOfAll == { 1: y \in x }
      |SetOfTuples == { <<a, b>> : a \in x, b \in x }
      |SubsetOf == { y \in x: TRUE }
      |Boolean == BOOLEAN
      |String == STRING
      |
      |================================
      |""".stripMargin

    // TODO: implement the following functions
    //        |nonRecursiveFcnSpec[y \in x] == y \cup x
    //        |recursiveFcnSpec[y \in x] == recursiveFcnSpec[y]

    // TODO: figure out what this operator is
    //
    //        |TemporalWhile == ????

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("builtins", Source.fromString(text))
    val mod = expectSingleModule("builtins", rootName, modules)
    val root = modules(rootName)

    def expectDecl(name: String, body: TlaEx): Unit =
      findAndExpectTlaDecl(locationStore, root, name, List(), body)

    val trueOperDecl = mod.declarations(1)
    expectDecl("True", ValEx(TlaTrue))
    val trueOper = TlaOperDecl("True", List(), ValEx(TlaTrue)).operator
    assert(OperEx(trueOper) == OperEx(trueOper))

    expectDecl("Eq", OperEx(TlaOper.eq, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Ne", OperEx(TlaOper.ne, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Prime", OperEx(TlaActionOper.prime, NameEx("x")))
    expectDecl("Not", OperEx(TlaBoolOper.not, NameEx("x")))
    expectDecl("Or", OperEx(TlaBoolOper.or, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("And", OperEx(TlaBoolOper.and, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Equiv", OperEx(TlaBoolOper.equiv, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Implies", OperEx(TlaBoolOper.implies, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Subset", OperEx(TlaSetOper.SUBSET, NameEx("x")))
    expectDecl("Union", OperEx(TlaSetOper.union, NameEx("x")))
    expectDecl("Domain", OperEx(TlaFunOper.domain, NameEx("x")))
    expectDecl("Subseteq", OperEx(TlaSetOper.subseteq, NameEx("x"), NameEx("x")))
    expectDecl("In", OperEx(TlaSetOper.in, NameEx("x"), NameEx("x")))
    expectDecl("Notin", OperEx(TlaSetOper.notin, NameEx("x"), NameEx("x")))
    expectDecl("Setminus", OperEx(TlaSetOper.setminus, NameEx("x"), NameEx("x")))
    expectDecl("Cap", OperEx(TlaSetOper.cap, NameEx("x"), NameEx("x")))
    expectDecl("Cup", OperEx(TlaSetOper.cup, NameEx("x"), NameEx("x")))
    expectDecl("Times", OperEx(TlaSetOper.times, NameEx("x"), NameEx("x")))
    expectDecl("LeadsTo", OperEx(TlaTempOper.leadsTo, ValEx(TlaTrue), ValEx(TlaTrue)))
    expectDecl("Box", OperEx(TlaTempOper.box, ValEx(TlaTrue)))
    expectDecl("Diamond", OperEx(TlaTempOper.diamond, ValEx(TlaTrue)))
    expectDecl("Enabled", OperEx(TlaActionOper.enabled, NameEx("x")))
    expectDecl("Unchanged", OperEx(TlaActionOper.unchanged, NameEx("x")))
    expectDecl("Cdot", OperEx(TlaActionOper.composition, OperEx(trueOper), OperEx(trueOper)))
    expectDecl("Guarantees",
      OperEx(TlaTempOper.guarantees, OperEx(trueOper), OperEx(trueOper)))
    expectDecl("Angleact",
      OperEx(TlaActionOper.nostutter, OperEx(trueOper), NameEx("x")))
    expectDecl("BoundedChoose",
      OperEx(TlaOper.chooseBounded, NameEx("y"), NameEx("x"), ValEx(TlaTrue)))
    expectDecl("BoundedExists",
      OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("x"), ValEx(TlaTrue)))
    expectDecl("BoundedForall",
      OperEx(TlaBoolOper.forall, NameEx("y"), NameEx("x"), ValEx(TlaTrue)))
    expectDecl("CartesianProd",
      OperEx(TlaSetOper.times, NameEx("x"), NameEx("x"), NameEx("x")))
    expectDecl("Pair", OperEx(TlaFunOper.tuple, ValEx(TlaInt(1)), ValEx(TlaInt(2))))
    expectDecl("Tuple",
      OperEx(TlaFunOper.tuple, ValEx(TlaInt(1)), ValEx(TlaInt(2)), ValEx(TlaInt(3))))
    expectDecl("Case",
      OperEx(TlaControlOper.caseNoOther, 1.to(6).map(i => ValEx(TlaInt(i))): _*))
    expectDecl("CaseOther",
      OperEx(TlaControlOper.caseWithOther, (7 +: 1.to(6)).map(i => ValEx(TlaInt(i))): _*))
    expectDecl("ConjList",
      OperEx(TlaBoolOper.and, List(TlaFalse, TlaTrue, TlaFalse).map(b => ValEx(b)): _*))
    expectDecl("DisjList",
      OperEx(TlaBoolOper.or, List(TlaFalse, TlaTrue, TlaFalse).map(b => ValEx(b)): _*))
    expectDecl("Except",
      OperEx(TlaFunOper.except,
        NameEx("x"), TlaFunOper.mkTuple(ValEx(TlaInt(0))), ValEx(TlaInt(1))
      ))
    expectDecl("ExceptAt",
      OperEx(TlaFunOper.except,
        NameEx("x"),
        TlaFunOper.mkTuple(ValEx(TlaInt(0))),
        OperEx(TlaBoolOper.and,
          OperEx(TlaFunOper.app, NameEx("x"), ValEx(TlaInt(0))),
          ValEx(TlaTrue))
      ))
    expectDecl("FcnApply", OperEx(TlaFunOper.app, NameEx("x"), ValEx(TlaInt(1))))
    val cup = OperEx(TlaSetOper.cup, NameEx("y"), NameEx("y"))
    expectDecl("FcnCtor",
      OperEx(TlaFunOper.funDef, cup, NameEx("y"), NameEx("x")))
    expectDecl("FcnCtor2",
      OperEx(TlaFunOper.funDef, OperEx(TlaFunOper.tuple, NameEx("a"), NameEx("b")),
        NameEx("a"), NameEx("x"), NameEx("b"), NameEx("x")))
    expectDecl("FcnCtor3",
      OperEx(TlaFunOper.funDef, OperEx(TlaFunOper.tuple, NameEx("a"), NameEx("b")),
        OperEx(TlaFunOper.tuple, NameEx("a"), NameEx("b")), NameEx("x")))
    expectDecl("IfThenElse",
      OperEx(TlaControlOper.ifThenElse, ValEx(TlaTrue), ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("RcdCtor",
      OperEx(TlaFunOper.enum,
        ValEx(TlaStr("a")), ValEx(TlaInt(1)), ValEx(TlaStr("b")), ValEx(TlaInt(2))))
    expectDecl("RcdSelect",
      OperEx(TlaFunOper.app,
        NameEx("x"), ValEx(TlaStr("foo"))))
    expectDecl("SetEnumerate",
      OperEx(TlaSetOper.enumSet, ValEx(TlaInt(1)), ValEx(TlaInt(2)), ValEx(TlaInt(3))))
    expectDecl("SetOfFcns", OperEx(TlaSetOper.funSet, NameEx("x"), NameEx("x")))
    expectDecl("SetOfRcds",
      OperEx(TlaSetOper.recSet,
        ValEx(TlaStr("a")), NameEx("x"), ValEx(TlaStr("b")), NameEx("x")))
    expectDecl("StrongFairness",
      OperEx(TlaTempOper.strongFairness, NameEx("x"), OperEx(trueOper)))
    expectDecl("WeakFairness",
      OperEx(TlaTempOper.weakFairness, NameEx("x"), OperEx(trueOper)))
    expectDecl("SquareAct",
      OperEx(TlaActionOper.stutter, OperEx(trueOper), NameEx("x")))
    expectDecl("TemporalExists",
      OperEx(TlaTempOper.EE, NameEx("y"), OperEx(trueOper)))
    expectDecl("TemporalForall",
      OperEx(TlaTempOper.AA, NameEx("y"), OperEx(trueOper)))
    expectDecl("UnboundedChoose",
      OperEx(TlaOper.chooseUnbounded, NameEx("y"), ValEx(TlaTrue)))
    expectDecl("UnboundedExists",
      OperEx(TlaBoolOper.existsUnbounded, NameEx("y"), ValEx(TlaTrue)))
    expectDecl("UnboundedForall",
      OperEx(TlaBoolOper.forallUnbounded, NameEx("y"), ValEx(TlaTrue)))
    expectDecl("SetOfAll",
      OperEx(TlaSetOper.map, ValEx(TlaInt(1)), NameEx("y"), NameEx("x")))
    expectDecl("SetOfTuples",
      OperEx(TlaSetOper.map, OperEx(TlaFunOper.tuple), NameEx("a"), NameEx("x"), NameEx("b"), NameEx("x")))
    expectDecl("SubsetOf",
      OperEx(TlaSetOper.filter, NameEx("y"), NameEx("x"), ValEx(TlaTrue)))
    expectDecl("Boolean", ValEx(TlaBoolSet))
    expectDecl("String", ValEx(TlaStrSet))
  }

  test("comprehensions with tuples") {
    // TLA+ allows the user to use some form of pattern matching with tuples
    val text =
    """---- MODULE comprehensions ----
      |VARIABLE XY, X, Y
      |FilterTuples == { <<x, y>> \in XY :  x = 1}
      |MapTuples == { x = 1 : <<x, y>> \in XY}
      |MapTuples2 == { x = 1 : x, y \in XY}
      |================================
      |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("comprehensions", Source.fromString(text))
    val mod = expectSingleModule("comprehensions", rootName, modules)
    val root = modules(rootName)

    def expectDecl(name: String, body: TlaEx): Unit =
      findAndExpectTlaDecl(locationStore, root, name, List(), body)

    expectDecl("FilterTuples",
      OperEx(TlaSetOper.filter,
        OperEx(TlaFunOper.tuple, NameEx("x"), NameEx("y")),
        NameEx("XY"),
        OperEx(TlaOper.eq, NameEx("x"), ValEx(TlaInt(1)))
      ))////
    expectDecl("MapTuples",
      OperEx(TlaSetOper.map,
        OperEx(TlaOper.eq, NameEx("x"), ValEx(TlaInt(1))),
        OperEx(TlaFunOper.tuple, NameEx("x"), NameEx("y")),
        NameEx("XY")
      ))////
    expectDecl("MapTuples2",
      OperEx(TlaSetOper.map,
        OperEx(TlaOper.eq, NameEx("x"), ValEx(TlaInt(1))),
        OperEx(TlaFunOper.tuple, NameEx("x"), NameEx("y")),
        NameEx("XY")
      ))////
  }

  test("funCtor quantifiers") {
    // One can write tricky combinations of bound variables in TLA+.
    // We translate all of them uniformly:
    //   every quantified expression has exactly one bound variable or a tuple of variables.
    val text =
    """---- MODULE composite ----
      |VARIABLE X, Z
      |Q1 == \E x \in X: \E y \in X: TRUE
      |Q2 == \E x, y \in X: TRUE
      |Q3 == \E x, y \in X, z \in Z: TRUE
      |Q4 == \E x, y \in X, <<a, b>> \in Z, z \in Z: TRUE
      |Q5 == \E x, y: TRUE
      |================================
      |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("composite", Source.fromString(text))
    val mod = expectSingleModule("composite", rootName, modules)

    def expectDecl(name: String, body: TlaEx): (TlaDecl => Unit) =
      expectTlaDecl(locationStore, name, List(), body)

    expectDecl("Q1",
      OperEx(TlaBoolOper.exists, NameEx("x"), NameEx("X"),
        OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("X"), ValEx(TlaTrue))))(mod.declarations(2))
    expectDecl("Q2",
      OperEx(TlaBoolOper.exists, NameEx("x"), NameEx("X"),
        OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("X"), ValEx(TlaTrue))))(mod.declarations(3))
    expectDecl("Q3",
      OperEx(TlaBoolOper.exists, NameEx("x"), NameEx("X"),
        OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("X"),
          OperEx(TlaBoolOper.exists, NameEx("z"), NameEx("Z"),
            ValEx(TlaTrue)))))(mod.declarations(4))
    expectDecl("Q4",
      OperEx(TlaBoolOper.exists, NameEx("x"), NameEx("X"),
        OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("X"),
          OperEx(TlaBoolOper.exists, OperEx(TlaFunOper.tuple, NameEx("a"), NameEx("b")), NameEx("Z"),
            OperEx(TlaBoolOper.exists, NameEx("z"), NameEx("Z"),
              ValEx(TlaTrue))))))(mod.declarations(5))
    expectDecl("Q5",
      OperEx(TlaBoolOper.existsUnbounded, NameEx("x"),
        OperEx(TlaBoolOper.existsUnbounded, NameEx("y"), ValEx(TlaTrue))))(mod.declarations(6))
  }

  test("function updates") {
    val text =
      """---- MODULE except ----
        |VARIABLE x
        |Except == [ x EXCEPT ![0] = 1 ]
        |ExceptAt == [ x EXCEPT ![0] = @ /\ TRUE]
        |ExceptTuple == [ x EXCEPT ![<<0>>] = 1 ]
        |ExceptManyAt == [ x EXCEPT ![1][2] = @ /\ TRUE]
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("except", Source.fromString(text))
    val mod = expectSingleModule("except", rootName, modules)

    def expectDecl(name: String, params: List[FormalParam], body: TlaEx) = expectTlaDecl(locationStore, name, params, body)

    expectDecl("Except",
      List(),
      OperEx(TlaFunOper.except,
        NameEx("x"), TlaFunOper.mkTuple(ValEx(TlaInt(0))), ValEx(TlaInt(1))
      ))(mod.declarations(1))

    expectDecl("ExceptAt",
      List(),
      OperEx(TlaFunOper.except,
        NameEx("x"),
        TlaFunOper.mkTuple(ValEx(TlaInt(0))),
        OperEx(TlaBoolOper.and,
          OperEx(TlaFunOper.app, NameEx("x"), ValEx(TlaInt(0))),
          ValEx(TlaTrue))
      ))(mod.declarations(2))

    expectDecl("ExceptTuple",
      List(),
      OperEx(TlaFunOper.except,
        NameEx("x"), TlaFunOper.mkTuple(TlaFunOper.mkTuple(ValEx(TlaInt(0)))), ValEx(TlaInt(1))
      ))(mod.declarations(3))

    // the importer automatically substitutes @ with function application and unfolds ![1][2] to a chain of EXCEPTS
    expectDecl("ExceptManyAt",
      List(),
      OperEx(TlaFunOper.except,
        NameEx("x"),
        TlaFunOper.mkTuple(ValEx(TlaInt(1))),
        OperEx(TlaFunOper.except,
          OperEx(TlaFunOper.app,
            NameEx("x"),
            ValEx(TlaInt(1))),
          TlaFunOper.mkTuple(ValEx(TlaInt(2))),
          OperEx(TlaBoolOper.and,
            OperEx(TlaFunOper.app,
              OperEx(TlaFunOper.app,
                NameEx("x"), ValEx(TlaInt(1))),
              ValEx(TlaInt(2))
            ), ///
            ValEx(TlaTrue))
        )))(mod.declarations(4))
  }

  test("expression labels") {
    val text =
      """---- MODULE labels ----
        |A == {FALSE} \cup (l1 :: {TRUE})
        |B == \E x \in BOOLEAN: l2(x) :: x /= FALSE
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("labels", Source.fromString(text))
    val mod = expectSingleModule("labels", rootName, modules)

    def expectDecl(n: String, p: List[FormalParam], b: TlaEx) = expectTlaDecl(locationStore, n, p, b)

    //    A == {FALSE} \cup (l1 :: {TRUE})
    val expectedABody =
      OperEx(TlaSetOper.cup,
        OperEx(TlaSetOper.enumSet, ValEx(TlaBool(false))),
        OperEx(TlaOper.label,
          OperEx(TlaSetOper.enumSet, ValEx(TlaBool(true))),
          ValEx(TlaStr("l1")))) ////
    expectDecl("A", List(), expectedABody)(mod.declarations.head)

    //    B == \E x \in BOOLEAN: l2(x) :: x /= FALSE
    // since we cannot access formal parameters, we ignore them
    val expectedBBody =
    OperEx(TlaBoolOper.exists,
      NameEx("x"),
      ValEx(TlaBoolSet),
      OperEx(TlaOper.label,
        OperEx(TlaOper.ne, NameEx("x"), ValEx(TlaBool(false))),
        ValEx(TlaStr("l2")))) ////
    expectDecl("B", List(), expectedBBody)(mod.declarations.tail.head)
  }

  test("tuples as arguments") {
    // test that single function arguments are translated consistently
    val text =
      """---- MODULE args ----
        |VARIABLE f, g, e, h
        |
        |A == f[2]
        |B == g["green"]
        |C == e[FALSE]
        |D == h[{}]
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("args", Source.fromString(text))
    val mod = expectSingleModule("args", rootName, modules)

    def expectDecl(n: String, p: List[FormalParam], b: TlaEx) = expectTlaDecl(locationStore, n, p, b)

    val expectedABody = OperEx(TlaFunOper.app, NameEx("f"), ValEx(TlaInt(2)))
    expectDecl("A", List(), expectedABody)(mod.declarations(4))

    val expectedBBody = OperEx(TlaFunOper.app, NameEx("g"), ValEx(TlaStr("green")))
    expectDecl("B", List(), expectedBBody)(mod.declarations(5))

    val expectedCBody = OperEx(TlaFunOper.app, NameEx("e"), ValEx(TlaBool(false)))
    expectDecl("C", List(), expectedCBody)(mod.declarations(6))

    val expectedDBody = OperEx(TlaFunOper.app, NameEx("h"), OperEx(TlaSetOper.enumSet))
    expectDecl("D", List(), expectedDBody)(mod.declarations(7))
  }

  test("complex updates") {
    // One can write tricky combinations of EXCEPT arguments.
    val text =
      """---- MODULE updates ----
        |VARIABLE f
        |E1 == [ f EXCEPT ![0] = 1, ![2] = 3 ]
        |E2 == [ f EXCEPT ![0][1][2] = 3 ]
        |E3 == [ f EXCEPT ![0,1,2] = 3 ]
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("updates", Source.fromString(text))
    val mod = expectSingleModule("updates", rootName, modules)

    def expectDecl(name: String, body: TlaEx): (TlaDecl => Unit) =
      expectTlaDecl(locationStore, name, List(), body)

    expectDecl("E1",
      OperEx(TlaFunOper.except,
        NameEx("f"),
        TlaFunOper.mkTuple(ValEx(TlaInt(0))), ValEx(TlaInt(1)),
        TlaFunOper.mkTuple(ValEx(TlaInt(2))), ValEx(TlaInt(3))
      ))(mod.declarations(1))
    expectDecl("E2",
      OperEx(TlaFunOper.except,
        NameEx("f"),
        TlaFunOper.mkTuple(ValEx(TlaInt(0))),
        OperEx(TlaFunOper.except,
          OperEx(TlaFunOper.app, NameEx("f"), ValEx(TlaInt(0))),
          TlaFunOper.mkTuple(ValEx(TlaInt(1))),
          OperEx(TlaFunOper.except,
            OperEx(TlaFunOper.app, OperEx(TlaFunOper.app, NameEx("f"), ValEx(TlaInt(0))), ValEx(TlaInt(1))),
            TlaFunOper.mkTuple(ValEx(TlaInt(2))),
            ValEx(TlaInt(3)))
        )//
      ))(mod.declarations(2))
    expectDecl("E3",
      OperEx(TlaFunOper.except,
        NameEx("f"),
        TlaFunOper.mkTuple(TlaFunOper.mkTuple(ValEx(TlaInt(0)), ValEx(TlaInt(1)), ValEx(TlaInt(2)))),
        ValEx(TlaInt(3))
      ))(mod.declarations(3))
  }

  test("complex record selects") {
    // Testing that multiple field accesses work as expected
    val text =
      """---- MODULE selects ----
        |VARIABLE f
        |S1 == f.a.b
        |S2 == f.a.b.c
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("selects", Source.fromString(text))
    val mod = expectSingleModule("selects", rootName, modules)

    def expectDecl(name: String, body: TlaEx): (TlaDecl => Unit) =
      expectTlaDecl(locationStore, name, List(), body)

    expectDecl("S1",
      OperEx(TlaFunOper.app,
        OperEx(TlaFunOper.app,
          NameEx("f"),
          ValEx(TlaStr("a"))),
        ValEx(TlaStr("b")))
    )(mod.declarations(1))
    expectDecl("S2",
      OperEx(TlaFunOper.app,
        OperEx(TlaFunOper.app,
          OperEx(TlaFunOper.app,
            NameEx("f"),
            ValEx(TlaStr("a"))),
          ValEx(TlaStr("b"))),
        ValEx(TlaStr("c")))
    )(mod.declarations(2))
  }

  test("complex function ctors") {
    // One can write tricky combinations of bound variables in function constructors.
    val text =
      """---- MODULE funCtor ----
        |VARIABLE X, Z
        |C1 == [ x \in X, y \in X |-> TRUE ]
        |C2 == [ x, y \in X |-> TRUE ]
        |C3 == [ x, y \in X, z \in Z |-> TRUE]
        |C4 == [ x, y \in X, <<a, b>> \in Z, z \in Z |-> TRUE ]
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("funCtor", Source.fromString(text))
    val mod = expectSingleModule("funCtor", rootName, modules)

    def expectDecl(name: String, body: TlaEx) = expectTlaDecl(locationStore, name, List(), body)

    expectDecl("C1",
      OperEx(TlaFunOper.funDef,
        ValEx(TlaTrue),
        NameEx("x"), NameEx("X"),
        NameEx("y"), NameEx("X")))(mod.declarations(2))
    expectDecl("C2",
      OperEx(TlaFunOper.funDef,
        ValEx(TlaTrue),
        NameEx("x"), NameEx("X"),
        NameEx("y"), NameEx("X")))(mod.declarations(3))
    expectDecl("C3",
      OperEx(TlaFunOper.funDef,
        ValEx(TlaTrue),
        NameEx("x"), NameEx("X"),
        NameEx("y"), NameEx("X"),
        NameEx("z"), NameEx("Z")
      ))(mod.declarations(4))
    expectDecl("C4",
      OperEx(TlaFunOper.funDef,
        ValEx(TlaTrue),
        NameEx("x"), NameEx("X"),
        NameEx("y"), NameEx("X"),
        TlaFunOper.mkTuple(NameEx("a"), NameEx("b")), NameEx("Z"),
        NameEx("z"), NameEx("Z")
      ))(mod.declarations(5))
  }

  test("level-1 operators") {
    // operators with basic parameters, that is, no operator is passed as a parameter
    val text =
      """---- MODULE level1Operators ----
        |VARIABLE x, y
        |A(i, j) == i \cup j
        |i ** j == i \cap j
        |C == A(1, 2)
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("level1Operators", Source.fromString(text))
    val mod = expectSingleModule("level1Operators", rootName, modules)

    def expectDecl(n: String, p: List[FormalParam], b: TlaEx) = expectTlaDecl(locationStore, n, p, b)

    expectDecl("A", List(SimpleFormalParam("i"), SimpleFormalParam("j")),
      OperEx(TlaSetOper.cup, NameEx("i"), NameEx("j")))(mod.declarations(2))
    expectDecl("**", List(SimpleFormalParam("i"), SimpleFormalParam("j")),
      OperEx(TlaSetOper.cap, NameEx("i"), NameEx("j")))(mod.declarations(3))
    val aDecl = mod.declarations(2).asInstanceOf[TlaOperDecl]
    expectDecl("C", List(),
      OperEx(aDecl.operator, ValEx(TlaInt(1)), ValEx(TlaInt(2))))(mod.declarations(4))
  }
<commit_msg>fixed the test after updating Desugarer<commit_after>package at.forsyte.apalache.tla.imp

import java.io.{File, PrintWriter}
import java.nio.file.Files

import at.forsyte.apalache.tla.imp.src.{SourcePosition, SourceRegion, SourceStore}
import at.forsyte.apalache.tla.lir.actions.TlaActionOper
import at.forsyte.apalache.tla.lir.control.{LetInOper, TlaControlOper}
import at.forsyte.apalache.tla.lir.convenience.tla
import at.forsyte.apalache.tla.lir.oper._
import at.forsyte.apalache.tla.lir.predef._
import at.forsyte.apalache.tla.lir.temporal.TlaTempOper
import at.forsyte.apalache.tla.lir.values._
import at.forsyte.apalache.tla.lir.{OperEx, ValEx, _}
import org.junit.runner.RunWith
import org.scalatest.FunSuite
import org.scalatest.junit.JUnitRunner

import scala.io.Source

/**
  * Tests for the SANY importer.
  *
  * @author konnov
  */
@RunWith(classOf[JUnitRunner])
class TestSanyImporter extends FunSuite {

  def expectTlaDecl(sourceStore: SourceStore, name: String, params: List[FormalParam],
                    body: TlaEx): (TlaDecl => Unit) = {
    case d: TlaOperDecl =>
      assert(name == d.name)
      assert(params == d.formalParams)
      assert(body == d.body)
      assert(d.body.ID.valid) // an identifier has been assigned
      assert(sourceStore.contains(d.body.safeId)) // and source file information has been saved

    case d@_ =>
      fail("Expected a TlaOperDecl, found: " + d)
  }

  def findAndExpectTlaDecl(sourceStore: SourceStore,
                           mod: TlaModule,
                           name: String,
                           params: List[FormalParam],
                           body: TlaEx): Unit = {
    mod.declarations.find {
      _.name == name
    } match {
      case Some(d: TlaOperDecl) =>
        expectTlaDecl(sourceStore, name, List(), body)

      case _ =>
        fail("Expected a TlaDecl")
    }
  }

  test("empty module") {
    val text =
      """---- MODULE justASimpleTest ----
        |================================
      """.stripMargin

    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, new SourceStore)
      .loadFromSource("justASimpleTest", Source.fromString(text))
    assert("justASimpleTest" == rootName)
  }

  test("empty module from a file") {
    val text =
      """---- MODULE justASimpleTest ----
        |================================
      """.stripMargin

    val tempDir = Files.createTempDirectory("sanyimp").toFile
    val temp = new File(tempDir, "justASimpleTest.tla")
    val source = Source.fromString(text)
    // write the contents to a temporary file
    val pw = new PrintWriter(temp)
    try {
      source.getLines().foreach(line => pw.println(line))
    } finally {
      pw.close()
    }

    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, new SourceStore)
      .loadFromFile(temp)
    assert("justASimpleTest" == rootName)
  }

  test("one variable") {
    val text =
      """---- MODULE onevar ----
        |VARIABLE x
        |================================
      """.stripMargin

    val (rootName, modules) =
      new SanyImporter(EnvironmentHandlerGenerator.makeEH, new SourceStore)
        .loadFromSource("onevar", Source.fromString(text))
    val mod = expectSingleModule("onevar", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case vd: TlaVarDecl =>
        assert("x" == vd.name)

      case _ =>
        fail("Expected a declaration of variable x")
    }
  }

  test("one constant") {
    val text =
      """---- MODULE oneconst ----
        |CONSTANT n
        |================================
      """.stripMargin

    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, new SourceStore)
      .loadFromSource("oneconst", Source.fromString(text))
    val mod = expectSingleModule("oneconst", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case constDecl: TlaConstDecl =>
        assert("n" == constDecl.name)
    }
  }

  test("constant operator (int)") {
    val text =
      """---- MODULE constop ----
        |MyOp == 1
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaInt(1)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 9)) == loc.region)
    }
  }

  test("constant operator (decimal)") {
    val text =
      """---- MODULE constop ----
        |MyOp == 123.456
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaDecimal(BigDecimal("123.456"))) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 15)) == loc.region)
    }
  }

  test("constant operator (string)") {
    val text =
      """---- MODULE constop ----
        |MyOp == "Hello"
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaStr("Hello")) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 15)) == loc.region)
    }
  }

  test("constant operator (FALSE)") {
    val text =
      """---- MODULE constop ----
        |MyOp == FALSE
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaBool(false)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert("constop" == loc.filename)
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 13)) == loc.region)
    }
  }

  test("constant operator (TRUE)") {
    val text =
      """---- MODULE constop ----
        |MyOp == TRUE
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(ValEx(TlaBool(true)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
    }
  }

  test("constant operator (a single variable)") {
    val text =
      """---- MODULE constop ----
        |VARIABLE x
        |
        |MyOp == x
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(2 == mod.declarations.size)

    mod.declarations(1) match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(NameEx("x") == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(4, 9), SourcePosition(4, 9)) == loc.region)
    }
  }

  test("constant operator (a single constant)") {
    val text =
      """---- MODULE constop ----
        |CONSTANT n
        |
        |MyOp == n
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("constop", Source.fromString(text))
    val mod = expectSingleModule("constop", rootName, modules)
    assert(2 == mod.declarations.size)

    mod.declarations(1) match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(NameEx("n") == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
    }
  }

  test("constant operator (a builtin operator)") {
    val text =
      """---- MODULE builtinop ----
        |MyOp == FALSE \/ TRUE
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("builtinop", Source.fromString(text))
    val mod = expectSingleModule("builtinop", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(OperEx(TlaBoolOper.or, ValEx(TlaFalse), ValEx(TlaTrue)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
      val loc = locationStore.find(actionDecl.body.safeId).get
        assert(SourceRegion(SourcePosition(2, 9), SourcePosition(2, 21)) == loc.region)
    }
  }

  test("empty set") {
    val text =
      """---- MODULE emptyset ----
        |MyOp == {}
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("emptyset", Source.fromString(text))
    val mod = expectSingleModule("emptyset", rootName, modules)
    assert(1 == mod.declarations.size)

    mod.declarations.head match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(OperEx(TlaSetOper.enumSet) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
    }
  }


  test("constant operator (a builtin op and variables)") {
    val text =
      """---- MODULE builtinop ----
        |VARIABLE x
        |MyOp == x /\ TRUE
        |================================
      """.stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("builtinop", Source.fromString(text))
    val mod = expectSingleModule("builtinop", rootName, modules)
    assert(2 == mod.declarations.size)

    mod.declarations(1) match {
      case actionDecl: TlaOperDecl =>
        assert("MyOp" == actionDecl.name)
        assert(0 == actionDecl.formalParams.length)
        assert(OperEx(TlaBoolOper.and, NameEx("x"), ValEx(TlaTrue)) == actionDecl.body)
        assert(actionDecl.body.ID.valid) // an identifier has been assigned
        assert(locationStore.contains(actionDecl.body.safeId)) // and source file information has been saved
    }
  }

  test("built-in operators") {
    // Following the operator table from tla2sany.semantic.BuiltInLevel.LevelData.
    // Here we test the very basic application of every operator.
    // For more advanced expressions, see the individual tests for each operator.
    val text =
    """---- MODULE builtins ----
      |VARIABLES x
      |True == TRUE
      |
      |Eq == FALSE = TRUE
      |Ne == FALSE /= TRUE
      |Prime == x'
      |Not == ~x
      |Or == FALSE \/ TRUE
      |And == FALSE /\ TRUE
      |Equiv == FALSE <=> TRUE
      |Implies == FALSE => TRUE
      |Subset == SUBSET x
      |Union == UNION x
      |Domain == DOMAIN x
      |Subseteq == x \subseteq x
      |In == x \in x
      |Notin == x \notin x
      |Setminus == x \ x
      |Cap == x \cap x
      |Cup == x \cup x
      |Times == x \X x
      |LeadsTo == TRUE ~> TRUE
      |Box == []TRUE
      |Diamond == <>TRUE
      |Enabled == ENABLED x
      |Unchanged == UNCHANGED x
      |Cdot == (True \cdot True)
      |Guarantees == True -+-> True
      |Angleact == <<True>>_x
      |BoundedChoose == CHOOSE y \in x: TRUE
      |BoundedExists == \E y \in x: TRUE
      |BoundedForall == \A y \in x: TRUE
      |CartesianProd == x \X x \X x
      |Pair == <<1, 2>>
      |Tuple == <<1, 2, 3>>
      |Case == CASE 1 -> 2 [] 3 -> 4 [] 5 -> 6
      |CaseOther == CASE 1 -> 2 [] 3 -> 4 [] 5 -> 6 [] OTHER -> 7
      |ConjList == /\ FALSE
      |            /\ TRUE
      |            /\ FALSE
      |DisjList == \/ FALSE
      |            \/ TRUE
      |            \/ FALSE
      |Except == [ x EXCEPT ![0] = 1 ]
      |ExceptAt == [ x EXCEPT ![0] = @ /\ TRUE]
      |FcnApply == x[1]
      |FcnCtor == [ y \in x |-> y \cup y ]
      |FcnCtor2 == [ a \in x, b \in BOOLEAN |-> <<a, b>> ]
      |FcnCtor3 == [ <<a, b>> \in x \X BOOLEAN |-> <<a, b>> ]
      |IfThenElse == IF TRUE THEN FALSE ELSE TRUE
      |RcdCtor == [ a |-> 1, b |-> 2 ]
      |RcdSelect == x.foo
      |SetEnumerate == { 1, 2, 3 }
      |SetOfFcns == [ x -> x ]
      |SetOfRcds == [ a: x, b: x ]
      |StrongFairness == SF_x(True)
      |WeakFairness == WF_x(True)
      |SquareAct == [True]_x
      |TemporalExists == \EE y : True
      |TemporalForall == \AA y : True
      |UnboundedChoose == CHOOSE y: TRUE
      |UnboundedExists == \E y: TRUE
      |UnboundedForall == \A y: TRUE
      |SetOfAll == { 1: y \in x }
      |SetOfTuples == { <<a, b>> : a \in x, b \in x }
      |SubsetOf == { y \in x: TRUE }
      |Boolean == BOOLEAN
      |String == STRING
      |
      |================================
      |""".stripMargin

    // TODO: implement the following functions
    //        |nonRecursiveFcnSpec[y \in x] == y \cup x
    //        |recursiveFcnSpec[y \in x] == recursiveFcnSpec[y]

    // TODO: figure out what this operator is
    //
    //        |TemporalWhile == ????

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("builtins", Source.fromString(text))
    val mod = expectSingleModule("builtins", rootName, modules)
    val root = modules(rootName)

    def expectDecl(name: String, body: TlaEx): Unit =
      findAndExpectTlaDecl(locationStore, root, name, List(), body)

    val trueOperDecl = mod.declarations(1)
    expectDecl("True", ValEx(TlaTrue))
    val trueOper = TlaOperDecl("True", List(), ValEx(TlaTrue)).operator
    assert(OperEx(trueOper) == OperEx(trueOper))

    expectDecl("Eq", OperEx(TlaOper.eq, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Ne", OperEx(TlaOper.ne, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Prime", OperEx(TlaActionOper.prime, NameEx("x")))
    expectDecl("Not", OperEx(TlaBoolOper.not, NameEx("x")))
    expectDecl("Or", OperEx(TlaBoolOper.or, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("And", OperEx(TlaBoolOper.and, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Equiv", OperEx(TlaBoolOper.equiv, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Implies", OperEx(TlaBoolOper.implies, ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("Subset", OperEx(TlaSetOper.SUBSET, NameEx("x")))
    expectDecl("Union", OperEx(TlaSetOper.union, NameEx("x")))
    expectDecl("Domain", OperEx(TlaFunOper.domain, NameEx("x")))
    expectDecl("Subseteq", OperEx(TlaSetOper.subseteq, NameEx("x"), NameEx("x")))
    expectDecl("In", OperEx(TlaSetOper.in, NameEx("x"), NameEx("x")))
    expectDecl("Notin", OperEx(TlaSetOper.notin, NameEx("x"), NameEx("x")))
    expectDecl("Setminus", OperEx(TlaSetOper.setminus, NameEx("x"), NameEx("x")))
    expectDecl("Cap", OperEx(TlaSetOper.cap, NameEx("x"), NameEx("x")))
    expectDecl("Cup", OperEx(TlaSetOper.cup, NameEx("x"), NameEx("x")))
    expectDecl("Times", OperEx(TlaSetOper.times, NameEx("x"), NameEx("x")))
    expectDecl("LeadsTo", OperEx(TlaTempOper.leadsTo, ValEx(TlaTrue), ValEx(TlaTrue)))
    expectDecl("Box", OperEx(TlaTempOper.box, ValEx(TlaTrue)))
    expectDecl("Diamond", OperEx(TlaTempOper.diamond, ValEx(TlaTrue)))
    expectDecl("Enabled", OperEx(TlaActionOper.enabled, NameEx("x")))
    expectDecl("Unchanged", OperEx(TlaActionOper.unchanged, NameEx("x")))
    expectDecl("Cdot", OperEx(TlaActionOper.composition, OperEx(trueOper), OperEx(trueOper)))
    expectDecl("Guarantees",
      OperEx(TlaTempOper.guarantees, OperEx(trueOper), OperEx(trueOper)))
    expectDecl("Angleact",
      OperEx(TlaActionOper.nostutter, OperEx(trueOper), NameEx("x")))
    expectDecl("BoundedChoose",
      OperEx(TlaOper.chooseBounded, NameEx("y"), NameEx("x"), ValEx(TlaTrue)))
    expectDecl("BoundedExists",
      OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("x"), ValEx(TlaTrue)))
    expectDecl("BoundedForall",
      OperEx(TlaBoolOper.forall, NameEx("y"), NameEx("x"), ValEx(TlaTrue)))
    expectDecl("CartesianProd",
      OperEx(TlaSetOper.times, NameEx("x"), NameEx("x"), NameEx("x")))
    expectDecl("Pair", OperEx(TlaFunOper.tuple, ValEx(TlaInt(1)), ValEx(TlaInt(2))))
    expectDecl("Tuple",
      OperEx(TlaFunOper.tuple, ValEx(TlaInt(1)), ValEx(TlaInt(2)), ValEx(TlaInt(3))))
    expectDecl("Case",
      OperEx(TlaControlOper.caseNoOther, 1.to(6).map(i => ValEx(TlaInt(i))): _*))
    expectDecl("CaseOther",
      OperEx(TlaControlOper.caseWithOther, (7 +: 1.to(6)).map(i => ValEx(TlaInt(i))): _*))
    expectDecl("ConjList",
      OperEx(TlaBoolOper.and, List(TlaFalse, TlaTrue, TlaFalse).map(b => ValEx(b)): _*))
    expectDecl("DisjList",
      OperEx(TlaBoolOper.or, List(TlaFalse, TlaTrue, TlaFalse).map(b => ValEx(b)): _*))
    expectDecl("Except",
      OperEx(TlaFunOper.except,
        NameEx("x"), TlaFunOper.mkTuple(ValEx(TlaInt(0))), ValEx(TlaInt(1))
      ))
    expectDecl("ExceptAt",
      OperEx(TlaFunOper.except,
        NameEx("x"),
        TlaFunOper.mkTuple(ValEx(TlaInt(0))),
        OperEx(TlaBoolOper.and,
          OperEx(TlaFunOper.app, NameEx("x"), ValEx(TlaInt(0))),
          ValEx(TlaTrue))
      ))
    expectDecl("FcnApply", OperEx(TlaFunOper.app, NameEx("x"), ValEx(TlaInt(1))))
    val cup = OperEx(TlaSetOper.cup, NameEx("y"), NameEx("y"))
    expectDecl("FcnCtor",
      OperEx(TlaFunOper.funDef, cup, NameEx("y"), NameEx("x")))
    expectDecl("FcnCtor2",
      OperEx(TlaFunOper.funDef, OperEx(TlaFunOper.tuple, NameEx("a"), NameEx("b")),
        NameEx("a"), NameEx("x"), NameEx("b"), NameEx("x")))
    expectDecl("FcnCtor3",
      OperEx(TlaFunOper.funDef, OperEx(TlaFunOper.tuple, NameEx("a"), NameEx("b")),
        OperEx(TlaFunOper.tuple, NameEx("a"), NameEx("b")), NameEx("x")))
    expectDecl("IfThenElse",
      OperEx(TlaControlOper.ifThenElse, ValEx(TlaTrue), ValEx(TlaFalse), ValEx(TlaTrue)))
    expectDecl("RcdCtor",
      OperEx(TlaFunOper.enum,
        ValEx(TlaStr("a")), ValEx(TlaInt(1)), ValEx(TlaStr("b")), ValEx(TlaInt(2))))
    expectDecl("RcdSelect",
      OperEx(TlaFunOper.app,
        NameEx("x"), ValEx(TlaStr("foo"))))
    expectDecl("SetEnumerate",
      OperEx(TlaSetOper.enumSet, ValEx(TlaInt(1)), ValEx(TlaInt(2)), ValEx(TlaInt(3))))
    expectDecl("SetOfFcns", OperEx(TlaSetOper.funSet, NameEx("x"), NameEx("x")))
    expectDecl("SetOfRcds",
      OperEx(TlaSetOper.recSet,
        ValEx(TlaStr("a")), NameEx("x"), ValEx(TlaStr("b")), NameEx("x")))
    expectDecl("StrongFairness",
      OperEx(TlaTempOper.strongFairness, NameEx("x"), OperEx(trueOper)))
    expectDecl("WeakFairness",
      OperEx(TlaTempOper.weakFairness, NameEx("x"), OperEx(trueOper)))
    expectDecl("SquareAct",
      OperEx(TlaActionOper.stutter, OperEx(trueOper), NameEx("x")))
    expectDecl("TemporalExists",
      OperEx(TlaTempOper.EE, NameEx("y"), OperEx(trueOper)))
    expectDecl("TemporalForall",
      OperEx(TlaTempOper.AA, NameEx("y"), OperEx(trueOper)))
    expectDecl("UnboundedChoose",
      OperEx(TlaOper.chooseUnbounded, NameEx("y"), ValEx(TlaTrue)))
    expectDecl("UnboundedExists",
      OperEx(TlaBoolOper.existsUnbounded, NameEx("y"), ValEx(TlaTrue)))
    expectDecl("UnboundedForall",
      OperEx(TlaBoolOper.forallUnbounded, NameEx("y"), ValEx(TlaTrue)))
    expectDecl("SetOfAll",
      OperEx(TlaSetOper.map, ValEx(TlaInt(1)), NameEx("y"), NameEx("x")))
    expectDecl("SetOfTuples",
      OperEx(TlaSetOper.map, OperEx(TlaFunOper.tuple), NameEx("a"), NameEx("x"), NameEx("b"), NameEx("x")))
    expectDecl("SubsetOf",
      OperEx(TlaSetOper.filter, NameEx("y"), NameEx("x"), ValEx(TlaTrue)))
    expectDecl("Boolean", ValEx(TlaBoolSet))
    expectDecl("String", ValEx(TlaStrSet))
  }

  test("comprehensions with tuples") {
    // TLA+ allows the user to use some form of pattern matching with tuples
    val text =
    """---- MODULE comprehensions ----
      |VARIABLE XY, X, Y
      |FilterTuples == { <<x, y>> \in XY :  x = 1}
      |MapTuples == { x = 1 : <<x, y>> \in XY}
      |MapTuples2 == { x = 1 : x, y \in XY}
      |================================
      |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("comprehensions", Source.fromString(text))
    val mod = expectSingleModule("comprehensions", rootName, modules)
    val root = modules(rootName)

    def expectDecl(name: String, body: TlaEx): Unit =
      findAndExpectTlaDecl(locationStore, root, name, List(), body)

    expectDecl("FilterTuples",
      OperEx(TlaSetOper.filter,
        OperEx(TlaFunOper.tuple, NameEx("x"), NameEx("y")),
        NameEx("XY"),
        OperEx(TlaOper.eq, NameEx("x"), ValEx(TlaInt(1)))
      ))////
    expectDecl("MapTuples",
      OperEx(TlaSetOper.map,
        OperEx(TlaOper.eq, NameEx("x"), ValEx(TlaInt(1))),
        OperEx(TlaFunOper.tuple, NameEx("x"), NameEx("y")),
        NameEx("XY")
      ))////
    expectDecl("MapTuples2",
      OperEx(TlaSetOper.map,
        OperEx(TlaOper.eq, NameEx("x"), ValEx(TlaInt(1))),
        OperEx(TlaFunOper.tuple, NameEx("x"), NameEx("y")),
        NameEx("XY")
      ))////
  }

  test("funCtor quantifiers") {
    // One can write tricky combinations of bound variables in TLA+.
    // We translate all of them uniformly:
    //   every quantified expression has exactly one bound variable or a tuple of variables.
    val text =
    """---- MODULE composite ----
      |VARIABLE X, Z
      |Q1 == \E x \in X: \E y \in X: TRUE
      |Q2 == \E x, y \in X: TRUE
      |Q3 == \E x, y \in X, z \in Z: TRUE
      |Q4 == \E x, y \in X, <<a, b>> \in Z, z \in Z: TRUE
      |Q5 == \E x, y: TRUE
      |================================
      |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("composite", Source.fromString(text))
    val mod = expectSingleModule("composite", rootName, modules)

    def expectDecl(name: String, body: TlaEx): (TlaDecl => Unit) =
      expectTlaDecl(locationStore, name, List(), body)

    expectDecl("Q1",
      OperEx(TlaBoolOper.exists, NameEx("x"), NameEx("X"),
        OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("X"), ValEx(TlaTrue))))(mod.declarations(2))
    expectDecl("Q2",
      OperEx(TlaBoolOper.exists, NameEx("x"), NameEx("X"),
        OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("X"), ValEx(TlaTrue))))(mod.declarations(3))
    expectDecl("Q3",
      OperEx(TlaBoolOper.exists, NameEx("x"), NameEx("X"),
        OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("X"),
          OperEx(TlaBoolOper.exists, NameEx("z"), NameEx("Z"),
            ValEx(TlaTrue)))))(mod.declarations(4))
    expectDecl("Q4",
      OperEx(TlaBoolOper.exists, NameEx("x"), NameEx("X"),
        OperEx(TlaBoolOper.exists, NameEx("y"), NameEx("X"),
          OperEx(TlaBoolOper.exists, OperEx(TlaFunOper.tuple, NameEx("a"), NameEx("b")), NameEx("Z"),
            OperEx(TlaBoolOper.exists, NameEx("z"), NameEx("Z"),
              ValEx(TlaTrue))))))(mod.declarations(5))
    expectDecl("Q5",
      OperEx(TlaBoolOper.existsUnbounded, NameEx("x"),
        OperEx(TlaBoolOper.existsUnbounded, NameEx("y"), ValEx(TlaTrue))))(mod.declarations(6))
  }

  test("function updates") {
    val text =
      """---- MODULE except ----
        |VARIABLE x
        |Except == [ x EXCEPT ![0] = 1 ]
        |ExceptAt == [ x EXCEPT ![0] = @ /\ TRUE]
        |ExceptTuple == [ x EXCEPT ![<<0>>] = 1 ]
        |ExceptManyAt == [ x EXCEPT ![1][2] = @ /\ TRUE]
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("except", Source.fromString(text))
    val mod = expectSingleModule("except", rootName, modules)

    def expectDecl(name: String, params: List[FormalParam], body: TlaEx) = expectTlaDecl(locationStore, name, params, body)

    expectDecl("Except",
      List(),
      OperEx(TlaFunOper.except,
        NameEx("x"), TlaFunOper.mkTuple(ValEx(TlaInt(0))), ValEx(TlaInt(1))
      ))(mod.declarations(1))

    expectDecl("ExceptAt",
      List(),
      OperEx(TlaFunOper.except,
        NameEx("x"),
        TlaFunOper.mkTuple(ValEx(TlaInt(0))),
        OperEx(TlaBoolOper.and,
          OperEx(TlaFunOper.app, NameEx("x"), ValEx(TlaInt(0))),
          ValEx(TlaTrue))
      ))(mod.declarations(2))

    expectDecl("ExceptTuple",
      List(),
      OperEx(TlaFunOper.except,
        NameEx("x"), TlaFunOper.mkTuple(TlaFunOper.mkTuple(ValEx(TlaInt(0)))), ValEx(TlaInt(1))
      ))(mod.declarations(3))

    // the importer automatically substitutes @ with function application and unfolds ![1][2] to a chain of EXCEPTS
    expectDecl("ExceptManyAt",
      List(),
      OperEx(TlaFunOper.except,
        NameEx("x"),
        TlaFunOper.mkTuple(ValEx(TlaInt(1))),
        OperEx(TlaFunOper.except,
          OperEx(TlaFunOper.app,
            NameEx("x"),
            ValEx(TlaInt(1))),
          TlaFunOper.mkTuple(ValEx(TlaInt(2))),
          OperEx(TlaBoolOper.and,
            OperEx(TlaFunOper.app,
              OperEx(TlaFunOper.app,
                NameEx("x"), ValEx(TlaInt(1))),
              ValEx(TlaInt(2))
            ), ///
            ValEx(TlaTrue))
        )))(mod.declarations(4))
  }

  test("expression labels") {
    val text =
      """---- MODULE labels ----
        |A == {FALSE} \cup (l1 :: {TRUE})
        |B == \E x \in BOOLEAN: l2(x) :: x /= FALSE
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("labels", Source.fromString(text))
    val mod = expectSingleModule("labels", rootName, modules)

    def expectDecl(n: String, p: List[FormalParam], b: TlaEx) = expectTlaDecl(locationStore, n, p, b)

    //    A == {FALSE} \cup (l1 :: {TRUE})
    val expectedABody =
      OperEx(TlaSetOper.cup,
        OperEx(TlaSetOper.enumSet, ValEx(TlaBool(false))),
        OperEx(TlaOper.label,
          OperEx(TlaSetOper.enumSet, ValEx(TlaBool(true))),
          ValEx(TlaStr("l1")))) ////
    expectDecl("A", List(), expectedABody)(mod.declarations.head)

    //    B == \E x \in BOOLEAN: l2(x) :: x /= FALSE
    // since we cannot access formal parameters, we ignore them
    val expectedBBody =
    OperEx(TlaBoolOper.exists,
      NameEx("x"),
      ValEx(TlaBoolSet),
      OperEx(TlaOper.label,
        OperEx(TlaOper.ne, NameEx("x"), ValEx(TlaBool(false))),
        ValEx(TlaStr("l2")))) ////
    expectDecl("B", List(), expectedBBody)(mod.declarations.tail.head)
  }

  test("tuples as arguments") {
    // test that single function arguments are translated consistently
    val text =
      """---- MODULE args ----
        |VARIABLE f, g, e, h
        |
        |A == f[2]
        |B == g["green"]
        |C == e[FALSE]
        |D == h[{}]
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("args", Source.fromString(text))
    val mod = expectSingleModule("args", rootName, modules)

    def expectDecl(n: String, p: List[FormalParam], b: TlaEx) = expectTlaDecl(locationStore, n, p, b)

    val expectedABody = OperEx(TlaFunOper.app, NameEx("f"), ValEx(TlaInt(2)))
    expectDecl("A", List(), expectedABody)(mod.declarations(4))

    val expectedBBody = OperEx(TlaFunOper.app, NameEx("g"), ValEx(TlaStr("green")))
    expectDecl("B", List(), expectedBBody)(mod.declarations(5))

    val expectedCBody = OperEx(TlaFunOper.app, NameEx("e"), ValEx(TlaBool(false)))
    expectDecl("C", List(), expectedCBody)(mod.declarations(6))

    val expectedDBody = OperEx(TlaFunOper.app, NameEx("h"), OperEx(TlaSetOper.enumSet))
    expectDecl("D", List(), expectedDBody)(mod.declarations(7))
  }

  test("complex updates") {
    // One can write tricky combinations of EXCEPT arguments.
    val text =
      """---- MODULE updates ----
        |VARIABLE f
        |E1 == [ f EXCEPT ![0] = 1, ![2] = 3 ]
        |E2 == [ f EXCEPT ![0][1][2] = 3 ]
        |E3 == [ f EXCEPT ![0,1,2] = 3 ]
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("updates", Source.fromString(text))
    val mod = expectSingleModule("updates", rootName, modules)

    def expectDecl(name: String, body: TlaEx): (TlaDecl => Unit) =
      expectTlaDecl(locationStore, name, List(), body)

    expectDecl("E1",
      OperEx(TlaFunOper.except,
        NameEx("f"),
        TlaFunOper.mkTuple(ValEx(TlaInt(0))), ValEx(TlaInt(1)),
        TlaFunOper.mkTuple(ValEx(TlaInt(2))), ValEx(TlaInt(3))
      ))(mod.declarations(1))
    expectDecl("E2",
      OperEx(TlaFunOper.except,
        NameEx("f"),
        TlaFunOper.mkTuple(ValEx(TlaInt(0))),
        OperEx(TlaFunOper.except,
          OperEx(TlaFunOper.app, NameEx("f"), ValEx(TlaInt(0))),
          TlaFunOper.mkTuple(ValEx(TlaInt(1))),
          OperEx(TlaFunOper.except,
            OperEx(TlaFunOper.app, OperEx(TlaFunOper.app, NameEx("f"), ValEx(TlaInt(0))), ValEx(TlaInt(1))),
            TlaFunOper.mkTuple(ValEx(TlaInt(2))),
            ValEx(TlaInt(3)))
        )//
      ))(mod.declarations(2))
    expectDecl("E3",
      OperEx(TlaFunOper.except,
        NameEx("f"),
        TlaFunOper.mkTuple(TlaFunOper.mkTuple(ValEx(TlaInt(0)), ValEx(TlaInt(1)), ValEx(TlaInt(2)))),
        ValEx(TlaInt(3))
      ))(mod.declarations(3))
  }

  test("complex record selects") {
    // Testing that multiple field accesses work as expected
    val text =
      """---- MODULE selects ----
        |VARIABLE f
        |S1 == f.a.b
        |S2 == f.a.b.c
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("selects", Source.fromString(text))
    val mod = expectSingleModule("selects", rootName, modules)

    def expectDecl(name: String, body: TlaEx): (TlaDecl => Unit) =
      expectTlaDecl(locationStore, name, List(), body)

    expectDecl("S1",
      OperEx(TlaFunOper.app,
        OperEx(TlaFunOper.app,
          NameEx("f"),
          ValEx(TlaStr("a"))),
        ValEx(TlaStr("b")))
    )(mod.declarations(1))
    expectDecl("S2",
      OperEx(TlaFunOper.app,
        OperEx(TlaFunOper.app,
          OperEx(TlaFunOper.app,
            NameEx("f"),
            ValEx(TlaStr("a"))),
          ValEx(TlaStr("b"))),
        ValEx(TlaStr("c")))
    )(mod.declarations(2))
  }

  test("complex function ctors") {
    // One can write tricky combinations of bound variables in function constructors.
    val text =
      """---- MODULE funCtor ----
        |VARIABLE X, Z
        |C1 == [ x \in X, y \in X |-> TRUE ]
        |C2 == [ x, y \in X |-> TRUE ]
        |C3 == [ x, y \in X, z \in Z |-> TRUE]
        |C4 == [ x, y \in X, <<a, b>> \in Z, z \in Z |-> TRUE ]
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("funCtor", Source.fromString(text))
    val mod = expectSingleModule("funCtor", rootName, modules)

    def expectDecl(name: String, body: TlaEx) = expectTlaDecl(locationStore, name, List(), body)

    expectDecl("C1",
      OperEx(TlaFunOper.funDef,
        ValEx(TlaTrue),
        NameEx("x"), NameEx("X"),
        NameEx("y"), NameEx("X")))(mod.declarations(2))
    expectDecl("C2",
      OperEx(TlaFunOper.funDef,
        ValEx(TlaTrue),
        NameEx("x"), NameEx("X"),
        NameEx("y"), NameEx("X")))(mod.declarations(3))
    expectDecl("C3",
      OperEx(TlaFunOper.funDef,
        ValEx(TlaTrue),
        NameEx("x"), NameEx("X"),
        NameEx("y"), NameEx("X"),
        NameEx("z"), NameEx("Z")
      ))(mod.declarations(4))
    expectDecl("C4",
      OperEx(TlaFunOper.funDef,
        ValEx(TlaTrue),
        NameEx("x"), NameEx("X"),
        NameEx("y"), NameEx("X"),
        NameEx("a_b"), NameEx("Z"), // the tuple <<a, b>> is collapsed to a_b by Desugarer
        NameEx("z"), NameEx("Z")
      ))(mod.declarations(5))
  }

  test("level-1 operators") {
    // operators with basic parameters, that is, no operator is passed as a parameter
    val text =
      """---- MODULE level1Operators ----
        |VARIABLE x, y
        |A(i, j) == i \cup j
        |i ** j == i \cap j
        |C == A(1, 2)
        |================================
        |""".stripMargin

    val locationStore = new SourceStore
    val (rootName, modules) = new SanyImporter(EnvironmentHandlerGenerator.makeEH, locationStore)
      .loadFromSource("level1Operators", Source.fromString(text))
    val mod = expectSingleModule("level1Operators", rootName, modules)

    def expectDecl(n: String, p: List[FormalParam], b: TlaEx) = expectTlaDecl(locationStore, n, p, b)

    expectDecl("A", List(SimpleFormalParam("i"), SimpleFormalParam("j")),
      OperEx(TlaSetOper.cup, NameEx("i"), NameEx("j")))(mod.declarations(2))
    expectDecl("**", List(SimpleFormalParam("i"), SimpleFormalParam("j")),
      OperEx(TlaSetOper.cap, NameEx("i"), NameEx("j")))(mod.declarations(3))
    val aDecl = mod.declarations(2).asInstanceOf[TlaOperDecl]
    expectDecl("C", List(),
      OperEx(aDecl.operator, ValEx(TlaInt(1)), ValEx(TlaInt(2))))(mod.declarations(4))
  }

{'fraction_non_alphanumeric': 0.17210654910500878, 'fraction_numerical': 0.006559728050140755, 'mean_word_length': 3.099787685774947, 'pattern_counts': {'":': 0, '<': 57, '<?xml version=': 0, '>': 149, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '19592', 'n_tokens_mistral': 25180, 'n_tokens_neox': 24535, 'n_words': 4860}