Use checked arithmetic in RuneUpdater (#3423)

src/index.rs

@@ -5,7 +5,8 @@ use {
       OutPointValue, RuneEntryValue, RuneIdValue, SatPointValue, SatRange, TxidValue,
     },
     event::Event,
-    reorg::*,
+    lot::Lot,
+    reorg::Reorg,
     updater::Updater,
   },
   super::*,
@@ -39,6 +40,7 @@ pub use self::entry::RuneEntry;
 pub(crate) mod entry;
 pub mod event;
 mod fetcher;
+mod lot;
 mod reorg;
 mod rtx;
 mod updater;
@@ -615,14 +617,14 @@ impl Index {
           log::info!("{}", err.to_string());
 
           match err.downcast_ref() {
-            Some(&ReorgError::Recoverable { height, depth }) => {
+            Some(&reorg::Error::Recoverable { height, depth }) => {
               Reorg::handle_reorg(self, height, depth)?;
             }
-            Some(&ReorgError::Unrecoverable) => {
+            Some(&reorg::Error::Unrecoverable) => {
               self
                 .unrecoverably_reorged
                 .store(true, atomic::Ordering::Relaxed);
-              return Err(anyhow!(ReorgError::Unrecoverable));
+              return Err(anyhow!(reorg::Error::Unrecoverable));
             }
             _ => return Err(err),
           };

src/index/lot.rs

@@ -0,0 +1,163 @@
+use {
+  super::*,
+  std::{
+    cmp::{PartialEq, PartialOrd},
+    ops::{Add, AddAssign, Div, Rem, Sub, SubAssign},
+  },
+};
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Default, Serialize, Deserialize)]
+pub(super) struct Lot(pub(super) u128);
+
+impl Lot {
+  #[cfg(test)]
+  const MAX: Self = Self(u128::MAX);
+
+  pub(super) fn n(self) -> u128 {
+    self.0
+  }
+
+  fn checked_add(self, rhs: Self) -> Option<Self> {
+    Some(Self(self.0.checked_add(rhs.0)?))
+  }
+
+  fn checked_sub(self, rhs: Self) -> Option<Self> {
+    Some(Self(self.0.checked_sub(rhs.0)?))
+  }
+}
+
+impl TryFrom<Lot> for usize {
+  type Error = <usize as TryFrom<u128>>::Error;
+  fn try_from(lot: Lot) -> Result<Self, Self::Error> {
+    usize::try_from(lot.0)
+  }
+}
+
+impl Add for Lot {
+  type Output = Self;
+  fn add(self, other: Self) -> Self::Output {
+    self.checked_add(other).expect("lot overflow")
+  }
+}
+
+impl AddAssign for Lot {
+  fn add_assign(&mut self, other: Self) {
+    *self = *self + other;
+  }
+}
+
+impl Add<u128> for Lot {
+  type Output = Self;
+  fn add(self, other: u128) -> Self::Output {
+    self + Lot(other)
+  }
+}
+
+impl AddAssign<u128> for Lot {
+  fn add_assign(&mut self, other: u128) {
+    *self += Lot(other);
+  }
+}
+
+impl Sub for Lot {
+  type Output = Self;
+  fn sub(self, other: Self) -> Self::Output {
+    self.checked_sub(other).expect("lot underflow")
+  }
+}
+
+impl SubAssign for Lot {
+  fn sub_assign(&mut self, other: Self) {
+    *self = *self - other;
+  }
+}
+
+impl Div<u128> for Lot {
+  type Output = Self;
+  fn div(self, other: u128) -> Self::Output {
+    Lot(self.0 / other)
+  }
+}
+
+impl Rem<u128> for Lot {
+  type Output = Self;
+  fn rem(self, other: u128) -> Self::Output {
+    Lot(self.0 % other)
+  }
+}
+
+impl PartialEq<u128> for Lot {
+  fn eq(&self, other: &u128) -> bool {
+    self.0 == *other
+  }
+}
+
+impl PartialOrd<u128> for Lot {
+  fn partial_cmp(&self, other: &u128) -> Option<std::cmp::Ordering> {
+    self.0.partial_cmp(other)
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  #[should_panic(expected = "lot overflow")]
+  fn add() {
+    let _ = Lot::MAX + 1;
+  }
+
+  #[test]
+  #[should_panic(expected = "lot overflow")]
+  fn add_assign() {
+    let mut l = Lot::MAX;
+    l += Lot(1);
+  }
+
+  #[test]
+  #[should_panic(expected = "lot overflow")]
+  fn add_u128() {
+    let _ = Lot::MAX + 1;
+  }
+
+  #[test]
+  #[should_panic(expected = "lot overflow")]
+  fn add_assign_u128() {
+    let mut l = Lot::MAX;
+    l += 1;
+  }
+
+  #[test]
+  #[should_panic(expected = "lot underflow")]
+  fn sub() {
+    let _ = Lot(0) - Lot(1);
+  }
+
+  #[test]
+  #[should_panic(expected = "lot underflow")]
+  fn sub_assign() {
+    let mut l = Lot(0);
+    l -= Lot(1);
+  }
+
+  #[test]
+  fn div() {
+    assert_eq!(Lot(100) / 2, Lot(50));
+  }
+
+  #[test]
+  fn rem() {
+    assert_eq!(Lot(77) % 8, Lot(5));
+  }
+
+  #[test]
+  fn partial_eq() {
+    assert_eq!(Lot(100), 100);
+  }
+
+  #[test]
+  fn partial_ord() {
+    assert!(Lot(100) > 10);
+  }
+}

src/index/reorg.rs

@@ -1,23 +1,23 @@
 use {super::*, updater::BlockData};
 
 #[derive(Debug, PartialEq)]
-pub(crate) enum ReorgError {
+pub(crate) enum Error {
   Recoverable { height: u32, depth: u32 },
   Unrecoverable,
 }
 
-impl Display for ReorgError {
+impl Display for Error {
   fn fmt(&self, f: &mut Formatter) -> fmt::Result {
     match self {
-      ReorgError::Recoverable { height, depth } => {
+      Self::Recoverable { height, depth } => {
         write!(f, "{depth} block deep reorg detected at height {height}")
       }
-      ReorgError::Unrecoverable => write!(f, "unrecoverable reorg detected"),
+      Self::Unrecoverable => write!(f, "unrecoverable reorg detected"),
     }
   }
 }
 
-impl std::error::Error for ReorgError {}
+impl std::error::Error for Error {}
 
 const MAX_SAVEPOINTS: u32 = 2;
 const SAVEPOINT_INTERVAL: u32 = 10;
@@ -43,11 +43,11 @@ impl Reorg {
             .into_option()?;
 
           if index_block_hash == bitcoind_block_hash {
-            return Err(anyhow!(ReorgError::Recoverable { height, depth }));
+            return Err(anyhow!(reorg::Error::Recoverable { height, depth }));
           }
         }
 
-        Err(anyhow!(ReorgError::Unrecoverable))
+        Err(anyhow!(reorg::Error::Unrecoverable))
       }
       _ => Ok(()),
     }

src/index/updater/rune_updater.rs

@@ -2,7 +2,7 @@ use super::*;
 
 pub(super) struct RuneUpdater<'a, 'tx, 'client> {
   pub(super) block_time: u32,
-  pub(super) burned: HashMap<RuneId, u128>,
+  pub(super) burned: HashMap<RuneId, Lot>,
   pub(super) client: &'client Client,
   pub(super) height: u32,
   pub(super) id_to_entry: &'a mut Table<'tx, RuneIdValue, RuneEntryValue>,
@@ -22,7 +22,7 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
 
     let mut unallocated = self.unallocated(tx)?;
 
-    let mut allocated: Vec<HashMap<RuneId, u128>> = vec![HashMap::new(); tx.output.len()];
+    let mut allocated: Vec<HashMap<RuneId, Lot>> = vec![HashMap::new(); tx.output.len()];
 
     if let Some(artifact) = &artifact {
       if let Some(id) = artifact.mint() {
@@ -40,6 +40,8 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
         }
 
         for Edict { id, amount, output } in runestone.edicts.iter().copied() {
+          let amount = Lot(amount);
+
           // edicts with output values greater than the number of outputs
           // should never be produced by the edict parser
           let output = usize::try_from(output).unwrap();
@@ -59,7 +61,7 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
             continue;
           };
 
-          let mut allocate = |balance: &mut u128, amount: u128, output: usize| {
+          let mut allocate = |balance: &mut Lot, amount: Lot, output: usize| {
             if amount > 0 {
               *balance -= amount;
               *allocated[output].entry(id).or_default() += amount;
@@ -113,7 +115,7 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
       }
     }
 
-    let mut burned: HashMap<RuneId, u128> = HashMap::new();
+    let mut burned: HashMap<RuneId, Lot> = HashMap::new();
 
     if let Some(Artifact::Cenotaph(_)) = artifact {
       for (id, balance) in unallocated {
@@ -165,20 +167,20 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
       // increment burned balances
       if tx.output[vout].script_pubkey.is_op_return() {
         for (id, balance) in &balances {
-          *burned.entry(*id).or_default() += balance;
+          *burned.entry(*id).or_default() += *balance;
         }
         continue;
       }
 
       buffer.clear();
 
-      let mut balances = balances.into_iter().collect::<Vec<(RuneId, u128)>>();
+      let mut balances = balances.into_iter().collect::<Vec<(RuneId, Lot)>>();
 
       // Sort balances by id so tests can assert balances in a fixed order
       balances.sort();
 
       for (id, balance) in balances {
-        Index::encode_rune_balance(id, balance, &mut buffer);
+        Index::encode_rune_balance(id, balance.n(), &mut buffer);
       }
 
       self.outpoint_to_balances.insert(
@@ -202,7 +204,7 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
   pub(super) fn update(self) -> Result {
     for (rune_id, burned) in self.burned {
       let mut entry = RuneEntry::load(self.id_to_entry.get(&rune_id.store())?.unwrap().value());
-      entry.burned += burned;
+      entry.burned = entry.burned.checked_add(burned.n()).unwrap();
       self.id_to_entry.insert(&rune_id.store(), entry.store())?;
     }
 
@@ -336,7 +338,7 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
     )))
   }
 
-  fn mint(&mut self, id: RuneId) -> Result<Option<u128>> {
+  fn mint(&mut self, id: RuneId) -> Result<Option<Lot>> {
     let Some(entry) = self.id_to_entry.get(&id.store())? else {
       return Ok(None);
     };
@@ -353,7 +355,7 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
 
     self.id_to_entry.insert(&id.store(), rune_entry.store())?;
 
-    Ok(Some(amount))
+    Ok(Some(Lot(amount)))
   }
 
   fn tx_commits_to_rune(&self, tx: &Transaction, rune: Rune) -> Result<bool> {
@@ -408,9 +410,9 @@ impl<'a, 'tx, 'client> RuneUpdater<'a, 'tx, 'client> {
     Ok(false)
   }
 
-  fn unallocated(&mut self, tx: &Transaction) -> Result<HashMap<RuneId, u128>> {
+  fn unallocated(&mut self, tx: &Transaction) -> Result<HashMap<RuneId, Lot>> {
     // map of rune ID to un-allocated balance of that rune
-    let mut unallocated: HashMap<RuneId, u128> = HashMap::new();
+    let mut unallocated: HashMap<RuneId, Lot> = HashMap::new();
 
     // increment unallocated runes with the runes in tx inputs
     for input in &tx.input {