Misc tidying of serai-db calls

This commit is contained in:
Luke Parker
2024-09-11 09:12:00 -04:00
parent 59fa49f750
commit 9b8c8f8231
6 changed files with 46 additions and 49 deletions

View File

@@ -47,9 +47,13 @@ macro_rules! create_db {
}) => {
$(
#[derive(Clone, Debug)]
pub(crate) struct $field_name;
impl $field_name {
pub(crate) fn key$(<$($generic_name: $generic_type),+>)?($($arg: $arg_type),*) -> Vec<u8> {
pub(crate) struct $field_name$(
<$($generic_name: $generic_type),+>
)?$(
(core::marker::PhantomData<($($generic_name),+)>)
)?;
impl$(<$($generic_name: $generic_type),+>)? $field_name$(<$($generic_name),+>)? {
pub(crate) fn key($($arg: $arg_type),*) -> Vec<u8> {
use scale::Encode;
$crate::serai_db_key(
stringify!($db_name).as_bytes(),
@@ -57,38 +61,38 @@ macro_rules! create_db {
($($arg),*).encode()
)
}
pub(crate) fn set$(<$($generic_name: $generic_type),+>)?(
pub(crate) fn set(
txn: &mut impl DbTxn
$(, $arg: $arg_type)*,
data: &$field_type
) {
let key = $field_name::key$(::<$($generic_name),+>)?($($arg),*);
let key = Self::key($($arg),*);
txn.put(&key, borsh::to_vec(data).unwrap());
}
pub(crate) fn get$(<$($generic_name: $generic_type),+>)?(
pub(crate) fn get(
getter: &impl Get,
$($arg: $arg_type),*
) -> Option<$field_type> {
getter.get($field_name::key$(::<$($generic_name),+>)?($($arg),*)).map(|data| {
getter.get(Self::key($($arg),*)).map(|data| {
borsh::from_slice(data.as_ref()).unwrap()
})
}
// Returns a PhantomData of all generic types so if the generic was only used in the value,
// not the keys, this doesn't have unused generic types
#[allow(dead_code)]
pub(crate) fn del$(<$($generic_name: $generic_type),+>)?(
pub(crate) fn del(
txn: &mut impl DbTxn
$(, $arg: $arg_type)*
) -> core::marker::PhantomData<($($($generic_name),+)?)> {
txn.del(&$field_name::key$(::<$($generic_name),+>)?($($arg),*));
txn.del(&Self::key($($arg),*));
core::marker::PhantomData
}
pub(crate) fn take$(<$($generic_name: $generic_type),+>)?(
pub(crate) fn take(
txn: &mut impl DbTxn
$(, $arg: $arg_type)*
) -> Option<$field_type> {
let key = $field_name::key$(::<$($generic_name),+>)?($($arg),*);
let key = Self::key($($arg),*);
let res = txn.get(&key).map(|data| borsh::from_slice(data.as_ref()).unwrap());
if res.is_some() {
txn.del(key);
@@ -119,14 +123,14 @@ macro_rules! db_channel {
}
}
impl $field_name {
pub(crate) fn send$(<$($generic_name: $generic_type),+>)?(
impl$(<$($generic_name: $generic_type),+>)? $field_name$(<$($generic_name),+>)? {
pub(crate) fn send(
txn: &mut impl DbTxn
$(, $arg: $arg_type)*
, value: &$field_type
) {
// Use index 0 to store the amount of messages
let messages_sent_key = $field_name::key$(::<$($generic_name),+>)?($($arg,)* 0);
let messages_sent_key = Self::key($($arg,)* 0);
let messages_sent = txn.get(&messages_sent_key).map(|counter| {
u32::from_le_bytes(counter.try_into().unwrap())
}).unwrap_or(0);
@@ -137,22 +141,22 @@ macro_rules! db_channel {
// at the same time
let index_to_use = messages_sent + 2;
$field_name::set$(::<$($generic_name),+>)?(txn, $($arg,)* index_to_use, value);
Self::set(txn, $($arg,)* index_to_use, value);
}
pub(crate) fn try_recv$(<$($generic_name: $generic_type),+>)?(
pub(crate) fn try_recv(
txn: &mut impl DbTxn
$(, $arg: $arg_type)*
) -> Option<$field_type> {
let messages_recvd_key = $field_name::key$(::<$($generic_name),+>)?($($arg,)* 1);
let messages_recvd_key = Self::key($($arg,)* 1);
let messages_recvd = txn.get(&messages_recvd_key).map(|counter| {
u32::from_le_bytes(counter.try_into().unwrap())
}).unwrap_or(0);
let index_to_read = messages_recvd + 2;
let res = $field_name::get$(::<$($generic_name),+>)?(txn, $($arg,)* index_to_read);
let res = Self::get(txn, $($arg,)* index_to_read);
if res.is_some() {
$field_name::del$(::<$($generic_name),+>)?(txn, $($arg,)* index_to_read);
Self::del(txn, $($arg,)* index_to_read);
txn.put(&messages_recvd_key, (messages_recvd + 1).to_le_bytes());
}
res