From ffc2437ccfa4d63ffca028d727cdc958c27f1e4c Mon Sep 17 00:00:00 2001
From: bugarela <gabrielamoreira05@gmail.com>
Date: Fri, 8 Mar 2024 18:30:05 -0300
Subject: [PATCH] Use type given by quint to build lambda bodies

---
 .../at/forsyte/apalache/io/quint/Quint.scala     | 16 ++++++++--------
 .../forsyte/apalache/io/quint/TestQuintEx.scala  |  6 ++++++
 .../apalache/tla/typecomp/ScopedBuilder.scala    | 15 +++++++++++++++
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/tla-io/src/main/scala/at/forsyte/apalache/io/quint/Quint.scala b/tla-io/src/main/scala/at/forsyte/apalache/io/quint/Quint.scala
index f975d1b847..35f0f9dbd9 100644
--- a/tla-io/src/main/scala/at/forsyte/apalache/io/quint/Quint.scala
+++ b/tla-io/src/main/scala/at/forsyte/apalache/io/quint/Quint.scala
@@ -82,18 +82,18 @@ class Quint(quintOutput: QuintOutput) {
   // operators that take parameters, but these require different constructs
   // in Apalache's IR. Thus, we need to decompose the parts of a QuintLambda
   // for two different purposes.
-  private val lambdaBodyAndParams: QuintLambda => NullaryOpReader[(TBuilderInstruction, Seq[(OperParam, TlaType1)])] = {
+  private val lambdaBodyAndParams: QuintLambda => NullaryOpReader[(TBuilderInstruction, Seq[(OperParam, TlaType1)], TlaType1)] = {
     case ex @ QuintLambda(id, paramNames, _, body) =>
-      val quintParamTypes = types(id).typ match {
-        case QuintOperT(types, _) => types
-        case invalidType          => throw new QuintIRParseError(s"lambda ${ex} has invalid type ${invalidType}")
+      val (quintParamTypes, quintResType) = types(id).typ match {
+        case QuintOperT(types, resType) => (types, resType)
+        case invalidType                => throw new QuintIRParseError(s"lambda ${ex} has invalid type ${invalidType}")
       }
       val operParams = paramNames.zip(quintParamTypes).map(operParam)
       val paramTypes = quintParamTypes.map(typeConv.convert(_))
       val typedParams = operParams.zip(paramTypes)
       for {
         tlaBody <- tlaExpression(body)
-      } yield (tlaBody, typedParams)
+      } yield (tlaBody, typedParams, typeConv.convert(quintResType))
   }
 
   private def typeTagOfId(id: BigInt): TypeTag = {
@@ -668,7 +668,7 @@ class Quint(quintOutput: QuintOutput) {
       (expr match {
         // Parameterized operators are defined in Quint using Lambdas
         case lam: QuintLambda =>
-          lambdaBodyAndParams(lam)
+          lambdaBodyAndParams(lam).map { case (body, params, _) => (body, params) }
         // Otherwise it's an operator with no params
         case other => tlaExpression(other).map(b => (b, List()))
       }).map {
@@ -716,8 +716,8 @@ class Quint(quintOutput: QuintOutput) {
             .map(tlaExpr => tla.letIn(tlaExpr, tlaOpDef))
         }
       case lam: QuintLambda =>
-        lambdaBodyAndParams(lam).map { case (body, typedParams) =>
-          tla.lambda(nameGen.uniqueLambdaName(), body, typedParams: _*)
+        lambdaBodyAndParams(lam).map { case (body, typedParams, resultType) =>
+          tla.lambda(nameGen.uniqueLambdaName(), body, resultType, typedParams: _*)
         }
       case app: QuintApp => tlaApplication(app)
     }
diff --git a/tla-io/src/test/scala/at/forsyte/apalache/io/quint/TestQuintEx.scala b/tla-io/src/test/scala/at/forsyte/apalache/io/quint/TestQuintEx.scala
index 9799732fa0..857ae6a15e 100644
--- a/tla-io/src/test/scala/at/forsyte/apalache/io/quint/TestQuintEx.scala
+++ b/tla-io/src/test/scala/at/forsyte/apalache/io/quint/TestQuintEx.scala
@@ -637,6 +637,12 @@ class TestQuintEx extends AnyFunSuite {
     assert(convert(quintMatch) == expected)
   }
 
+  test("prefers quint types over inferred types") {
+    // Regression on https://github.com/informalsystems/quint/issues/1393
+    val expr = Q.lam(Seq("x" -> QuintBoolT()), Q.nam("x", QuintVarT("t0")), QuintBoolT())
+    assert(translate(expr).typeTag == Typed(OperT1(Seq(BoolT1), BoolT1)))
+  }
+
   test("can convert builtin assert operator") {
     assert(convert(Q.app("assert", Q.nIsGreaterThan0)(QuintBoolT())) == "n > 0")
   }
diff --git a/tlair/src/main/scala/at/forsyte/apalache/tla/typecomp/ScopedBuilder.scala b/tlair/src/main/scala/at/forsyte/apalache/tla/typecomp/ScopedBuilder.scala
index 0a407d5319..cde8bf7af9 100644
--- a/tlair/src/main/scala/at/forsyte/apalache/tla/typecomp/ScopedBuilder.scala
+++ b/tlair/src/main/scala/at/forsyte/apalache/tla/typecomp/ScopedBuilder.scala
@@ -299,6 +299,21 @@ class ScopedBuilder(val strict: Boolean = true)
     } yield ex
   }
 
+  /** Alternative to lambda, which accepts an explicit result type */
+  def lambda(
+      uniqueName: String,
+      body: TBuilderInstruction,
+      resultType: TlaType1,
+      params: TypedParam*): TBuilderInstruction = {
+    params.foreach(validateParamType)
+    for {
+      bodyEx <- body
+      paramTypes = params.map(_._2)
+      operType = OperT1(paramTypes, resultType) // use given type instead of type tag
+      ex <- letIn(name(uniqueName, operType), decl(uniqueName, body, params: _*))
+    } yield ex
+  }
+
   /** {{{LET decl(...) = ... IN body}}} */
   def letIn(body: TBuilderInstruction, decl: TBuilderOperDeclInstruction): TBuilderInstruction = for {
     usedBeforeDecl <- getAllUsed // decl name may not appear in decl body