From 6851e0361e65be5c523ccc3d6f6c47000a8aaaf4 Mon Sep 17 00:00:00 2001
From: finch <finch@plaidfinch.net>
Date: Fri, 11 Aug 2023 18:00:46 -0400
Subject: [PATCH] Follow TODO instruction to replace amount update with
 checked_add_signed

---
 .../shielded-pool/src/component/supply.rs     | 43 ++++---------------
 crates/core/num/src/amount.rs                 |  6 +++
 2 files changed, 15 insertions(+), 34 deletions(-)

diff --git a/crates/core/component/shielded-pool/src/component/supply.rs b/crates/core/component/shielded-pool/src/component/supply.rs
index 730a24b00e..d75d47338f 100644
--- a/crates/core/component/shielded-pool/src/component/supply.rs
+++ b/crates/core/component/shielded-pool/src/component/supply.rs
@@ -62,43 +62,18 @@ pub trait SupplyWrite: StateWrite {
         let key = state_key::token_supply(asset_id);
         let current_supply: Amount = self.get_proto(&key).await?.unwrap_or(0u64).into();
 
-        // TODO: if the token was a staking token, update the total supply of staking tokens
-        // TODO: if the token was a delegation token, update the total supply of staking tokens
-        // using the correct exchange rate for that validator (via `unbonded_amount`)
-        // TODO: if the token was an unbonding token, update the total supply of staking tokens
-        // using a 1:1 exchange rate
+        let new_supply = current_supply.checked_add_signed(change).ok_or_else(|| {
+            anyhow::anyhow!(
+                "over-/under-flow updating token {} supply {} with delta {}",
+                asset_id,
+                current_supply,
+                change
+            )
+        })?;
 
-        // TODO: replace with a single checked_add_signed call when mixed_integer_ops lands in stable (1.66)
-        let new_supply: Amount = if change < 0 {
-            current_supply
-                .value()
-                .checked_sub(change.unsigned_abs())
-                .ok_or_else(|| {
-                    anyhow::anyhow!(
-                        "underflow updating token {} supply {} with delta {}",
-                        asset_id,
-                        current_supply,
-                        change
-                    )
-                })?
-                .into()
-        } else {
-            current_supply
-                .value()
-                .checked_add(change as u128)
-                .ok_or_else(|| {
-                    anyhow::anyhow!(
-                        "overflow updating token {} supply {} with delta {}",
-                        asset_id,
-                        current_supply,
-                        change
-                    )
-                })?
-                .into()
-        };
         tracing::debug!(?current_supply, ?new_supply, ?change);
-
         self.put(key, new_supply);
+
         Ok(())
     }
 }
diff --git a/crates/core/num/src/amount.rs b/crates/core/num/src/amount.rs
index 9c16ec67f9..d747920a03 100644
--- a/crates/core/num/src/amount.rs
+++ b/crates/core/num/src/amount.rs
@@ -57,6 +57,12 @@ impl Amount {
             .map(|inner| Self { inner })
     }
 
+    pub fn checked_add_signed(&self, rhs: i128) -> Option<Self> {
+        Some(Amount {
+            inner: self.inner.checked_add_signed(rhs)?,
+        })
+    }
+
     pub fn saturating_add(&self, rhs: &Self) -> Self {
         Self {
             inner: self.inner.saturating_add(rhs.inner),