From 7af076945581ed6bc32c02ef29873c6ff5d65579 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Zemanovi=C4=8D?= Date: Thu, 25 Jan 2024 14:49:16 +0000 Subject: [PATCH] vp/multitoken: refactor maps and sets usage --- .../namada/src/ledger/native_vp/multitoken.rs | 98 +++++++------------ 1 file changed, 37 insertions(+), 61 deletions(-) diff --git a/crates/namada/src/ledger/native_vp/multitoken.rs b/crates/namada/src/ledger/native_vp/multitoken.rs index afb96c5f81..9ad76a5598 100644 --- a/crates/namada/src/ledger/native_vp/multitoken.rs +++ b/crates/namada/src/ledger/native_vp/multitoken.rs @@ -60,8 +60,10 @@ where let pre: Amount = self.ctx.read_pre(key)?.unwrap_or_default(); let post: Amount = self.ctx.read_post(key)?.unwrap_or_default(); match post.checked_sub(pre) { - Some(diff) => match inc_changes.get_mut(token) { - Some(change) => { + Some(diff) => { + let change = + inc_changes.entry(token.clone()).or_default(); + *change = change.checked_add(diff).ok_or_else(|| { Error::NativeVpError( native_vp::Error::SimpleMessage( @@ -69,67 +71,49 @@ where ), ) })?; - } - None => { - inc_changes.insert(token.clone(), diff); - } - }, + } None => { let diff = pre .checked_sub(post) .expect("Underflow shouldn't happen here"); - match dec_changes.get_mut(token) { - Some(change) => { - change.checked_add(diff).ok_or_else(|| { - Error::NativeVpError( - native_vp::Error::SimpleMessage( - "Overflowed in balance check", - ), - ) - })?; - } - None => { - dec_changes.insert(token.clone(), diff); - } - } - } - } - } else if let Some(token) = is_any_minted_balance_key(key) { - let pre: Amount = self.ctx.read_pre(key)?.unwrap_or_default(); - let post: Amount = self.ctx.read_post(key)?.unwrap_or_default(); - match post.checked_sub(pre) { - Some(diff) => match inc_mints.get_mut(token) { - Some(mint) => { - mint.checked_add(diff).ok_or_else(|| { + let change = + dec_changes.entry(token.clone()).or_default(); + *change = + change.checked_add(diff).ok_or_else(|| { Error::NativeVpError( native_vp::Error::SimpleMessage( "Overflowed in balance check", ), ) })?; - } - None => { - inc_mints.insert(token.clone(), diff); - } - }, + } + } + } else if let Some(token) = is_any_minted_balance_key(key) { + let pre: Amount = self.ctx.read_pre(key)?.unwrap_or_default(); + let post: Amount = self.ctx.read_post(key)?.unwrap_or_default(); + match post.checked_sub(pre) { + Some(diff) => { + let mint = inc_mints.entry(token.clone()).or_default(); + *mint = mint.checked_add(diff).ok_or_else(|| { + Error::NativeVpError( + native_vp::Error::SimpleMessage( + "Overflowed in balance check", + ), + ) + })?; + } None => { let diff = pre .checked_sub(post) .expect("Underflow shouldn't happen here"); - match dec_mints.get_mut(token) { - Some(mint) => { - mint.checked_add(diff).ok_or_else(|| { - Error::NativeVpError( - native_vp::Error::SimpleMessage( - "Overflowed in balance check", - ), - ) - })?; - } - None => { - dec_mints.insert(token.clone(), diff); - } - } + let mint = dec_mints.entry(token.clone()).or_default(); + *mint = mint.checked_add(diff).ok_or_else(|| { + Error::NativeVpError( + native_vp::Error::SimpleMessage( + "Overflowed in balance check", + ), + ) + })?; } } // Check if the minter is set @@ -152,18 +136,10 @@ where } let mut all_tokens = BTreeSet::new(); - inc_changes.keys().for_each(|k| { - all_tokens.insert(k.clone()); - }); - dec_changes.keys().for_each(|k| { - all_tokens.insert(k.clone()); - }); - inc_mints.keys().for_each(|k| { - all_tokens.insert(k.clone()); - }); - dec_mints.keys().for_each(|k| { - all_tokens.insert(k.clone()); - }); + all_tokens.extend(inc_changes.keys().cloned()); + all_tokens.extend(dec_changes.keys().cloned()); + all_tokens.extend(inc_mints.keys().cloned()); + all_tokens.extend(dec_mints.keys().cloned()); Ok(all_tokens.iter().all(|token| { let inc_change =