use commonware_parallel::{Sequential, Strategy};
use commonware_runtime::{Spawner, Supervisor, deterministic};
use commonware_utils::{
channel::{fallible::OneshotExt as _, oneshot},
sync::Mutex,
};
use std::{cmp::Ordering, fmt, future::Future, sync::Arc};
pub struct Gate {
pub started: oneshot::Receiver<()>,
pub release: oneshot::Sender<()>,
}
struct Pending {
started: oneshot::Sender<()>,
release: oneshot::Receiver<()>,
}
#[derive(Clone)]
pub struct GatedStrategy {
context: Arc<deterministic::Context>,
pending: Arc<Mutex<Option<Pending>>>,
}
impl fmt::Debug for GatedStrategy {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("GatedStrategy").finish_non_exhaustive()
}
}
impl GatedStrategy {
pub fn new(context: deterministic::Context) -> Self {
Self {
context: Arc::new(context),
pending: Arc::new(Mutex::new(None)),
}
}
pub fn arm(&self) -> Gate {
let (started_tx, started) = oneshot::channel();
let (release, release_rx) = oneshot::channel();
let mut pending = self.pending.lock();
assert!(pending.is_none(), "a job is already armed");
*pending = Some(Pending {
started: started_tx,
release: release_rx,
});
Gate { started, release }
}
}
impl Strategy for GatedStrategy {
fn spawn<F, T>(&self, _len: usize, f: F) -> impl Future<Output = T> + Send + 'static
where
F: FnOnce(Self) -> T + Send + 'static,
T: Send + 'static,
{
let pending = self.pending.lock().take();
let strategy = self.clone();
let (result_tx, result_rx) = oneshot::channel();
self.context
.child("job")
.shared(true)
.spawn(move |_| async move {
if let Some(pending) = pending {
pending.started.send_lossy(());
pending.release.await.expect("job release dropped");
}
let result = f(strategy);
result_tx.send_lossy(result);
});
async move { result_rx.await.expect("job result dropped") }
}
fn run<R, SEQ, PAR>(&self, len: usize, serial: SEQ, parallel: PAR) -> R
where
R: Send,
SEQ: FnOnce() -> R + Send,
PAR: FnOnce() -> R + Send,
{
Sequential.run(len, serial, parallel)
}
fn try_run<R, E, SEQ, PAR>(&self, len: usize, serial: SEQ, parallel: PAR) -> Result<R, E>
where
R: Send,
E: Send,
SEQ: FnOnce() -> Result<R, E> + Send,
PAR: FnOnce() -> Result<R, E> + Send,
{
Sequential.try_run(len, serial, parallel)
}
fn fold_init<I, INIT, T, R, ID, F, RD>(
&self,
iter: I,
init: INIT,
identity: ID,
fold_op: F,
reduce_op: RD,
) -> R
where
I: IntoIterator<IntoIter: Send, Item: Send> + Send,
INIT: Fn() -> T + Send + Sync,
T: Send,
R: Send,
ID: Fn() -> R + Send + Sync,
F: Fn(R, &mut T, I::Item) -> R + Send + Sync,
RD: Fn(R, R) -> R + Send + Sync,
{
Sequential.fold_init(iter, init, identity, fold_op, reduce_op)
}
fn try_fold<I, R, E, ID, F, RD>(
&self,
iter: I,
identity: ID,
fold_op: F,
reduce_op: RD,
) -> Result<R, E>
where
I: IntoIterator<IntoIter: Send, Item: Send> + Send,
R: Send,
E: Send,
ID: Fn() -> R + Send + Sync,
F: Fn(R, I::Item) -> Result<R, E> + Send + Sync,
RD: Fn(R, R) -> R + Send + Sync,
{
Sequential.try_fold(iter, identity, fold_op, reduce_op)
}
fn join<A, B, RA, RB>(&self, a: A, b: B) -> (RA, RB)
where
A: FnOnce() -> RA + Send,
B: FnOnce() -> RB + Send,
RA: Send,
RB: Send,
{
Sequential.join(a, b)
}
fn sort_by<T, F>(&self, items: &mut [T], compare: F)
where
T: Send,
F: Fn(&T, &T) -> Ordering + Send + Sync,
{
Sequential.sort_by(items, compare)
}
}