Multisig db (#444)

* implement db macro for processor/multisigs/db.rs

* ran fmt

* cargo +nightly fmt

---------

Co-authored-by: Luke Parker <lukeparker5132@gmail.com>

processor/src/multisigs/db.rs

@@ -1,56 +1,39 @@
-use core::marker::PhantomData;
-
 use ciphersuite::Ciphersuite;
-
 pub use serai_db::*;
 
 use scale::{Encode, Decode};
 use serai_client::{primitives::Balance, in_instructions::primitives::InInstructionWithBalance};
 
 use crate::{
-  Get, Db, Plan,
+  Get, Plan,
   networks::{Transaction, Network},
 };
 
-#[derive(Debug)]
-pub struct MultisigsDb<N: Network, D: Db>(PhantomData<N>, PhantomData<D>);
-impl<N: Network, D: Db> MultisigsDb<N, D> {
-  fn multisigs_key(dst: &'static [u8], key: impl AsRef<[u8]>) -> Vec<u8> {
-    D::key(b"MULTISIGS", dst, key)
+create_db!(
+  MultisigsDb {
+    NextBatchDb: () -> u32,
+    PlanDb: (id: &[u8]) -> Vec<u8>,
+    PlansFromScanningDb: (block_number: u64) -> Vec<u8>,
+    OperatingCostsDb: () -> u64,
+    ResolvedDb: (tx: &[u8]) -> [u8; 32],
+    SigningDb: (key: &[u8]) -> Vec<u8>,
+    ForwardedOutputDb: (balance: Balance) -> Vec<u8>,
+    DelayedOutputDb: () -> Vec<u8>
   }
+);
 
-  fn next_batch_key() -> Vec<u8> {
-    Self::multisigs_key(b"next_batch", [])
-  }
-  // Set the next batch ID to use
-  pub fn set_next_batch_id(txn: &mut D::Transaction<'_>, batch: u32) {
-    txn.put(Self::next_batch_key(), batch.to_le_bytes());
-  }
-  // Get the next batch ID
-  pub fn next_batch_id<G: Get>(getter: &G) -> u32 {
-    getter.get(Self::next_batch_key()).map_or(0, |v| u32::from_le_bytes(v.try_into().unwrap()))
-  }
-
-  fn plan_key(id: &[u8]) -> Vec<u8> {
-    Self::multisigs_key(b"plan", id)
-  }
-  fn resolved_key(tx: &[u8]) -> Vec<u8> {
-    Self::multisigs_key(b"resolved", tx)
-  }
-  fn signing_key(key: &[u8]) -> Vec<u8> {
-    Self::multisigs_key(b"signing", key)
-  }
-  pub fn save_active_plan(
-    txn: &mut D::Transaction<'_>,
+impl PlanDb {
+  pub fn save_active_plan<N: Network>(
+    txn: &mut impl DbTxn,
     key: &[u8],
-    block_number: u64,
+    block_number: usize,
     plan: &Plan<N>,
     operating_costs_at_time: u64,
   ) {
     let id = plan.id();
 
     {
-      let mut signing = txn.get(Self::signing_key(key)).unwrap_or(vec![]);
+      let mut signing = SigningDb::get(txn, key).unwrap_or_default();
 
       // If we've already noted we're signing this, return
       assert_eq!(signing.len() % 32, 0);
@@ -61,25 +44,25 @@ impl<N: Network, D: Db> MultisigsDb<N, D> {
       }
 
       signing.extend(&id);
-      txn.put(Self::signing_key(key), id);
+      SigningDb::set(txn, key, &id);
     }
 
     {
       let mut buf = block_number.to_le_bytes().to_vec();
       plan.write(&mut buf).unwrap();
       buf.extend(&operating_costs_at_time.to_le_bytes());
-      txn.put(Self::plan_key(&id), &buf);
+      Self::set(txn, &id, &buf);
     }
   }
 
-  pub fn active_plans<G: Get>(getter: &G, key: &[u8]) -> Vec<(u64, Plan<N>, u64)> {
-    let signing = getter.get(Self::signing_key(key)).unwrap_or(vec![]);
+  pub fn active_plans<N: Network>(getter: &impl Get, key: &[u8]) -> Vec<(u64, Plan<N>, u64)> {
+    let signing = SigningDb::get(getter, key).unwrap_or_default();
     let mut res = vec![];
 
     assert_eq!(signing.len() % 32, 0);
     for i in 0 .. (signing.len() / 32) {
       let id = &signing[(i * 32) .. ((i + 1) * 32)];
-      let buf = getter.get(Self::plan_key(id)).unwrap();
+      let buf = Self::get(getter, id).unwrap();
 
       let block_number = u64::from_le_bytes(buf[.. 8].try_into().unwrap());
       let plan = Plan::<N>::read::<&[u8]>(&mut &buf[8 ..]).unwrap();
@@ -87,50 +70,41 @@ impl<N: Network, D: Db> MultisigsDb<N, D> {
       let operating_costs = u64::from_le_bytes(buf[(buf.len() - 8) ..].try_into().unwrap());
       res.push((block_number, plan, operating_costs));
     }
-
     res
   }
 
-  fn operating_costs_key() -> Vec<u8> {
-    Self::multisigs_key(b"operating_costs", [])
-  }
-  pub fn take_operating_costs(txn: &mut D::Transaction<'_>) -> u64 {
-    let existing = txn
-      .get(Self::operating_costs_key())
-      .map(|bytes| u64::from_le_bytes(bytes.try_into().unwrap()))
-      .unwrap_or(0);
-    txn.del(Self::operating_costs_key());
-    existing
-  }
-  pub fn set_operating_costs(txn: &mut D::Transaction<'_>, amount: u64) {
-    if amount != 0 {
-      txn.put(Self::operating_costs_key(), amount.to_le_bytes());
-    }
-  }
-
-  pub fn resolved_plan<G: Get>(
-    getter: &G,
-    tx: <N::Transaction as Transaction<N>>::Id,
-  ) -> Option<[u8; 32]> {
-    getter.get(Self::resolved_key(tx.as_ref())).map(|id| id.try_into().unwrap())
-  }
-  pub fn plan_by_key_with_self_change<G: Get>(
-    getter: &G,
+  pub fn plan_by_key_with_self_change<N: Network>(
+    getter: &impl Get,
     key: <N::Curve as Ciphersuite>::G,
     id: [u8; 32],
   ) -> bool {
-    let plan =
-      Plan::<N>::read::<&[u8]>(&mut &getter.get(Self::plan_key(&id)).unwrap()[8 ..]).unwrap();
+    let plan = Plan::<N>::read::<&[u8]>(&mut &Self::get(getter, &id).unwrap()[8 ..]).unwrap();
     assert_eq!(plan.id(), id);
     (key == plan.key) && (Some(N::change_address(plan.key)) == plan.change)
   }
-  pub fn resolve_plan(
-    txn: &mut D::Transaction<'_>,
+}
+
+impl OperatingCostsDb {
+  pub fn take_operating_costs(txn: &mut impl DbTxn) -> u64 {
+    let existing = Self::get(txn).unwrap_or_default();
+    txn.del(Self::key());
+    existing
+  }
+  pub fn set_operating_costs(txn: &mut impl DbTxn, amount: u64) {
+    if amount != 0 {
+      Self::set(txn, &amount);
+    }
+  }
+}
+
+impl ResolvedDb {
+  pub fn resolve_plan<N: Network>(
+    txn: &mut impl DbTxn,
     key: &[u8],
     plan: [u8; 32],
     resolution: <N::Transaction as Transaction<N>>::Id,
   ) {
-    let mut signing = txn.get(Self::signing_key(key)).unwrap_or(vec![]);
+    let mut signing = SigningDb::get(txn, key).unwrap_or_default();
     assert_eq!(signing.len() % 32, 0);
 
     let mut found = false;
@@ -147,17 +121,14 @@ impl<N: Network, D: Db> MultisigsDb<N, D> {
     if !found {
       log::warn!("told to finish signing {} yet wasn't actively signing it", hex::encode(plan));
     }
-
-    txn.put(Self::signing_key(key), signing);
-
-    txn.put(Self::resolved_key(resolution.as_ref()), plan);
+    SigningDb::set(txn, key, &signing);
+    Self::set(txn, resolution.as_ref(), &plan);
   }
+}
 
-  fn plans_from_scanning_key(block_number: usize) -> Vec<u8> {
-    Self::multisigs_key(b"plans_from_scanning", u32::try_from(block_number).unwrap().to_le_bytes())
-  }
-  pub fn set_plans_from_scanning(
-    txn: &mut D::Transaction<'_>,
+impl PlansFromScanningDb {
+  pub fn set_plans_from_scanning<N: Network>(
+    txn: &mut impl DbTxn,
     block_number: usize,
     plans: Vec<Plan<N>>,
   ) {
@@ -165,14 +136,15 @@ impl<N: Network, D: Db> MultisigsDb<N, D> {
     for plan in plans {
       plan.write(&mut buf).unwrap();
     }
-    txn.put(Self::plans_from_scanning_key(block_number), buf);
+    Self::set(txn, block_number.try_into().unwrap(), &buf);
   }
-  pub fn take_plans_from_scanning(
-    txn: &mut D::Transaction<'_>,
+
+  pub fn take_plans_from_scanning<N: Network>(
+    txn: &mut impl DbTxn,
     block_number: usize,
   ) -> Option<Vec<Plan<N>>> {
-    let key = Self::plans_from_scanning_key(block_number);
-    let res = txn.get(&key).map(|plans| {
+    let block_number = u64::try_from(block_number).unwrap();
+    let res = Self::get(txn, block_number).map(|plans| {
       let mut plans_ref = plans.as_slice();
       let mut res = vec![];
       while !plans_ref.is_empty() {
@@ -181,56 +153,46 @@ impl<N: Network, D: Db> MultisigsDb<N, D> {
       res
     });
     if res.is_some() {
-      txn.del(key);
+      txn.del(Self::key(block_number));
     }
     res
   }
+}
 
-  fn forwarded_output_key(balance: Balance) -> Vec<u8> {
-    Self::multisigs_key(b"forwarded_output", balance.encode())
-  }
-  pub fn save_forwarded_output(
-    txn: &mut D::Transaction<'_>,
-    instruction: InInstructionWithBalance,
-  ) {
-    let key = Self::forwarded_output_key(instruction.balance);
-    let mut existing = txn.get(&key).unwrap_or(vec![]);
+impl ForwardedOutputDb {
+  pub fn save_forwarded_output(txn: &mut impl DbTxn, instruction: InInstructionWithBalance) {
+    let mut existing = Self::get(txn, instruction.balance).unwrap_or_default();
     existing.extend(instruction.encode());
-    txn.put(key, existing);
+    Self::set(txn, instruction.balance, &existing);
   }
+
   pub fn take_forwarded_output(
-    txn: &mut D::Transaction<'_>,
+    txn: &mut impl DbTxn,
     balance: Balance,
   ) -> Option<InInstructionWithBalance> {
-    let key = Self::forwarded_output_key(balance);
-
-    let outputs = txn.get(&key)?;
+    let outputs = Self::get(txn, balance)?;
     let mut outputs_ref = outputs.as_slice();
-
     let res = InInstructionWithBalance::decode(&mut outputs_ref).unwrap();
     assert!(outputs_ref.len() < outputs.len());
     if outputs_ref.is_empty() {
-      txn.del(&key);
+      txn.del(&Self::key(balance));
     } else {
-      txn.put(&key, outputs_ref);
+      Self::set(txn, balance, &outputs);
     }
     Some(res)
   }
+}
 
-  fn delayed_output_keys() -> Vec<u8> {
-    Self::multisigs_key(b"delayed_outputs", [])
-  }
-  pub fn save_delayed_output(txn: &mut D::Transaction<'_>, instruction: InInstructionWithBalance) {
-    let key = Self::delayed_output_keys();
-    let mut existing = txn.get(&key).unwrap_or(vec![]);
+impl DelayedOutputDb {
+  pub fn save_delayed_output(txn: &mut impl DbTxn, instruction: InInstructionWithBalance) {
+    let mut existing = Self::get(txn).unwrap_or_default();
     existing.extend(instruction.encode());
-    txn.put(key, existing);
+    Self::set(txn, &existing);
   }
-  pub fn take_delayed_outputs(txn: &mut D::Transaction<'_>) -> Vec<InInstructionWithBalance> {
-    let key = Self::delayed_output_keys();
 
-    let Some(outputs) = txn.get(&key) else { return vec![] };
-    txn.del(key);
+  pub fn take_delayed_outputs(txn: &mut impl DbTxn) -> Vec<InInstructionWithBalance> {
+    let Some(outputs) = Self::get(txn) else { return vec![] };
+    txn.del(Self::key());
 
     let mut outputs_ref = outputs.as_slice();
     let mut res = vec![];