From ffe0d971fa2d7bdfa41e27f936a7351b3f2e2c01 Mon Sep 17 00:00:00 2001 From: Andrus Salumets Date: Mon, 22 Jun 2026 14:31:24 +0200 Subject: [PATCH] feat(wallet): track leader reward outputs (#2978) --- core/src/events/mod.rs | 18 +++++- core/src/mantle/channel.rs | 36 ++++++----- core/src/mantle/ops/leader_claim.rs | 66 +++++++++++++++++++- ledger/src/lib.rs | 39 +++++++----- services/wallet/src/lib.rs | 23 ++++++- tests/src/tests/mantle/channel.rs | 22 ++++--- wallet/src/lib.rs | 94 ++++++++++++++++++++++++----- zone-sdk/src/adapter.rs | 2 +- 8 files changed, 240 insertions(+), 60 deletions(-) diff --git a/core/src/events/mod.rs b/core/src/events/mod.rs index 6ab592bab..7530fbc15 100644 --- a/core/src/events/mod.rs +++ b/core/src/events/mod.rs @@ -5,8 +5,11 @@ use crate::{ codec::{DeserializeOp as _, SerializeOp as _}, crypto::Hash, mantle::{ - TxHash, Value, - ops::channel::{ChannelId, deposit::Metadata}, + TxHash, Utxo, Value, + ops::{ + channel::{ChannelId, deposit::Metadata}, + leader_claim::VoucherNullifier, + }, }, }; @@ -45,6 +48,12 @@ impl FromIterator for Events { } } +impl From for Events { + fn from(event: Event) -> Self { + Self(vec![event]) + } +} + /// An event generated by the execution of block/transaction/operation #[derive(Debug, Clone, Serialize, Deserialize)] pub enum Event { @@ -79,6 +88,11 @@ pub enum EventPayload { amount: Value, metadata: Metadata, }, + /// A leader claim operation created the reward note for its beneficiary. + LeaderRewardClaimed { + voucher_nullifier: VoucherNullifier, + utxo: Utxo, + }, } impl TryFrom for Events { diff --git a/core/src/mantle/channel.rs b/core/src/mantle/channel.rs index 042b621bf..2d73816cb 100644 --- a/core/src/mantle/channel.rs +++ b/core/src/mantle/channel.rs @@ -411,24 +411,32 @@ mod tests { ); assert_eq!(events.len(), 1); - let Event::Tx { + let Some(Event::Tx { tx_hash, op_id, - payload, - } = events.iter().next().cloned().unwrap() + payload: + EventPayload::Deposit { + channel_id: event_channel_id, + amount, + metadata, + }, + }) = events.iter().find(|event| { + matches!( + event, + Event::Tx { + payload: EventPayload::Deposit { .. }, + .. + } + ) + }) else { - panic!("expected Tx event") + panic!("events should include deposit event") }; - assert_eq!(tx_hash, [0; 32].into()); - assert_eq!(op_id, deposit_op.op_id()); - let EventPayload::Deposit { - channel_id, - amount, - metadata, - } = payload; - assert_eq!(channel_id, deposit_op.channel_id); - assert_eq!(amount, utxo.note.value); - assert_eq!(metadata, deposit_op.metadata); + assert_eq!(*tx_hash, [0; 32].into()); + assert_eq!(*op_id, deposit_op.op_id()); + assert_eq!(*event_channel_id, deposit_op.channel_id); + assert_eq!(*amount, utxo.note.value); + assert_eq!(*metadata, deposit_op.metadata); } #[test] diff --git a/core/src/mantle/ops/leader_claim.rs b/core/src/mantle/ops/leader_claim.rs index d58912c48..c39dbcfa8 100644 --- a/core/src/mantle/ops/leader_claim.rs +++ b/core/src/mantle/ops/leader_claim.rs @@ -8,7 +8,7 @@ use thiserror::Error; use crate::{ crypto::ZkHasher, - events::Events, + events::{Event, EventPayload, Events}, mantle::{ Note, TxHash, Utxo, Value, encoding::encode_leader_claim, @@ -168,6 +168,7 @@ pub struct LeaderClaimExecutionContext { pub reward_amount: Value, pub claimable_rewards: Value, pub utxos: Utxos, + pub tx_hash: TxHash, } impl Operation> for LeaderClaimOp { @@ -212,8 +213,20 @@ impl Operation> for LeaderClaimOp { // Remove the distributed rewards from the pool ctx.claimable_rewards -= ctx.reward_amount; + let tx_hash = ctx.tx_hash; - Ok((ctx, Events::new())) + Ok(( + ctx, + Event::from_tx( + tx_hash, + self.op_id(), + EventPayload::LeaderRewardClaimed { + voucher_nullifier: self.voucher_nullifier, + utxo, + }, + ) + .into(), + )) } } @@ -254,4 +267,53 @@ mod tests { assert_eq!(op.validate(&ctx), Ok(())); } + + #[test] + fn execute_emits_leader_reward_claimed_event() { + let voucher_secret = VoucherSecret::from(Fr::from(7u64)); + let reward_amount = 38; + let pk = ZkPublicKey::zero(); + let tx_hash = TxHash::from([11u8; 32]); + let op = LeaderClaimOp { + rewards_root: RewardsRoot::default(), + voucher_nullifier: VoucherNullifier::from_secret(voucher_secret), + pk, + }; + + let (ctx, events) = op + .execute(LeaderClaimExecutionContext { + nullifiers: rpds::HashTrieSetSync::new_sync(), + reward_amount, + claimable_rewards: 100, + utxos: Utxos::new(), + tx_hash, + }) + .expect("leader claim execution should succeed"); + + assert!(ctx.nullifiers.contains(&op.voucher_nullifier)); + assert_eq!(ctx.claimable_rewards, 62); + assert_eq!( + ctx.utxos.get(&op.utxo(reward_amount).id()), + Some(op.utxo(reward_amount)) + ); + + let mut events = events.iter(); + let Some(Event::Tx { + tx_hash: event_tx_hash, + op_id, + payload: + EventPayload::LeaderRewardClaimed { + voucher_nullifier, + utxo, + }, + }) = events.next() + else { + panic!("expected LeaderRewardClaimed tx event"); + }; + assert_eq!(*event_tx_hash, tx_hash); + assert_eq!(*op_id, op.op_id()); + assert_eq!(*voucher_nullifier, op.voucher_nullifier); + assert_eq!(*utxo, op.utxo(reward_amount)); + assert!(events.next().is_none()); + } } diff --git a/ledger/src/lib.rs b/ledger/src/lib.rs index a00d1dccb..31faff368 100644 --- a/ledger/src/lib.rs +++ b/ledger/src/lib.rs @@ -681,6 +681,7 @@ impl LedgerState { reward_amount: self.mantle_ledger.leaders.reward_amount(), claimable_rewards: self.mantle_ledger.leaders.claimable_rewards(), utxos: self.cryptarchia_ledger.latest_utxos().clone(), + tx_hash, }) .map_err(mantle::Error::LeaderClaim)?; self.mantle_ledger @@ -1120,24 +1121,32 @@ mod tests { assert_eq!(balance, Balance::from(0)); assert_eq!(events.len(), 1); - let Event::Tx { - tx_hash, + let Some(Event::Tx { + tx_hash: event_tx_hash, op_id, - payload, - } = events.iter().next().unwrap().clone() + payload: + EventPayload::Deposit { + channel_id: event_channel_id, + amount, + metadata, + }, + }) = events.iter().find(|event| { + matches!( + event, + Event::Tx { + payload: EventPayload::Deposit { .. }, + .. + } + ) + }) else { - panic!("expected a Tx event") + panic!("events should include deposit event") }; - assert_eq!(tx_hash, tx.hash()); - assert_eq!(op_id, deposit.op_id()); - let EventPayload::Deposit { - channel_id, - amount, - metadata, - } = payload; - assert_eq!(channel_id, deposit.channel_id); - assert_eq!(amount, utxo.note.value); - assert_eq!(metadata, deposit.metadata); + assert_eq!(*event_tx_hash, tx.hash()); + assert_eq!(*op_id, deposit.op_id()); + assert_eq!(*event_channel_id, deposit.channel_id); + assert_eq!(*amount, utxo.note.value); + assert_eq!(*metadata, deposit.metadata); } #[test] diff --git a/services/wallet/src/lib.rs b/services/wallet/src/lib.rs index 669d5d409..7568f00d5 100644 --- a/services/wallet/src/lib.rs +++ b/services/wallet/src/lib.rs @@ -1262,8 +1262,9 @@ where return; }; + let events = Self::load_block_events(header_id, storage_adapter).await; let wallet_block = - WalletBlock::from_block(&block, epoch_config.epoch(block.header().slot())); + WalletBlock::from_block(&block, epoch_config.epoch(block.header().slot()), &events); match state.apply_block(&wallet_block) { Ok(()) => { trace!(target: LOG_TARGET, block_id = ?wallet_block.id, "Applied block to wallet"); @@ -1311,6 +1312,23 @@ where .ok_or(WalletServiceError::BlockNotFoundInStorage(header_id)) } + async fn load_block_events( + header_id: HeaderId, + storage_adapter: &StorageAdapter, + ) -> Events { + storage_adapter + .get_block_events(&header_id) + .await + .unwrap_or_else(|| { + warn!( + target: LOG_TARGET, + block_id = ?header_id, + "Failed to load block events for wallet" + ); + Events::new() + }) + } + async fn handle_lib_update( lib_update: &LibUpdate, storage_adapter: &StorageAdapter, @@ -1417,8 +1435,9 @@ where } let block = Self::load_block(header_id, storage_adapter).await?; + let events = Self::load_block_events(header_id, storage_adapter).await; let wallet_block = - WalletBlock::from_block(&block, epoch_config.epoch(block.header().slot())); + WalletBlock::from_block(&block, epoch_config.epoch(block.header().slot()), &events); if let Err(e) = state.apply_block(&wallet_block) { error!( diff --git a/tests/src/tests/mantle/channel.rs b/tests/src/tests/mantle/channel.rs index 528537ac0..aa833383b 100644 --- a/tests/src/tests/mantle/channel.rs +++ b/tests/src/tests/mantle/channel.rs @@ -152,20 +152,22 @@ async fn channel_deposit() { .expect("timed out waiting for the deposit tx to be included in a block"); let events = fetch_block_events(&validator.client, deposit_block_id).await; - let payload = events + let (channel_id, amount, metadata) = events .iter() .find_map(|event| match event { Event::Tx { - tx_hash, payload, .. - } => (tx_hash == &deposit_tx_hash).then(|| payload.clone()), - Event::Ledger(_) => None, + tx_hash, + payload: + EventPayload::Deposit { + channel_id, + amount, + metadata, + }, + .. + } if tx_hash == &deposit_tx_hash => Some((*channel_id, *amount, metadata.clone())), + _ => None, }) - .expect("block events should include the deposit tx"); - let EventPayload::Deposit { - channel_id, - amount, - metadata, - } = payload; + .expect("block events should include the deposit event"); assert_eq!(channel_id, deposit_op.channel_id); assert_eq!(amount, deposit_amount); assert_eq!(metadata, deposit_op.metadata); diff --git a/wallet/src/lib.rs b/wallet/src/lib.rs index db99217a2..96e6cd033 100644 --- a/wallet/src/lib.rs +++ b/wallet/src/lib.rs @@ -12,6 +12,7 @@ pub use error::WalletError; use lb_core::{ block::Block, crypto::ZkHasher, + events::{Event, EventPayload, Events}, header::HeaderId, mantle::{ AuthenticatedMantleTx, GasConstants, NoteId, Utxo, Value, @@ -43,6 +44,7 @@ pub struct WalletBlock { pub epoch: Epoch, pub voucher_cm: VoucherCm, pub spent_notes: Vec, + pub leader_reward_utxos: Vec, pub transfers: Vec, pub locked_notes: HashSet, pub unlocked_notes: HashSet, @@ -50,12 +52,13 @@ pub struct WalletBlock { impl WalletBlock { #[must_use] - pub fn from_block(block: &Block, epoch: Epoch) -> Self + pub fn from_block(block: &Block, epoch: Epoch, events: &Events) -> Self where Tx: AuthenticatedMantleTx, { // TODO: handle inputs/outputs of ALL operations: https://github.com/logos-blockchain/logos-blockchain/issues/2627 let mut spent_notes = Vec::new(); + let mut leader_reward_utxos = Vec::new(); let mut transfers = Vec::new(); let mut locked_notes = HashSet::new(); let mut unlocked_notes = HashSet::new(); @@ -81,12 +84,15 @@ impl WalletBlock { } } + leader_reward_utxos.extend(leader_reward_utxos_from_events(events)); + Self { id: block.header().id(), parent: block.header().parent(), epoch, voucher_cm: *block.header().leader_proof().voucher_cm(), spent_notes, + leader_reward_utxos, transfers, locked_notes, unlocked_notes, @@ -247,20 +253,14 @@ impl WalletState { remove_spent_utxo(spent_id, &mut utxos, &mut pk_index); } + for utxo in &block.leader_reward_utxos { + insert_utxo_if_owned(*utxo, known_keys, &mut utxos, &mut pk_index); + } + for transfer in &block.transfers { // Add new UTXOs (outputs) - only if they belong to our known keys for utxo in transfer.outputs.utxos(transfer) { - if known_keys.contains_key(&utxo.note.pk) { - let note_id = utxo.id(); - utxos = utxos.insert(note_id, utxo); - - let note_set = pk_index - .get(&utxo.note.pk) - .cloned() - .unwrap_or_else(rpds::HashTrieSetSync::new_sync) - .insert(note_id); - pk_index = pk_index.insert(utxo.note.pk, note_set); - } + insert_utxo_if_owned(utxo, known_keys, &mut utxos, &mut pk_index); } } @@ -340,6 +340,37 @@ impl WalletState { } } +fn insert_utxo_if_owned( + utxo: Utxo, + known_keys: &HashMap, + utxos: &mut rpds::HashTrieMapSync, + pk_index: &mut rpds::HashTrieMapSync>, +) { + if !known_keys.contains_key(&utxo.note.pk) { + return; + } + + let note_id = utxo.id(); + utxos.insert_mut(note_id, utxo); + + let note_set = pk_index + .get(&utxo.note.pk) + .cloned() + .unwrap_or_else(rpds::HashTrieSetSync::new_sync) + .insert(note_id); + pk_index.insert_mut(utxo.note.pk, note_set); +} + +fn leader_reward_utxos_from_events(events: &Events) -> impl Iterator + '_ { + events.iter().filter_map(|event| match event { + Event::Tx { + payload: EventPayload::LeaderRewardClaimed { utxo, .. }, + .. + } => Some(*utxo), + _ => None, + }) +} + fn remove_spent_utxo( spent_id: &NoteId, utxos: &mut rpds::HashTrieMapSync, @@ -581,7 +612,7 @@ mod tests { use lb_core::{ crypto::{Hash, ZkDigest as _}, mantle::{ - Note, + Note, TxHash, channel::Channels, gas::MainnetGasConstants as Gas, ledger::{Inputs, Outputs}, @@ -721,6 +752,7 @@ mod tests { epoch: 1.into(), voucher_cm: v1_cm, spent_notes: vec![], + leader_reward_utxos: vec![], transfers: vec![transfer1.clone()], locked_notes: HashSet::from([locked_note]), // Unknown unlocked note that will be ignored @@ -743,6 +775,7 @@ mod tests { epoch: 2.into(), voucher_cm: v2_cm, spent_notes: vec![alice_100_nmo_utxo.id()], + leader_reward_utxos: vec![], transfers: vec![TransferOp { inputs: Inputs::new([alice_100_nmo_utxo.id()]), outputs: Outputs::new([Note::new(20, bob), Note::new(80, alice)]), @@ -792,6 +825,7 @@ mod tests { epoch: 2.into(), voucher_cm: v3_cm, spent_notes: vec![alice_80_nmo_utxo.id()], + leader_reward_utxos: vec![Utxo::new(tx_hash(9), 0, Note::new(38, alice))], transfers: vec![], locked_notes: HashSet::new(), unlocked_notes: HashSet::new(), @@ -800,7 +834,7 @@ mod tests { assert_eq!( wallet.balance(block_3.id, alice).unwrap().unwrap().balance, - 4 + 42 ); assert_eq!( wallet.balance(block_3.id, bob).unwrap().unwrap().balance, @@ -815,6 +849,38 @@ mod tests { assert_not_tracked_voucher(&wallet, block_3.id, &v3_cm); } + #[test] + fn extracts_leader_reward_utxos_from_events() { + let alice = pk(1); + let bob = pk(2); + let alice_utxo = Utxo::new(tx_hash(9), 0, Note::new(38, alice)); + let bob_utxo = Utxo::new(tx_hash(10), 0, Note::new(21, bob)); + let events = [ + Event::from_tx( + TxHash::from(tx_hash(1)), + tx_hash(9), + EventPayload::LeaderRewardClaimed { + voucher_nullifier: voucher(1, 0).1, + utxo: alice_utxo, + }, + ), + Event::from_tx( + TxHash::from(tx_hash(2)), + tx_hash(10), + EventPayload::LeaderRewardClaimed { + voucher_nullifier: voucher(2, 0).1, + utxo: bob_utxo, + }, + ), + ] + .into_iter() + .collect(); + + let utxos = leader_reward_utxos_from_events(&events).collect::>(); + + assert_eq!(utxos, vec![alice_utxo, bob_utxo]); + } + #[test] fn test_fund_tx_with_change() { let alice = pk(1); diff --git a/zone-sdk/src/adapter.rs b/zone-sdk/src/adapter.rs index 65f158d51..ea9daa19a 100644 --- a/zone-sdk/src/adapter.rs +++ b/zone-sdk/src/adapter.rs @@ -237,7 +237,7 @@ pub(crate) fn build_deposit_amounts(events: &Events) -> HashMap<(TxHash, Hash), op_id, payload: EventPayload::Deposit { amount, .. }, } => Some(((*tx_hash, *op_id), *amount)), - Event::Ledger(_) => None, + Event::Tx { .. } | Event::Ledger(_) => None, }) .collect() }