use crate::{
indexer::Indexer, key::TxKey, lock_guards::LockGuards, lock_policies::lock_policy::LockPolicy,
multi_shard_ops::MultiShardOps, new_types::BitMask, prepared::schema::TxKeySelector,
shard_ops::ShardOps,
};
use std::{
hash::{BuildHasher, Hash},
ops::{Deref, DerefMut},
};
#[allow(clippy::type_complexity)]
pub(crate) enum PreparedOp<'tx, K, V, KEYS, PARAMS, STATE>
where
K: Clone + Hash + Eq,
STATE: Default,
{
Get {
key_selector: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
get: Box<dyn Fn(&K, Option<&V>, &PARAMS, &mut STATE) + 'tx>,
},
InsertWith {
key_selector: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
value_generator: Box<dyn Fn(&K, &PARAMS, &mut STATE) -> V + 'tx>,
},
InsertWithIfAbsent {
key_selector: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
value_generator: Box<dyn Fn(&K, &PARAMS, &mut STATE) -> V + 'tx>,
},
Modify {
key_selector: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
mutate: Box<dyn Fn(&K, &mut V, &PARAMS, &mut STATE) + 'tx>,
},
MoveValue {
key_selector_from: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
key_selector_to: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
},
Remove {
key_selector: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
},
RemoveIf {
key_selector: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
condition: Box<dyn Fn(&K, &V, &PARAMS, &mut STATE) -> bool + 'tx>,
},
SwapValue {
key_selector_a: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
key_selector_b: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
},
Update {
key_selector: Box<dyn TxKeySelector<TxKey<K>, KEYS> + 'tx>,
transform: Box<dyn Fn(&K, Option<&V>, &PARAMS, &mut STATE) -> Option<V> + 'tx>,
},
}
impl<'tx, K, V, KEYS, PARAMS, STATE> PreparedOp<'tx, K, V, KEYS, PARAMS, STATE>
where
K: Clone + Hash + Eq,
STATE: Default,
{
pub fn read_write_bitmasks(&self, keys: &KEYS) -> (BitMask, BitMask) {
match self {
Self::Get { key_selector, .. } => {
(key_selector.get(keys).shard_index.bitmask(), BitMask::ZERO)
}
Self::InsertWith { key_selector, .. }
| Self::InsertWithIfAbsent { key_selector, .. }
| Self::Modify { key_selector, .. }
| Self::Remove { key_selector, .. }
| Self::RemoveIf { key_selector, .. }
| Self::Update { key_selector, .. } => {
(BitMask::ZERO, key_selector.get(keys).shard_index.bitmask())
}
Self::MoveValue {
key_selector_from,
key_selector_to,
..
} => (
BitMask::ZERO,
key_selector_from.get(keys).shard_index.bitmask()
| key_selector_to.get(keys).shard_index.bitmask(),
),
Self::SwapValue {
key_selector_a,
key_selector_b,
..
} => (
BitMask::ZERO,
key_selector_a.get(keys).shard_index.bitmask()
| key_selector_b.get(keys).shard_index.bitmask(),
),
}
}
pub fn apply<L, S>(
&self,
lock_guards: &mut LockGuards<'_, K, V, L>,
keys: &KEYS,
params: &PARAMS,
indexer: &Indexer<S>,
state: &mut STATE,
) where
L: LockPolicy,
S: BuildHasher,
{
match self {
Self::Get { key_selector, get } => {
let key = key_selector.get(keys);
let shard =
if (key.shard_index.bitmask() & lock_guards.write_bitmask) != BitMask::ZERO {
lock_guards.write_guard(key).deref_mut()
} else {
lock_guards.read_guard(key).deref()
};
let value_ref = ShardOps::value_ref(shard, key);
(get)(&key.key, value_ref, params, state)
}
Self::InsertWith {
key_selector,
value_generator,
} => {
let key = key_selector.get(keys);
let new_value = (value_generator)(&key.key, params, state);
let write_guard = lock_guards.write_guard(key);
ShardOps::insert::<K, V, S>(write_guard, key, new_value, indexer);
}
Self::InsertWithIfAbsent {
key_selector,
value_generator,
} => {
let key = key_selector.get(keys);
let write_guard = lock_guards.write_guard(key);
ShardOps::insert_if_absent::<K, V, S>(
write_guard,
key,
|| (value_generator)(&key.key, params, state),
indexer,
);
}
Self::Modify {
key_selector,
mutate,
} => {
let key = key_selector.get(keys);
let shard = lock_guards.write_guard(key);
ShardOps::modify(shard, key, |k, v| mutate(k, v, params, state));
}
Self::MoveValue {
key_selector_from,
key_selector_to,
} => {
let key_from = key_selector_from.get(keys);
let key_to = key_selector_to.get(keys);
MultiShardOps::move_value::<K, V, L, S>(
&mut lock_guards.write,
key_from,
key_to,
indexer,
);
}
Self::Remove { key_selector } => {
let key = key_selector.get(keys);
let shard = lock_guards.write_guard(key);
ShardOps::remove_entry::<K, V>(shard, key);
}
Self::RemoveIf {
key_selector,
condition,
} => {
let key = key_selector.get(keys);
let shard = lock_guards.write_guard(key);
ShardOps::remove_if(shard, key, |k, v| condition(k, v, params, state), indexer);
}
Self::SwapValue {
key_selector_a,
key_selector_b,
} => {
let key_a = key_selector_a.get(keys);
let key_b = key_selector_b.get(keys);
MultiShardOps::swap_value::<K, V, L, S>(
&mut lock_guards.write,
key_a,
key_b,
indexer,
);
}
Self::Update {
key_selector,
transform,
} => {
let key = key_selector.get(keys);
let shard = lock_guards.write_guard(key);
ShardOps::update(
shard,
key,
|k, v_opt| transform(k, v_opt, params, state),
indexer,
);
}
}
}
}