use crate::shared::new_rng;
use log::warn;
use rand::prelude::*;
use sc_cli::{Error, Result};
pub(crate) type EmptyStorage = Error;
pub(crate) fn select_entries<E, I1, I2, F1, F2, FKey>(
keys_limit: Option<usize>,
random_seed: Option<u64>,
iter_from_first_key: F1,
iter_from_start: F2,
key_of: FKey,
) -> Result<(Vec<E>, impl Rng)>
where
I1: Iterator<Item = E>,
I2: Iterator<Item = E>,
F1: FnOnce(Option<&[u8]>) -> Result<I1>,
F2: FnOnce() -> Result<I2>,
FKey: Fn(&E) -> &[u8],
{
let mut entries: Vec<E> = if let Some(limit) = keys_limit {
let first_key =
random_seed.map(|seed| sp_core::blake2_256(&seed.to_be_bytes()[..]).to_vec());
let from_first = iter_from_first_key(first_key.as_deref())?;
let mut collected: Vec<E> = from_first.take(limit).collect();
if collected.len() < limit {
let need_more = limit - collected.len();
if let Some(ref fk) = first_key {
let extra: Vec<E> = iter_from_start()?
.take_while(|e| key_of(e) < fk.as_slice())
.take(need_more)
.collect();
collected.extend(extra);
}
if collected.len() < limit {
warn!("Only {} entries available (requested {})", collected.len(), limit);
}
}
collected
} else {
iter_from_start()?.collect()
};
if entries.is_empty() {
return Err(EmptyStorage::Input("Can't process benchmarking with empty storage".into()));
}
let (mut rng, _) = new_rng(random_seed);
entries.shuffle(&mut rng);
Ok((entries, rng))
}