Merge pull request #3207 from apalache-mc/igor/fix3204

Fix transitive inlining of polytypes when receiving Quint as input

Author: konnov

.unreleased/bug-fixes/transitive-inlining.md

@@ -0,0 +1 @@
+ - Fix transitive inlining of polymorphic Quint definitions (#3207)

tla-pp/src/main/scala/at/forsyte/apalache/tla/pp/Inliner.scala

@@ -9,7 +9,7 @@ import at.forsyte.apalache.tla.lir.transformations.standard.{
 }
 import at.forsyte.apalache.tla.lir.transformations.{TlaExTransformation, TransformationTracker}
 import at.forsyte.apalache.tla.pp.Inliner.{DeclFilter, FilterFun}
-import at.forsyte.apalache.tla.types.{Substitution, TypeUnifier, TypeVarPool}
+import at.forsyte.apalache.tla.types.{EqClass, Substitution, TypeUnifier, TypeVarPool}
 
 import scala.collection.immutable.SortedMap
 
@@ -96,31 +96,45 @@ class Inliner(
   private def nonNullaryFilter(d: TlaOperDecl): Boolean =
     !keepNullaryMono || isPolyTag(d.typeTag) || d.formalParams.nonEmpty
 
-  // Given a declaration (possibly holding a polymorphic type) and a monotyped target, computes
-  // a substitution of the two. A substitution is assumed to exist, otherwise TypingException is thrown.
-  private def getSubstitution(targetType: TlaType1, decl: TlaOperDecl): (Substitution, TlaType1) = {
-    val genericType = decl.typeTag.asTlaType1()
-    val maxUsedVar = Math.max(genericType.usedNames.foldLeft(0)(Math.max), targetType.usedNames.foldLeft(0)(Math.max))
-    new TypeUnifier(new TypeVarPool(maxUsedVar + 1)).unify(Substitution.empty, genericType, targetType) match {
+  // Given a declaration `callee` (possibly holding a polymorphic type) and the type at a call site,
+  // which can also be a polytype, computes a substitution of the two. A substitution is assumed to exist,
+  // otherwise TypingException is thrown.
+  private def getSubstitution(callSiteType: TlaType1, callee: TlaOperDecl): (Substitution, TlaType1) = {
+    val calleeType = callee.typeTag.asTlaType1()
+    // To fix transitive inlining of polymorphic operators, rename all type variables in `callee`
+    // to fresh ones, so these type variables have indices larger than those in `callSiteType`.
+    // The unification algorithm prefers to keep type variables with smaller indices,
+    // so this ensures that type variables from `callSiteType` are kept.
+    // See the issue #3204.
+    val maxUsedVar = ((calleeType.usedNames ++ callSiteType.usedNames) ++ Set(0)).max
+    val typeVarPool = new TypeVarPool(maxUsedVar + 1)
+    val calleeSub = Substitution(calleeType.usedNames.map(v => EqClass(v) -> typeVarPool.fresh).toMap)
+    val calleeTypeRenamed = calleeSub.subRec(calleeType)
+
+    new TypeUnifier(typeVarPool).unify(Substitution.empty, callSiteType, calleeTypeRenamed) match {
       case None =>
-        throw new TypingException(
-            s"Inliner: Unable to unify generic signature $genericType of ${decl.name} with the concrete type $targetType",
-            decl.ID)
-
-      case Some(pair) => pair
+        throw new TypingException(s"Inliner: Unable to unify the signature $calleeType of ${callee.name} "
+              + "with the type $callSiteType at call site", callee.ID)
+
+      case Some((unifierSub, unifiedType)) =>
+        // Now, we have to add the renamed type variables back to the substitution.
+        // To this end, we compose `calleeSub` with `unifierSub`.
+        // Not that we are not merging them. Otherwise, the type variables of `callee` might take over.
+        val composed = Substitution(calleeSub.mapping ++ unifierSub.mapping)
+        (composed, composed.subRec(unifiedType))
     }
   }
 
-  // Assume an operator declaration named name is in scope.
+  // Assume an operator declaration named `name` is in scope.
   // Creates a fresh copy of the operator body and replaces formal parameter instances with the argument instances.
   private def instantiateWithArgs(scope: Scope)(nameEx: NameEx, args: Seq[TlaEx]): TlaEx = {
     val name = nameEx.name
-    val decl = scope(name)
+    val callee = scope(name)
 
-    val freshBody = deepCopy(decl.body)
+    val freshBody = deepCopy(callee.body)
 
     // All formal parameters get instantiated at once, to avoid parameter-name issues, see #1903
-    val paramMap = decl.formalParams
+    val paramMap = callee.formalParams
       .zip(args)
       .map({ case (OperParam(name, _), arg) =>
         name -> arg
@@ -140,28 +154,30 @@ class Inliner(
     // To cover both cases at once, we run an additional transform on the replaced body
     val newBody = transform(scope)(replacedBody)
 
-    // Note: it can happen that the new body and the decl have desynced types (poly vs mono).
-    // We fix that below with type unification.
-    // If the operator has a parametric signature, we have to substitute type parameters with concrete parameters
-    // 1. Unify the operator type with the arguments.
-    // 2. Apply the resulting substitution to the types in all subexpressions.
-    val actualType = nameEx.typeTag.asTlaType1()
+    // Note: it can happen that the type at the call site and the type of `callee` have different types,
+    // e.g., the `callee` has a more general polytype, or they are both polytypes.
+    // We fix that below with type unification:
+    //  1. Unify the operator type with the arguments.
+    //  2. Apply the resulting substitution to the types in all subexpressions.
+    //  3. Importantly, we prefer the type variables of the call site, as they are the type variables of
+    //     the caller context.
+    val callSiteType = nameEx.typeTag.asTlaType1()
 
-    val (substitution, _) = getSubstitution(actualType, decl)
+    val (substitution, _) = getSubstitution(callSiteType, callee)
 
     if (substitution.isEmpty) newBody
     else new TypeSubstitutor(tracker, substitution)(newBody)
   }
 
   // Assume name is in scope. Creates a local LET-IN for pass-by-name operators.
   private def embedPassByName(scope: Scope)(nameEx: NameEx): TlaEx = {
-    val decl = scope(nameEx.name)
+    val callee = scope(nameEx.name)
 
     // like in instantiateWithArgs, we compare the declaration type to the expected monotype
-    val freshBody = deepCopy(decl.body)
+    val freshBody = deepCopy(callee.body)
     val monoOperType = nameEx.typeTag.asTlaType1()
 
-    val (substitution, tp) = getSubstitution(monoOperType, decl)
+    val (substitution, tp) = getSubstitution(monoOperType, callee)
 
     val tpTag = Typed(tp)
 
@@ -170,8 +186,8 @@ class Inliner(
       else new TypeSubstitutor(tracker, substitution)(freshBody)
 
     // To make a local definition, we use a fresh name, derived from the original name, but renamed to get a fresh $N
-    val newName = renaming.apply(NameEx(decl.name)(decl.typeTag)).asInstanceOf[NameEx].name
-    val newLocalDecl = TlaOperDecl(newName, decl.formalParams, newBody)(tpTag)
+    val newName = renaming.apply(NameEx(callee.name)(callee.typeTag)).asInstanceOf[NameEx].name
+    val newLocalDecl = TlaOperDecl(newName, callee.formalParams, newBody)(tpTag)
 
     LetInEx(NameEx(newName)(tpTag), newLocalDecl)(tpTag)
   }

tla-pp/src/test/scala/at/forsyte/apalache/tla/pp/TestInliner.scala

@@ -3,6 +3,7 @@ package at.forsyte.apalache.tla.pp
 import at.forsyte.apalache.tla.lir.TypedPredefs._
 import at.forsyte.apalache.tla.lir._
 import at.forsyte.apalache.tla.lir.convenience.tla
+import at.forsyte.apalache.tla.types.{tla => ttla}
 import at.forsyte.apalache.tla.lir.oper.ApalacheOper
 import at.forsyte.apalache.tla.lir.transformations.impl.IdleTracker
 import at.forsyte.apalache.tla.lir.transformations.standard.IncrementalRenaming
@@ -11,11 +12,13 @@ import org.junit.runner.RunWith
 import org.scalatest.BeforeAndAfterEach
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatestplus.junit.JUnitRunner
+import org.scalatestplus.scalacheck.ScalaCheckPropertyChecks
+import org.scalatest.prop.TableFor2
 
 import scala.collection.immutable.SortedMap
 
 @RunWith(classOf[JUnitRunner])
-class TestInliner extends AnyFunSuite with BeforeAndAfterEach {
+class TestInliner extends AnyFunSuite with BeforeAndAfterEach with ScalaCheckPropertyChecks {
 
   import Inliner.Scope
 
@@ -351,4 +354,95 @@ class TestInliner extends AnyFunSuite with BeforeAndAfterEach {
     assert(expected == actual)
   }
 
+  val typeVariableOrders: TableFor2[Int, Int] = Table(("a", "b"), (0, 1), (1, 0), (0, 2), (2, 0), (1, 2), (2, 1))
+
+  test("Unify generic definitions under various type variable namings") {
+    // Regression test for the issue #3204:
+    // Inliner failed to unify generics definitions under certain type variable namings.
+    forAll(typeVariableOrders) { (indexOfA: Int, indexOfB: Int) =>
+      // P(x) == { x }; Q(y) == P(y); x' := Q({}) ~~> x' := { {} }
+      val a = VarT1(indexOfA)
+      val b = VarT1(indexOfB)
+      // \* @type: a => Set(a);
+      // P(x) == { x }
+      val PBody = ttla.enumSet(ttla.name("x", a))
+      val PDecl = ttla.decl("P", PBody, (OperParam("x"), a))
+      // \* @type: b => b;
+      // Q(y) == P(y)
+      val PTypeInQ = OperT1(Seq(b), SetT1(b))
+      val QBody = ttla.appOp(ttla.name("P", PTypeInQ), ttla.name("y", b))
+      val QDecl = ttla.decl("Q", QBody, (OperParam("y"), b))
+
+      // @type: () => Bool;
+      // X() == x' := Q({})
+      val intSetSet = SetT1(SetT1(IntT1))
+      val bodyOfX =
+        ttla.assign(ttla.prime(ttla.name("x", intSetSet)),
+            ttla
+              .appOp(
+                  ttla.name("Q", OperT1(Seq(SetT1(IntT1)), SetT1(SetT1(IntT1)))),
+                  ttla.emptySet(IntT1),
+              ))
+      val declOfX = ttla.decl("X", bodyOfX)
+
+      val decls = List(PDecl, QDecl, declOfX)
+      val inputModule = mkModule(decls: _*)
+      val outputModule = inlinerKeepNullary.transformModule(inputModule)
+      val actualBodyOfX = outputModule.declarations(2).asInstanceOf[TlaOperDecl].body
+
+      // extract the type of {{}}
+      val typeOfEmptySet = actualBodyOfX.asInstanceOf[OperEx].args(1).typeTag
+      assert(typeOfEmptySet == Typed(SetT1(SetT1(IntT1))))
+    }
+  }
+
+  test("Unify generic definitions under various type variable namings in LET-IN") {
+    // Regression test for the issue #3204:
+    // Inliner failed to unify generics definitions under certain type variable namings.
+    // This is the LET-IN version of the previous test.
+    forAll(typeVariableOrders) { (indexOfA: Int, indexOfB: Int) =>
+      // X ==
+      //   LET P(x) == { x }
+      //     Q(y) == P(y)
+      //   IN
+      //   x' := Q({}) ~~> x' := { {} }
+      val a = VarT1(indexOfA)
+      val b = VarT1(indexOfB)
+      // \* @type: a => Set(a);
+      // P(x) == { x }
+      val PBody = ttla.enumSet(ttla.name("x", a))
+      val PDecl = ttla.decl("P", PBody, (OperParam("x"), a))
+      // \* @type: b => b;
+      // Q(y) == P(y)
+      val PTypeInQ = OperT1(Seq(b), SetT1(b))
+      val QBody = ttla.appOp(ttla.name("P", PTypeInQ), ttla.name("y", b))
+      val QDecl = ttla.decl("Q", QBody, (OperParam("y"), b))
+
+      // @type: () => Bool;
+      // X() ==
+      //   LET P(x) == { x }
+      //       Q(y) == P(y)
+      //   IN
+      //   x' := Q({})
+      val intSetSet = SetT1(SetT1(IntT1))
+      val bodyOfX =
+        ttla.assign(ttla.prime(ttla.name("x", intSetSet)),
+            ttla
+              .appOp(
+                  ttla.name("Q", OperT1(Seq(SetT1(IntT1)), SetT1(SetT1(IntT1)))),
+                  ttla.emptySet(IntT1),
+              ))
+      val xUnderLetIn = ttla.letIn(ttla.letIn(bodyOfX, QDecl), PDecl)
+      val declOfX = ttla.decl("X", xUnderLetIn)
+
+      val decls = List(declOfX)
+      val inputModule = mkModule(decls: _*)
+      val outputModule = inlinerKeepNullary.transformModule(inputModule)
+      val actualBodyOfX = outputModule.declarations.head.asInstanceOf[TlaOperDecl].body
+
+      // extract the type of {{}}
+      val typeOfEmptySet = actualBodyOfX.asInstanceOf[OperEx].args(1).typeTag
+      assert(typeOfEmptySet == Typed(SetT1(SetT1(IntT1))))
+    }
+  }
 }

tlair/src/main/scala/at/forsyte/apalache/tla/types/TypeUnifier.scala

@@ -35,6 +35,10 @@ class TypeUnifier(varPool: TypeVarPool) {
    * is the solution set showing how to unify lhs and rhs and t is the type resulting from successfully unifying lhs and
    * rhs using mgu. Note that apart from variable substitution, our unification also involves merging record types. When
    * there is no unifier, it returns None.
+   *
+   * <b>WARNING:</b> When two type variables are unified into a single equivalence class, the variable with the
+   * <i>smaller</i> index becomes the representative of the class. Thus, unification may substitute `b` with `a`, but
+   * never `a` with `b`. This property MUST be preserved, as the calling code relies on it.
    */
   def unify(substitution: Substitution, lhs: TlaType1, rhs: TlaType1): Option[(Substitution, TlaType1)] = {
     // Copy the equivalence classes and the mapping from the substitution.