odem_rs_util/random/

rayon_support.rs

//! Contains the `IndexedParallelIterator` implementation for `RngStream`.
//!
//! This module provides [`RngStreamParIter`], which is created by the
//! [`RngStream::par_iter`] method.

use super::{RngGen, RngStream};
use rayon::iter::{
	IndexedParallelIterator, IntoParallelIterator, ParallelIterator,
	plumbing::{Consumer, Producer, ProducerCallback, UnindexedConsumer, bridge},
};

/// An infinite parallel iterator that yields independent `RngStream`s.
///
/// This iterator is the primary mechanism for running multiple simulations in
/// parallel. It is an `IndexedParallelIterator`, which allows for more
/// efficient work splitting by Rayon and enables adapters like `zip` and
/// `enumerate`.
///
/// It works by maintaining a single `RngStream` and dynamically splitting its
/// state as Rayon requests chunks of work, ensuring perfect independence and
/// reproducibility with minimal overhead.
///
/// It is created by calling [`RngStream::par_iter`] and should be chained with
/// [`take(n)`](IndexedParallelIterator::take) to limit the number of runs.
///
/// # Usage
///
/// Use it during the set-up of a simulation run to generate a sequence of
/// independent `RngStream`s to pass into the simulation to repeat a set of
/// experiments in parallel.
///
/// ```no_run
/// # use rand::Rng;
/// # use odem_rs_util::random::{DefaultRng, RngStream};
/// # fn run_sim(s: RngStream) -> usize { 0 }
/// use rayon::prelude::*;
/// let sum: usize = RngStream::<DefaultRng>::par_iter()
///     .take(100)
///     .map(|rng_stream| run_sim(rng_stream))
///     .sum();
/// ```
pub struct RngStreamParIter<R: RngGen> {
	stream: RngStream<R>,
	/// The number of streams this iterator will produce.
	/// Starts at `usize::MAX` to be effectively infinite until limited by
	/// `take()`.
	reruns: usize,
}

impl<R: RngGen> RngStreamParIter<R> {
	/// Creates a new, effectively infinite parallel iterator.
	pub(super) fn new(stream: RngStream<R>) -> Self {
		Self {
			stream,
			reruns: usize::MAX,
		}
	}
}

impl<R> ParallelIterator for RngStreamParIter<R>
where
	R: RngGen + Send,
	R::State: Send,
{
	type Item = RngStream<R>;

	fn drive_unindexed<C>(self, consumer: C) -> C::Result
	where
		C: UnindexedConsumer<Self::Item>,
	{
		bridge(self, consumer)
	}

	fn opt_len(&self) -> Option<usize> {
		Some(self.reruns)
	}
}

impl<R> IndexedParallelIterator for RngStreamParIter<R>
where
	R: RngGen + Send,
	R::State: Send,
{
	fn len(&self) -> usize {
		self.reruns
	}

	fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
		bridge(self, consumer)
	}

	fn with_producer<CB>(self, callback: CB) -> CB::Output
	where
		CB: ProducerCallback<Self::Item>,
	{
		let producer = StreamProducer(self);
		callback.callback(producer)
	}
}

impl<R> IntoParallelIterator for RngStream<R>
where
	R: RngGen + Send,
	R::State: Send,
{
	type Iter = RngStreamParIter<R>;
	type Item = RngStream<R>;

	fn into_par_iter(self) -> Self::Iter {
		RngStreamParIter::new(self)
	}
}

impl<R: RngGen> Clone for RngStreamParIter<R>
where
	R::State: Clone,
{
	fn clone(&self) -> Self {
		Self {
			stream: self.stream.clone(),
			reruns: self.reruns,
		}
	}
}

/// The stateful producer for our `RngStream` indexed parallel iterator.
///
/// This producer implements Rayon's split-on-demand logic. When `split_at`
/// is called, it uses `RngGen::split_at` to carve off a new, independent
/// `RngStream` for the left-hand producer, while the right-hand producer
/// continues with the advanced state.
struct StreamProducer<R: RngGen>(RngStreamParIter<R>);

impl<R: RngGen + Send> Producer for StreamProducer<R>
where
	R::State: Send,
{
	type Item = RngStream<R>;
	type IntoIter = StreamIter<R>;

	fn into_iter(self) -> Self::IntoIter {
		StreamIter(self.0)
	}

	fn split_at(mut self, index: usize) -> (Self, Self) {
		// Carve out a new stream for the left side of the split.
		let left_stream_state = R::split_at(self.0.stream.0.get_mut(), index);

		let left_iter = RngStreamParIter {
			stream: RngStream::new(left_stream_state),
			reruns: index,
		};

		// The right side (self) keeps the remaining runs.
		self.0.reruns -= index;

		(Self(left_iter), self)
	}
}

/// The concrete `Iterator` created by the `StreamProducer` for sequential
/// consumption.
struct StreamIter<R: RngGen>(RngStreamParIter<R>);

impl<R: RngGen> Iterator for StreamIter<R> {
	type Item = RngStream<R>;

	fn next(&mut self) -> Option<Self::Item> {
		if self.0.reruns > 0 {
			self.0.reruns -= 1;

			// Produce the new stream by splitting off from the iterator state.
			Some(RngStream::new(R::split_at(self.0.stream.0.get_mut(), 1)))
		} else {
			None
		}
	}

	fn size_hint(&self) -> (usize, Option<usize>) {
		(self.0.reruns, Some(self.0.reruns))
	}

	fn nth(&mut self, n: usize) -> Option<Self::Item> {
		if self.0.reruns > n {
			self.0.reruns -= n + 1;

			// Produce the new stream by splitting off from the iterator state.
			let _left_state = R::split_at(self.0.stream.0.get_mut(), n);
			Some(RngStream::new(R::split_at(self.0.stream.0.get_mut(), 1)))
		} else {
			self.0.reruns = 0;
			None
		}
	}
}

impl<R: RngGen> DoubleEndedIterator for StreamIter<R> {
	fn next_back(&mut self) -> Option<Self::Item> {
		use core::mem::replace;
		if self.0.reruns == 0 {
			return None;
		}

		// To get the *last* item, we must split off the first `n-1` items,
		// leaving the original stream in the desired final state.
		let new_state = self.0.stream.0.get_mut();
		let old_state = R::split_at(new_state, self.0.reruns - 1);

		// The state we want is what `new_state` has become. We must swap it
		// with the `old_state` we just carved off to restore the parent stream.
		let result_state = replace(new_state, old_state);

		self.0.reruns -= 1;
		Some(RngStream::new(result_state))
	}
}

impl<R: RngGen> ExactSizeIterator for StreamIter<R> {}