Fix: Migration Data Id in MigrationDataEvent

docs/spec/mode-a-spec.md

@@ -70,7 +70,7 @@ note_hash = poseidon2_hash_with_separator(
 
 The `#[event]` macro does not support generic structs, so `MigrationDataEvent` implements `EventInterface` manually with `#[derive(Serialize)]`.
 
-Events do not include a note-identifying hash. Instead, wallet clients match events to notes by **emission order**: each `lock_state` call emits a `create_note` followed immediately by the corresponding `MigrationDataEvent`, so the i-th note and i-th event always correspond. *(Source: `migration_lock.nr:46–68`)*
+Each event carries a `data_id` field that identifies the kind of migration data it contains. Within a single `MigrationLock` chain, `data_id` auto-increments from 0. When a contract uses multiple `MigrationLock` instances (e.g. separate entrypoints for private and public state), use `MigrationLock::new_with_offset` to assign non-overlapping `data_id` ranges. Wallet clients match events to notes by **emission order** within a transaction: each `lock_state` call emits a `create_note` followed immediately by the corresponding `MigrationDataEvent`, so the i-th note and i-th event always correspond. *(Source: `migration_lock.nr`, `migration_data_event.nr`)*
 
 ## Claim Flow (Library Level)
 

noir/aztec-state-migration/src/mode_a/migration_data_event.nr

@@ -10,6 +10,7 @@ use aztec::protocol::traits::Serialize;
 /// See https://github.com/AztecProtocol/aztec-packages/blob/c6b2615d5119cf306e8b9c903fb00dad08560d32/noir-projects/aztec-nr/aztec/src/macros/utils.nr#L154
 #[derive(Serialize)]
 pub struct MigrationDataEvent<T> {
+    pub data_id: Field,
     pub migration_data: T,
 }
 

noir/aztec-state-migration/src/mode_a/migration_lock.nr

@@ -16,7 +16,51 @@ use std::embedded_curve_ops::EmbeddedCurvePoint as Point;
 ///   MigrationLock::new(context, mpk, owner, destination_rollup)
 ///       .lock_state(migration_data)   // repeat for each piece of migration data
 ///       .finish();                    // finalize migration (no-op)
+/// ```
+///
+/// It is recommended to "batch" migrated state into one `.lock_state(...)` call:
+///
 ///```
+///   #[derive(Packable, Serialize)]
+///   struct MigrationData {
+///       state1: ...,
+///       state2: ...,
+///        ...
+///   }
+///
+/// ...
+///
+///   let migration_data = MigrationData { state1, state2, ... };
+///
+///   MigrationLock::new(context, mpk, owner, destination_rollup)
+///       .lock_state(migration_data)
+///       .finish();
+///```
+///
+/// ## Multiple MigrationLock instances:
+/// Each `.lock_state(...)` call creates a new migration note and emits a corresponding event
+/// with the migration data. To avoid data ID collisions in the events, when using multiple `MigrationLock` instances,
+/// you should use `MigrationLock::new_with_offset` to set a unique starting data ID for each instance.
+///
+/// ```
+///   #[external("private")]
+///   fn lock_for_migration_1(destination_rollup: Field, mpk: Point) {
+///       ...
+///
+///       MigrationLock::new(self.context, mpk, note_owner, destination_rollup)
+///           .lock_state(state1)
+///           .finish();
+///   }
+///
+///   #[external("private")]
+///   fn lock_for_migration_2(destination_rollup: Field, mpk: Point) {
+///       ...
+///
+///       MigrationLock::new_with_offset(self.context, mpk, note_owner, destination_rollup, 1)
+///           .lock_state(state2)
+///           .finish();
+///   }
+/// ```
 ///
 /// The `.finish()` method is a no-op.
 /// It's included to surpress warnings about unused variables and
@@ -27,6 +71,7 @@ pub struct MigrationLock {
     owner: AztecAddress,
     destination_rollup: Field,
     notes_creator: AztecAddress,
+    next_data_id: Field,
 }
 
 impl MigrationLock {
@@ -40,10 +85,29 @@ impl MigrationLock {
             (mpk.y * mpk.y == mpk.x * mpk.x * mpk.x - 17) & (!mpk.is_infinite),
             "mpk not on Grumpkin curve",
         );
-        Self { context, mpk, owner, destination_rollup, notes_creator: context.this_address() }
+        Self {
+            context,
+            mpk,
+            owner,
+            destination_rollup,
+            notes_creator: context.this_address(),
+            next_data_id: 0,
+        }
     }
 
-    pub fn lock_state<T: Packable + Serialize>(self, migration_data: T) -> Self {
+    pub fn new_with_offset(
+        context: &mut PrivateContext,
+        mpk: Point,
+        owner: AztecAddress,
+        destination_rollup: Field,
+        data_id_offset: Field,
+    ) -> Self {
+        let mut lock = Self::new(context, mpk, owner, destination_rollup);
+        lock.next_data_id = data_id_offset;
+        lock
+    }
+
+    pub fn lock_state<T: Packable + Serialize>(mut self, migration_data: T) -> Self {
         let migration_note = MigrationNote::new(
             self.notes_creator,
             self.mpk,
@@ -60,10 +124,12 @@ impl MigrationLock {
 
         // Emit migration data as an encrypted private event so the recipient can
         // reconstruct the original data when claiming on the new rollup.
-        emit_event_in_private(self.context, MigrationDataEvent { migration_data }).deliver_to(
-            self.owner,
-            MessageDelivery.ONCHAIN_CONSTRAINED,
-        );
+        emit_event_in_private(
+            self.context,
+            MigrationDataEvent { data_id: self.next_data_id, migration_data },
+        )
+            .deliver_to(self.owner, MessageDelivery.ONCHAIN_CONSTRAINED);
+        self.next_data_id += 1;
         self
     }
 

ts/aztec-state-migration/mode-a/types.ts

@@ -46,5 +46,6 @@ export const MigrationNote = {
 /** A Mode A migration note paired with its decoded migration data. */
 export interface MigrationNoteAndData<T> {
   note: NoteDao;
+  dataId: number;
   data: T;
 }

ts/aztec-state-migration/wallet/migration-base-wallet.ts

@@ -176,7 +176,10 @@ export abstract class MigrationBaseWallet extends BaseWallet {
    * filtering on the well-known {@link MIGRATION_NOTE_STORAGE_SLOT} storage slot.
    *
    * @typeParam T - The shape of the migration data (e.g. `bigint` for token amounts).
-   * @param abiType - A single ABI type applied to all events.
+   * @param contractAddress - The contract address to fetch migration notes from.
+   * @param owner - The owner address to filter migration notes.
+   * @param abiType - The ABI type to decode migration data.
+   * @param scopes - Optional scopes to filter migration notes (default: [owner]).
    */
   async getMigrationNotesAndData<T>(
     contractAddress: AztecAddress,
@@ -188,51 +191,43 @@ export abstract class MigrationBaseWallet extends BaseWallet {
       await this.getMigrationNotesEventSelector(contractAddress, owner, scopes);
     let notesAndData: MigrationNoteAndData<T>[] = [];
     for (const [txHash, notes] of noteByTxHash) {
-      const events = await this.getPrivateEvents<T>(
-        {
-          eventSelector,
-          abiType: abiType,
-          fieldNames: ["migration_data"],
-        },
+      const pxeEvents = await this.pxe.getPrivateEvents(
+        eventSelector,
         eventFilter(txHash),
       );
-      if (events.length !== notes.length) {
+      if (pxeEvents.length !== notes.length) {
         throw new Error(
-          `Mismatched number of events (${events.length}) and notes (${notes.length}) for tx ${txHash}.`,
+          `Mismatched number of events (${pxeEvents.length}) and notes (${notes.length}) for tx ${txHash}.`,
         );
       }
       for (let i = 0; i < notes.length; i++) {
-        notesAndData.push({ note: notes[i], data: events[i].event });
+        const dataId = pxeEvents[i].packedEvent[0].toNumber();
+        const rest = pxeEvents[i].packedEvent.slice(1);
+        const data = decodeFromAbi([abiType], rest) as T;
+        notesAndData.push({ note: notes[i], dataId, data });
       }
     }
     return notesAndData;
   }
+
   /**
    * Like {@link getMigrationNotesAndData}, but for mixed-type data structures.
    *
    * When a single `lock_state` chain emits events with different data structures,
-   * pass an ordered `AbiType[]` where `abiTypes[i]` decodes the i-th event
-   * within each tx (matching `lock_state` call order).
+   * pass a `Record<number, AbiType>` mapping each `dataId` to its decoder.
    *
-   * @param abiTypes - Ordered array of ABI types, one per `lock_state` call.
+   * @param abiTypes - Map from dataId to ABI type for decoding each event kind.
    */
   async getMixedMigrationNotesAndData(
     contractAddress: AztecAddress,
     owner: AztecAddress,
-    abiTypes: AbiType[],
+    abiTypes: Record<number, AbiType>,
     scopes?: AztecAddress[],
   ): Promise<MigrationNoteAndData<unknown>[]> {
     const { noteByTxHash, eventSelector, eventFilter } =
       await this.getMigrationNotesEventSelector(contractAddress, owner, scopes);
-
     let notesAndData: MigrationNoteAndData<unknown>[] = [];
-
     for (const [txHash, notes] of noteByTxHash) {
-      if (abiTypes.length !== notes.length) {
-        throw new Error(
-          `abiTypes array length (${abiTypes.length}) does not match number of notes (${notes.length}) for tx ${txHash}.`,
-        );
-      }
       const pxeEvents = await this.pxe.getPrivateEvents(
         eventSelector,
         eventFilter(txHash),
@@ -243,11 +238,17 @@ export abstract class MigrationBaseWallet extends BaseWallet {
         );
       }
       for (let i = 0; i < notes.length; i++) {
-        const decodedEvent = decodeFromAbi(
-          [abiTypes[i]],
-          pxeEvents[i].packedEvent,
-        ) as unknown;
-        notesAndData.push({ note: notes[i], data: decodedEvent });
+        const dataId = pxeEvents[i].packedEvent[0].toNumber();
+        const rest = pxeEvents[i].packedEvent.slice(1);
+        const abiType = abiTypes[dataId];
+        if (!abiType) {
+          this.log.warn(
+            `Unknown migration dataId ${dataId} in tx ${txHash}, skipping note ${notes[i].noteHash}.`,
+          );
+          continue;
+        }
+        const data = decodeFromAbi([abiType], rest) as unknown;
+        notesAndData.push({ note: notes[i], dataId, data });
       }
     }
     return notesAndData;