pub trait Generator {
type Item: Clone + Send;
type Config: Clone + Send + Sync;
fn new(config: Self::Config, seed: u64) -> Self
where
Self: Sized;
fn generate_one(&mut self) -> Self::Item;
fn generate_batch(&mut self, count: usize) -> Vec<Self::Item> {
(0..count).map(|_| self.generate_one()).collect()
}
fn generate_iter(&mut self, count: usize) -> GeneratorIterator<'_, Self>
where
Self: Sized,
{
GeneratorIterator {
generator: self,
remaining: count,
}
}
fn reset(&mut self);
fn count(&self) -> u64;
fn seed(&self) -> u64;
}
pub struct GeneratorIterator<'a, G: Generator> {
generator: &'a mut G,
remaining: usize,
}
impl<'a, G: Generator> Iterator for GeneratorIterator<'a, G> {
type Item = G::Item;
fn next(&mut self) -> Option<Self::Item> {
if self.remaining > 0 {
self.remaining -= 1;
Some(self.generator.generate_one())
} else {
None
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.remaining, Some(self.remaining))
}
}
impl<'a, G: Generator> ExactSizeIterator for GeneratorIterator<'a, G> {}
pub trait ParallelGenerator: Generator + Sized {
fn split(self, parts: usize) -> Vec<Self>;
fn merge_results(results: Vec<Vec<Self::Item>>) -> Vec<Self::Item> {
results.into_iter().flatten().collect()
}
}
#[derive(Debug, Clone)]
pub struct GenerationProgress {
pub total: u64,
pub completed: u64,
pub items_per_second: f64,
pub eta_seconds: Option<u64>,
pub phase: String,
}
impl GenerationProgress {
pub fn new(total: u64) -> Self {
Self {
total,
completed: 0,
items_per_second: 0.0,
eta_seconds: None,
phase: String::new(),
}
}
pub fn percentage(&self) -> f64 {
if self.total == 0 {
1.0
} else {
self.completed as f64 / self.total as f64
}
}
pub fn is_complete(&self) -> bool {
self.completed >= self.total
}
}
pub trait ProgressReporter {
fn report_progress(&self, progress: &GenerationProgress);
}
pub struct NoopProgressReporter;
impl ProgressReporter for NoopProgressReporter {
fn report_progress(&self, _progress: &GenerationProgress) {}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
struct SimpleGenerator {
seed: u64,
count: u64,
value: u64,
}
impl Generator for SimpleGenerator {
type Item = u64;
type Config = ();
fn new(_config: Self::Config, seed: u64) -> Self {
Self {
seed,
count: 0,
value: seed,
}
}
fn generate_one(&mut self) -> Self::Item {
self.count += 1;
self.value = self.value.wrapping_mul(6364136223846793005).wrapping_add(1);
self.value
}
fn reset(&mut self) {
self.count = 0;
self.value = self.seed;
}
fn count(&self) -> u64 {
self.count
}
fn seed(&self) -> u64 {
self.seed
}
}
#[test]
fn test_generator_batch() {
let mut gen = SimpleGenerator::new((), 42);
let batch = gen.generate_batch(10);
assert_eq!(batch.len(), 10);
assert_eq!(gen.count(), 10);
}
#[test]
fn test_generator_determinism() {
let mut gen1 = SimpleGenerator::new((), 42);
let mut gen2 = SimpleGenerator::new((), 42);
for _ in 0..100 {
assert_eq!(gen1.generate_one(), gen2.generate_one());
}
}
#[test]
fn test_generator_reset() {
let mut gen = SimpleGenerator::new((), 42);
let first_run: Vec<_> = gen.generate_iter(10).collect();
gen.reset();
let second_run: Vec<_> = gen.generate_iter(10).collect();
assert_eq!(first_run, second_run);
}
}