Fix handling of transient contract keys in kvutils (#3110)

* Add failing test case for transient contract keys

* Add kvutils tests for transient contracts and keys

* Fix handling of transient contracts and keys in kvutils

ledger/ledger-api-test-tool/src/main/scala/com/daml/ledger/api/testtool/tests/ContractKeys.scala

@@ -179,6 +179,7 @@ final class ContractKeys(session: LedgerSession) extends LedgerTestSuite(session
       for {
         ledger <- context.participant()
         owner <- ledger.allocateParty()
+
         delegated1TxTree <- ledger
           .submitAndWaitRequest(owner, Delegated(owner, key).create.command)
           .flatMap(ledger.submitAndWaitForTransactionTree)
@@ -199,11 +200,39 @@ final class ContractKeys(session: LedgerSession) extends LedgerTestSuite(session
       }
     }
 
+  val transientContractsArchiveKeys =
+    LedgerTest("CKTransients", "Contract keys created by transient contracts are properly archived") {
+      context =>
+        val key = s"${UUID.randomUUID.toString}-key"
+        val key2 = s"${UUID.randomUUID.toString}-key"
+
+        for {
+          ledger <- context.participant()
+          owner <- ledger.allocateParty()
+          delegation <- ledger.create(owner, Delegation(owner, owner))
+          delegated <- ledger.create(owner, Delegated(owner, key))
+
+          failedFetch <- ledger
+            .exercise(owner, delegation.exerciseFetchByKeyDelegated(_, owner, key2, None))
+            .failed
+
+          // Create a transient contract with a key that is created and archived in same transaction.
+          _ <- ledger.exercise(owner, delegated.exerciseCreateAnotherAndArchive(_, key2))
+
+          // Try it again, expecting it to succeed.
+          _ <- ledger.exercise(owner, delegated.exerciseCreateAnotherAndArchive(_, key2))
+
+        } yield {
+          assertGrpcError(failedFetch, Status.Code.INVALID_ARGUMENT, "couldn't find key")
+        }
+    }
+
   override val tests: Vector[LedgerTest] = Vector(
     fetchDivulgedContract,
     rejectFetchingUndisclosedContract,
     processContractKeys,
-    recreateContractKeys
+    recreateContractKeys,
+    transientContractsArchiveKeys
   )
 
 }

ledger/participant-state/kvutils/src/main/scala/com/daml/ledger/participant/state/kvutils/InputsAndEffects.scala

@@ -37,10 +37,10 @@ private[kvutils] object InputsAndEffects {
         */
       createdContracts: List[(DamlStateKey, NodeCreate[ContractId, VersionedValue[ContractId]])],
       /** The contract keys created or updated as part of the transaction. */
-      updatedContractKeys: List[(DamlStateKey, DamlContractKeyState)]
+      updatedContractKeys: Map[DamlStateKey, DamlContractKeyState]
   )
   object Effects {
-    val empty = Effects(List.empty, List.empty, List.empty)
+    val empty = Effects(List.empty, List.empty, Map.empty)
   }
 
   /** Compute the inputs to a DAML transaction, that is, the referenced contracts, keys
@@ -83,6 +83,8 @@ private[kvutils] object InputsAndEffects {
 
   /** Compute the effects of a DAML transaction, that is, the created and consumed contracts. */
   def computeEffects(entryId: DamlLogEntryId, tx: SubmittedTransaction): Effects = {
+    // TODO(JM): Skip transient contracts in createdContracts/updateContractKeys. E.g. rewrite this to
+    // fold bottom up (with reversed roots!) and skip creates of archived contracts.
     tx.fold(GenTransaction.TopDown, Effects.empty) {
       case (effects, (nodeId, node)) =>
         node match {
@@ -96,7 +98,8 @@ private[kvutils] object InputsAndEffects {
               updatedContractKeys = create.key
                 .fold(effects.updatedContractKeys)(
                   keyWithMaintainers =>
-                    contractKeyToStateKey(
+                    effects.updatedContractKeys +
+                      (contractKeyToStateKey(
                         GlobalKey(
                           create.coinst.template,
                           forceAbsoluteContractIds(keyWithMaintainers.key))) ->
@@ -104,25 +107,28 @@ private[kvutils] object InputsAndEffects {
                           .setContractId(encodeRelativeContractId(
                             entryId,
                             create.coid.asInstanceOf[RelativeContractId]))
-                        .build
+                          .build))
-                      :: effects.updatedContractKeys)
             )
+
           case exe: NodeExercises[_, ContractId, VersionedValue[ContractId]] =>
             if (exe.consuming) {
-              exe.targetCoid match {
+              val stateKey = exe.targetCoid match {
                 case acoid: AbsoluteContractId =>
+                  absoluteContractIdToStateKey(acoid)
+                case rcoid: RelativeContractId =>
+                  relativeContractIdToStateKey(entryId, rcoid)
+              }
               effects.copy(
-                    consumedContracts = absoluteContractIdToStateKey(acoid) :: effects.consumedContracts,
+                consumedContracts = stateKey :: effects.consumedContracts,
                 updatedContractKeys = exe.key
                   .fold(effects.updatedContractKeys)(
                     key =>
-                          contractKeyToStateKey(
+                      effects.updatedContractKeys +
+                        (contractKeyToStateKey(
                           GlobalKey(exe.templateId, forceAbsoluteContractIds(key))) ->
-                            DamlContractKeyState.newBuilder.build :: effects.updatedContractKeys)
+                          DamlContractKeyState.newBuilder.build)
+                  )
               )
-                case _ =>
-                  effects
-              }
             } else {
               effects
             }

ledger/participant-state/kvutils/src/main/scala/com/daml/ledger/participant/state/kvutils/Version.scala

@@ -10,6 +10,7 @@ package com.daml.ledger.participant.state.kvutils
   * [since 100.13.26]:
   * - Added metrics to track submission processing.
   * - Use InsertOrdMap to store resulting state in kvutils for deterministic ordering of state key-values.
+  * - Fix bug with transient contract keys, e.g. keys created and archived in same transaction.
   *
   * [since 100.13.21]:
   * - Added 'Envelope' for compressing and versioning kvutils messages that are transmitted

ledger/participant-state/kvutils/src/main/scala/com/daml/ledger/participant/state/kvutils/committing/ProcessTransactionSubmission.scala

@@ -214,17 +214,6 @@ private[kvutils] case class ProcessTransactionSubmission(
       }
 
     sequence2(
-      // Update contract state entries to mark contracts as consumed (checked by 'validateModelConformance')
-      sequence2(effects.consumedContracts.map { key =>
-        for {
-          cs <- getContractState(key).map { cs =>
-            cs.toBuilder
-              .setArchivedAt(buildTimestamp(txLet))
-              .setArchivedByEntry(entryId)
-          }
-          r <- set(key -> DamlStateValue.newBuilder.setContractState(cs).build)
-        } yield r
-      }: _*),
       // Add contract state entries to mark contract activeness (checked by 'validateModelConformance')
       set(effects.createdContracts.map {
         case (key, createNode) =>
@@ -249,6 +238,17 @@ private[kvutils] case class ProcessTransactionSubmission(
           }
           key -> DamlStateValue.newBuilder.setContractState(cs).build
       }),
+      // Update contract state entries to mark contracts as consumed (checked by 'validateModelConformance')
+      sequence2(effects.consumedContracts.map { key =>
+        for {
+          cs <- getContractState(key).map { cs =>
+            cs.toBuilder
+              .setArchivedAt(buildTimestamp(txLet))
+              .setArchivedByEntry(entryId)
+          }
+          r <- set(key -> DamlStateValue.newBuilder.setContractState(cs).build)
+        } yield r
+      }: _*),
       // Update contract state of divulged contracts
       sequence2(blindingInfo.globalDivulgence.map {
         case (absCoid, parties) =>

ledger/participant-state/kvutils/src/test/suite/scala/com/daml/ledger/participant/state/kvutils/KVTest.scala

@@ -141,14 +141,14 @@ object KVTest {
 
   val minMRTDelta: Duration = theDefaultConfig.timeModel.minTtl
 
-  def runCommand(submitter: Party, cmd: Command): KVTest[SubmittedTransaction] =
+  def runCommand(submitter: Party, cmds: Command*): KVTest[SubmittedTransaction] =
     for {
       s <- get[KVTestState]
       tx = s.engine
         .submit(
           Commands(
             submitter = submitter,
-            commands = ImmArray(cmd),
+            commands = ImmArray(cmds),
             ledgerEffectiveTime = s.recordTime,
             commandsReference = "cmds-ref"
           )

ledger/participant-state/kvutils/src/test/suite/scala/com/daml/ledger/participant/state/kvutils/KVUtilsTransactionSpec.scala

@@ -10,7 +10,12 @@ import com.daml.ledger.participant.state.kvutils.DamlKvutils.{
 }
 import com.daml.ledger.participant.state.kvutils.TestHelpers._
 import com.daml.ledger.participant.state.v1.Update
-import com.digitalasset.daml.lf.command.{Command, CreateCommand, ExerciseCommand}
+import com.digitalasset.daml.lf.command.{
+  Command,
+  CreateAndExerciseCommand,
+  CreateCommand,
+  ExerciseCommand
+}
 import com.digitalasset.daml.lf.data.Ref
 import com.digitalasset.daml.lf.transaction.Node.NodeCreate
 import com.digitalasset.daml.lf.value.Value.{AbsoluteContractId, ValueUnit}
@@ -31,6 +36,11 @@ class KVUtilsTransactionSpec extends WordSpec with Matchers {
         Ref.ContractIdString.assertFromString(coid),
         simpleConsumeChoiceid,
         ValueUnit)
+    val simpleCreateAndExerciseCmd: Command = CreateAndExerciseCommand(
+      simpleTemplateId,
+      mkSimpleTemplateArg("Alice"),
+      simpleConsumeChoiceid,
+      ValueUnit)
 
     val p0 = mkParticipantId(0)
     val p1 = mkParticipantId(1)
@@ -154,5 +164,34 @@ class KVUtilsTransactionSpec extends WordSpec with Matchers {
       }
     }
 
+    "transient contracts and keys are properly archived" in KVTest.runTestWithSimplePackage {
+      for {
+        tx1 <- runCommand(alice, simpleCreateAndExerciseCmd)
+        createAndExerciseTx1 <- submitTransaction(alice, tx1).map(_._2)
+        tx2 <- runCommand(alice, simpleCreateAndExerciseCmd)
+        createAndExerciseTx2 <- submitTransaction(alice, tx2).map(_._2)
+        finalState <- scalaz.State.get[KVTestState]
+      } yield {
+        createAndExerciseTx1.getPayloadCase shouldEqual DamlLogEntry.PayloadCase.TRANSACTION_ENTRY
+        createAndExerciseTx2.getPayloadCase shouldEqual DamlLogEntry.PayloadCase.TRANSACTION_ENTRY
+
+        // Check that all contracts and keys are in the archived state.
+        finalState.damlState.foreach {
+          case (_, v) =>
+            v.getValueCase match {
+              case DamlKvutils.DamlStateValue.ValueCase.CONTRACT_KEY_STATE =>
+                val cks = v.getContractKeyState
+                cks.hasContractId shouldBe false
+
+              case DamlKvutils.DamlStateValue.ValueCase.CONTRACT_STATE =>
+                val cs = v.getContractState
+                cs.hasArchivedAt shouldBe true
+
+              case _ =>
+                succeed
+            }
+        }
+      }
+    }
   }
 }

ledger/participant-state/kvutils/src/test/suite/scala/com/daml/ledger/participant/state/kvutils/TestHelpers.scala

@@ -31,7 +31,8 @@ object TestHelpers {
           agreement "",
           choices {
             choice Consume (x: Unit) : Unit by Cons @Party [Simple:SimpleTemplate {owner} this] (Nil @Party) to upure @Unit ()
-          }
+          },
+          key @Party (Simple:SimpleTemplate {owner} this) (\ (p: Party) -> Cons @Party [p] (Nil @Party))
         } ;
       }
     """

ledger/test-common/src/main/daml/Test.daml

@@ -427,6 +427,16 @@ template Delegated
         create this
         return ()
 
+      -- Non-consuming choice to test that contract keys created and archived
+      -- by transient contracts are properly handled.
+    controller owner can
+      nonconsuming CreateAnotherAndArchive: ()
+        with
+          k2: Text
+        do
+          cid <- create Delegated with owner = owner, k = k2
+          exercise cid Archive
+
 template Delegation
   with
     owner : Party