use std::hash::Hash;
use std::marker::PhantomData;
use nanorand::{tls_rng, Rng};
use crate::operator::{ExchangeData, KeyerFn};
use super::group_by_hash;
pub(crate) enum NextStrategy<Out, IndexFn = fn(&Out) -> u64>
where
IndexFn: KeyerFn<u64, Out>,
{
OnlyOne,
Random,
GroupBy(IndexFn, PhantomData<Out>),
All,
}
impl<Out, IndexFn> std::fmt::Debug for NextStrategy<Out, IndexFn>
where
IndexFn: KeyerFn<u64, Out>,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::OnlyOne => write!(f, "OnlyOne"),
Self::Random => write!(f, "Random"),
Self::GroupBy(_, _) => write!(f, "GroupBy"),
Self::All => write!(f, "All"),
}
}
}
impl<Out, IndexFn: Clone> Clone for NextStrategy<Out, IndexFn>
where
IndexFn: KeyerFn<u64, Out>,
{
fn clone(&self) -> Self {
match self {
Self::OnlyOne => Self::OnlyOne,
Self::Random => Self::Random,
Self::GroupBy(idx, _) => Self::GroupBy(idx.clone(), PhantomData),
Self::All => Self::All,
}
}
}
impl<Out: ExchangeData> NextStrategy<Out> {
pub(crate) fn group_by<Key: Hash, Keyer>(
keyer: Keyer,
) -> NextStrategy<Out, impl KeyerFn<u64, Out>>
where
Keyer: KeyerFn<Key, Out>,
{
NextStrategy::GroupBy(
move |item: &Out| group_by_hash(&keyer(item)),
Default::default(),
)
}
pub(crate) fn all() -> NextStrategy<Out> {
NextStrategy::All
}
pub(crate) fn only_one() -> NextStrategy<Out> {
NextStrategy::OnlyOne
}
pub(crate) fn random() -> NextStrategy<Out> {
NextStrategy::Random
}
}
impl<Out: ExchangeData, IndexFn> NextStrategy<Out, IndexFn>
where
IndexFn: KeyerFn<u64, Out>,
{
pub fn index(&self, message: &Out) -> usize {
match self {
NextStrategy::OnlyOne | NextStrategy::All => 0,
NextStrategy::Random => tls_rng().generate(),
NextStrategy::GroupBy(keyer, _) => keyer(message) as usize,
}
}
}