From 46b469115499593c615b0498d13a8c2a1bb45f81 Mon Sep 17 00:00:00 2001 From: JanKuczma Date: Mon, 2 Mar 2026 17:39:28 +0100 Subject: [PATCH 1/3] migraiton data id in event --- .../src/mode_a/migration_data_event.nr | 1 + .../src/mode_a/migration_lock.nr | 22 +++++--- ts/aztec-state-migration/mode-a/types.ts | 1 + .../wallet/migration-base-wallet.ts | 53 ++++++++++--------- 4 files changed, 45 insertions(+), 32 deletions(-) diff --git a/noir/aztec-state-migration/src/mode_a/migration_data_event.nr b/noir/aztec-state-migration/src/mode_a/migration_data_event.nr index cde238c..3e3cbf4 100644 --- a/noir/aztec-state-migration/src/mode_a/migration_data_event.nr +++ b/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 { + pub data_id: Field, pub migration_data: T, } diff --git a/noir/aztec-state-migration/src/mode_a/migration_lock.nr b/noir/aztec-state-migration/src/mode_a/migration_lock.nr index b29da4f..9641c2a 100644 --- a/noir/aztec-state-migration/src/mode_a/migration_lock.nr +++ b/noir/aztec-state-migration/src/mode_a/migration_lock.nr @@ -27,6 +27,7 @@ pub struct MigrationLock { owner: AztecAddress, destination_rollup: Field, notes_creator: AztecAddress, + next_data_id: Field, } impl MigrationLock { @@ -40,10 +41,17 @@ 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(self, migration_data: T) -> Self { + pub fn lock_state(mut self, migration_data: T) -> Self { let migration_note = MigrationNote::new( self.notes_creator, self.mpk, @@ -60,10 +68,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 } diff --git a/ts/aztec-state-migration/mode-a/types.ts b/ts/aztec-state-migration/mode-a/types.ts index 90e758a..316164e 100644 --- a/ts/aztec-state-migration/mode-a/types.ts +++ b/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 { note: NoteDao; + dataId: number; data: T; } diff --git a/ts/aztec-state-migration/wallet/migration-base-wallet.ts b/ts/aztec-state-migration/wallet/migration-base-wallet.ts index 08144d5..c9f693b 100644 --- a/ts/aztec-state-migration/wallet/migration-base-wallet.ts +++ b/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( contractAddress: AztecAddress, @@ -188,51 +191,43 @@ export abstract class MigrationBaseWallet extends BaseWallet { await this.getMigrationNotesEventSelector(contractAddress, owner, scopes); let notesAndData: MigrationNoteAndData[] = []; for (const [txHash, notes] of noteByTxHash) { - const events = await this.getPrivateEvents( - { - 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` 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, scopes?: AztecAddress[], ): Promise[]> { const { noteByTxHash, eventSelector, eventFilter } = await this.getMigrationNotesEventSelector(contractAddress, owner, scopes); - let notesAndData: MigrationNoteAndData[] = []; - 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; From 1f248f7b3979811955de241dca473aa1c34ac7dd Mon Sep 17 00:00:00 2001 From: JanKuczma Date: Mon, 2 Mar 2026 18:15:21 +0100 Subject: [PATCH 2/3] add "with_offset" lock state builder constructor --- .../src/mode_a/migration_lock.nr | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/noir/aztec-state-migration/src/mode_a/migration_lock.nr b/noir/aztec-state-migration/src/mode_a/migration_lock.nr index 9641c2a..5f546dc 100644 --- a/noir/aztec-state-migration/src/mode_a/migration_lock.nr +++ b/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 @@ -51,6 +95,18 @@ impl MigrationLock { } } + 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(mut self, migration_data: T) -> Self { let migration_note = MigrationNote::new( self.notes_creator, From 9128a610e4a38d03388aa60c77667d938432827b Mon Sep 17 00:00:00 2001 From: JanKuczma Date: Mon, 2 Mar 2026 18:33:54 +0100 Subject: [PATCH 3/3] upd doc --- docs/spec/mode-a-spec.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/spec/mode-a-spec.md b/docs/spec/mode-a-spec.md index ae6e9be..c9dbce1 100644 --- a/docs/spec/mode-a-spec.md +++ b/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)