use crate::{
Circuit, DBData, DynZWeight, RootCircuit, Runtime, Stream, TypedBox, ZWeight,
circuit::{
LocalStoreMarker, Scope,
metadata::OperatorLocation,
operator_traits::{Operator, SourceOperator},
},
dynamic::{DowncastTrait, DynData, DynPair, DynPairs, DynUnit, Erase, LeanVec},
operator::dynamic::{
input::{
AddInputIndexedZSetFactories, AddInputMapFactories, AddInputMapWithWaterlineFactories,
AddInputZSetFactories, CollectionHandle, UpsertHandle,
},
input_upsert::DynUpdate,
},
typed_batch::{OrdIndexedZSet, OrdZSet},
utils::Tup2,
};
use std::{
borrow::{Borrow, Cow},
collections::VecDeque,
fmt::Debug,
hash::{Hash, Hasher},
marker::PhantomData,
mem::{replace, take, transmute},
ops::{Deref, Range},
panic::Location,
sync::{Arc, Mutex},
};
use typedmap::TypedMapKey;
pub use crate::operator::dynamic::input_upsert::{PatchFunc, Update};
pub type IndexedZSetStream<K, V> = Stream<RootCircuit, OrdIndexedZSet<K, V>>;
pub type ZSetStream<K> = Stream<RootCircuit, OrdZSet<K>>;
pub trait StagedBuffers: Send + Sync {
fn flush(&mut self);
}
pub struct ZSetStagedBuffers {
input_handle: InputHandle<Vec<Box<DynPairs<DynPair<DynData, DynUnit>, DynZWeight>>>>,
vals: Vec<Box<DynPairs<DynPair<DynData, DynUnit>, DynZWeight>>>,
}
impl StagedBuffers for ZSetStagedBuffers {
fn flush(&mut self) {
for (worker, vals) in self.vals.drain(..).enumerate() {
self.input_handle.update_for_worker(worker, |tuples| {
tuples.push(vals);
});
}
}
}
#[repr(transparent)]
pub struct ZSetHandle<K> {
handle: CollectionHandle<DynPair<DynData, DynUnit>, DynZWeight>,
phantom: PhantomData<fn(&K)>,
}
impl<K> Clone for ZSetHandle<K> {
fn clone(&self) -> Self {
Self {
handle: self.handle.clone(),
phantom: PhantomData,
}
}
}
impl<K> Deref for ZSetHandle<K> {
type Target = CollectionHandle<DynPair<DynData, DynUnit>, DynZWeight>;
fn deref(&self) -> &Self::Target {
&self.handle
}
}
impl<K> ZSetHandle<K>
where
K: DBData,
{
fn new(handle: CollectionHandle<DynPair<DynData, DynUnit>, DynZWeight>) -> Self {
Self {
handle,
phantom: PhantomData,
}
}
pub fn push(&self, k: K, mut w: ZWeight) {
self.handle.dyn_push(Tup2(k, ()).erase_mut(), w.erase_mut())
}
pub fn append(&self, vals: &mut Vec<Tup2<K, ZWeight>>) {
let vals: &mut Vec<Tup2<Tup2<K, ()>, ZWeight>> = unsafe { transmute(vals) };
let vals = Box::new(LeanVec::from(take(vals)));
self.handle.dyn_append(&mut vals.erase_box())
}
pub fn stage(
&self,
buffers: impl IntoIterator<Item = VecDeque<Tup2<K, ZWeight>>>,
) -> ZSetStagedBuffers {
let num_partitions = self.handle.num_partitions();
let mut partitions = vec![self.handle.pairs_factory.default_box(); num_partitions];
let mut next_worker = 0;
for vals in buffers {
let mut vec = Vec::from(vals);
let vals: &mut Vec<Tup2<Tup2<K, ()>, ZWeight>> = unsafe { transmute(&mut vec) };
let vals = Box::new(LeanVec::from(take(vals)));
self.handle
.dyn_stage(&mut vals.erase_box(), &mut next_worker, &mut partitions);
}
ZSetStagedBuffers {
input_handle: self.handle.input_handle.clone(),
vals: partitions,
}
}
}
pub struct IndexedZSetStagedBuffers {
input_handle: InputHandle<Vec<Box<DynPairs<DynData, DynPair<DynData, DynZWeight>>>>>,
vals: Vec<Box<DynPairs<DynData, DynPair<DynData, DynZWeight>>>>,
}
impl StagedBuffers for IndexedZSetStagedBuffers {
fn flush(&mut self) {
for (worker, vals) in self.vals.drain(..).enumerate() {
self.input_handle.update_for_worker(worker, |tuples| {
tuples.push(vals);
});
}
}
}
#[derive(Clone)]
#[repr(transparent)]
pub struct IndexedZSetHandle<K, V> {
handle: CollectionHandle<DynData, DynPair<DynData, DynZWeight>>,
phantom: PhantomData<fn(&K, &V)>,
}
impl<K, V> IndexedZSetHandle<K, V>
where
K: DBData,
V: DBData,
{
fn new(handle: CollectionHandle<DynData, DynPair<DynData, DynZWeight>>) -> Self {
Self {
handle,
phantom: PhantomData,
}
}
pub fn push(&self, mut k: K, (v, w): (V, ZWeight)) {
self.handle.dyn_push(k.erase_mut(), Tup2(v, w).erase_mut())
}
pub fn append(&self, vals: &mut Vec<Tup2<K, Tup2<V, ZWeight>>>) {
let vals = Box::new(LeanVec::from(take(vals)));
self.handle.dyn_append(&mut vals.erase_box())
}
}
pub struct MapStagedBuffers {
input_handle: InputHandle<Vec<Box<DynPairs<DynData, DynUpdate<DynData, DynData>>>>>,
vals: Vec<Box<DynPairs<DynData, DynUpdate<DynData, DynData>>>>,
}
impl StagedBuffers for MapStagedBuffers {
fn flush(&mut self) {
for (worker, vals) in self.vals.drain(..).enumerate() {
self.input_handle.update_for_worker(worker, |tuples| {
tuples.push(vals);
});
}
}
}
#[repr(transparent)]
pub struct MapHandle<K, V, U> {
handle: UpsertHandle<DynData, DynUpdate<DynData, DynData>>,
phantom: PhantomData<fn(&K, &V, &U)>,
}
impl<K, V, U> Clone for MapHandle<K, V, U> {
fn clone(&self) -> Self {
Self {
handle: self.handle.clone(),
phantom: PhantomData,
}
}
}
impl<K, V, U> MapHandle<K, V, U>
where
K: DBData,
V: DBData,
U: DBData,
{
fn new(handle: UpsertHandle<DynData, DynUpdate<DynData, DynData>>) -> Self {
Self {
handle,
phantom: PhantomData,
}
}
pub fn push(&self, mut k: K, mut upd: Update<V, U>) {
self.handle.dyn_push(k.erase_mut(), upd.erase_mut())
}
pub fn append(&mut self, vals: &mut Vec<Tup2<K, Update<V, U>>>) {
let vals = Box::new(LeanVec::from(take(vals)));
self.handle.dyn_append(&mut vals.erase_box())
}
pub fn stage(
&self,
buffers: impl IntoIterator<Item = VecDeque<Tup2<K, Update<V, U>>>>,
) -> MapStagedBuffers {
let num_partitions = self.handle.num_partitions();
let mut partitions = vec![self.handle.pairs_factory.default_box(); num_partitions];
for vals in buffers {
let vec = Vec::from(vals);
let vals = Box::new(LeanVec::from(vec));
self.handle
.dyn_stage(&mut vals.erase_box(), &mut partitions);
}
MapStagedBuffers {
input_handle: self.handle.input_handle.clone(),
vals: partitions,
}
}
}
impl RootCircuit {
#[track_caller]
pub fn add_input_stream<T>(&self) -> (Stream<Self, T>, InputHandle<T>)
where
T: Default + Debug + Clone + Send + 'static,
{
let (input, input_handle) =
Input::new(Location::caller(), |x| x, Arc::new(|| Default::default()));
let stream = self.add_source(input);
(stream, input_handle)
}
#[track_caller]
pub fn add_input_zset<K>(&self) -> (Stream<RootCircuit, OrdZSet<K>>, ZSetHandle<K>)
where
K: DBData,
{
let factories = AddInputZSetFactories::new::<K>();
let (stream, handle) = self.dyn_add_input_zset_mono(&factories);
(stream.typed(), ZSetHandle::new(handle))
}
#[allow(clippy::type_complexity)]
#[track_caller]
pub fn add_input_indexed_zset<K, V>(
&self,
) -> (
Stream<RootCircuit, OrdIndexedZSet<K, V>>,
IndexedZSetHandle<K, V>,
)
where
K: DBData,
V: DBData,
{
let factories = AddInputIndexedZSetFactories::new::<K, V>();
let (stream, handle) = self.dyn_add_input_indexed_zset_mono(&factories);
(stream.typed(), IndexedZSetHandle::new(handle))
}
#[track_caller]
pub fn add_input_map<K, V, U, PF>(
&self,
patch_func: PF,
) -> (
Stream<RootCircuit, OrdIndexedZSet<K, V>>,
MapHandle<K, V, U>,
)
where
K: DBData,
V: DBData,
U: DBData + Erase<DynData>,
PF: Fn(&mut V, &U) + 'static,
{
self.add_input_map_persistent(None, patch_func)
}
#[track_caller]
pub fn add_input_map_persistent<K, V, U, PF>(
&self,
persistent_id: Option<&str>,
patch_func: PF,
) -> (
Stream<RootCircuit, OrdIndexedZSet<K, V>>,
MapHandle<K, V, U>,
)
where
K: DBData,
V: DBData,
U: DBData + Erase<DynData>,
PF: Fn(&mut V, &U) + 'static,
{
let factories = AddInputMapFactories::new::<K, V, U>();
let (stream, handle) = self.dyn_add_input_map_mono(
persistent_id,
&factories,
Box::new(move |v: &mut DynData, u: &DynData| unsafe {
patch_func(v.downcast_mut::<V>(), u.downcast::<U>())
}),
);
(stream.typed(), MapHandle::new(handle))
}
#[track_caller]
pub fn add_input_map_with_waterline<K, V, U, W, E, PF, IF, WF, LB, FF, RF>(
&self,
patch_func: PF,
init: IF,
extract_ts: WF,
least_upper_bound: LB,
filter_func: FF,
report_func: RF,
) -> (
Stream<RootCircuit, OrdIndexedZSet<K, V>>,
Stream<RootCircuit, OrdZSet<E>>,
Stream<RootCircuit, TypedBox<W, DynData>>,
MapHandle<K, V, U>,
)
where
K: DBData,
V: DBData,
U: DBData + Erase<DynData>,
W: DBData,
E: DBData,
PF: Fn(&mut V, &U) + 'static,
IF: Fn() -> W + 'static,
WF: Fn(&K, &V) -> W + 'static,
LB: Fn(&W, &W) -> W + Clone + 'static,
FF: Fn(&W, &K, &V) -> bool + 'static,
RF: Fn(&W, &K, &V, ZWeight) -> E + 'static,
{
self.add_input_map_with_waterline_persistent(
None,
patch_func,
init,
extract_ts,
least_upper_bound,
filter_func,
report_func,
)
}
#[allow(clippy::too_many_arguments)]
#[track_caller]
pub fn add_input_map_with_waterline_persistent<K, V, U, W, E, PF, IF, WF, LB, FF, RF>(
&self,
persistent_id: Option<&str>,
patch_func: PF,
init: IF,
extract_ts: WF,
least_upper_bound: LB,
filter_func: FF,
report_func: RF,
) -> (
Stream<RootCircuit, OrdIndexedZSet<K, V>>,
Stream<RootCircuit, OrdZSet<E>>,
Stream<RootCircuit, TypedBox<W, DynData>>,
MapHandle<K, V, U>,
)
where
K: DBData,
V: DBData,
U: DBData + Erase<DynData>,
W: DBData + Erase<DynData>,
E: DBData + Erase<DynData>,
PF: Fn(&mut V, &U) + 'static,
IF: Fn() -> W + 'static,
WF: Fn(&K, &V) -> W + 'static,
LB: Fn(&W, &W) -> W + Clone + 'static,
FF: Fn(&W, &K, &V) -> bool + 'static,
RF: Fn(&W, &K, &V, ZWeight) -> E + 'static,
{
let factories = AddInputMapWithWaterlineFactories::new::<K, V, U, E>();
let (stream, errors, waterline, handle) = self.dyn_add_input_map_with_waterline_mono(
persistent_id,
&factories,
Box::new(move |v: &mut DynData, u: &DynData| unsafe {
patch_func(v.downcast_mut::<V>(), u.downcast::<U>())
}),
Box::new(move || Box::new(init())),
Box::new(move |k: &DynData, v: &DynData, ts: &mut DynData| {
let k = unsafe { k.downcast::<K>() };
let v = unsafe { v.downcast::<V>() };
let w = unsafe { ts.downcast_mut::<W>() };
*w = extract_ts(k, v);
}),
Box::new(move |a: &DynData, b: &DynData, ts: &mut DynData| {
let a = unsafe { a.downcast::<W>() };
let b = unsafe { b.downcast::<W>() };
let ts = unsafe { ts.downcast_mut::<W>() };
*ts = least_upper_bound(a, b)
}),
Box::new(move |wl: &DynData, k: &DynData, v: &DynData| {
let wl = unsafe { wl.downcast::<W>() };
let k = unsafe { k.downcast::<K>() };
let v = unsafe { v.downcast::<V>() };
filter_func(wl, k, v)
}),
Box::new(
move |wl: &DynData, k: &DynData, v: &DynData, w: ZWeight, err: &mut DynData| {
let wl = unsafe { wl.downcast::<W>() };
let k = unsafe { k.downcast::<K>() };
let v = unsafe { v.downcast::<V>() };
let err = unsafe { err.downcast_mut::<E>() };
*err = report_func(wl, k, v, w);
},
),
);
(
stream.typed(),
errors.typed(),
unsafe { waterline.typed_data() },
MapHandle::new(handle),
)
}
}
struct InputId<T> {
id: usize,
_marker: PhantomData<T>,
}
unsafe impl<T> Sync for InputId<T> {}
impl<T> Hash for InputId<T> {
fn hash<H>(&self, state: &mut H)
where
H: Hasher,
{
self.id.hash(state);
}
}
impl<T> PartialEq for InputId<T> {
fn eq(&self, other: &Self) -> bool {
self.id == other.id
}
}
impl<T> Eq for InputId<T> {}
impl<T> InputId<T> {
fn new(id: usize) -> Self {
Self {
id,
_marker: PhantomData,
}
}
}
impl<T> TypedMapKey<LocalStoreMarker> for InputId<T>
where
T: 'static,
{
type Value = InputHandle<T>;
}
#[derive(Clone)]
pub(crate) struct Mailbox<T> {
empty: Arc<dyn Fn() -> T + Send + Sync>,
value: Arc<Mutex<T>>,
}
impl<T: Clone> Mailbox<T> {
pub(in crate::operator) fn new(empty: Arc<dyn Fn() -> T + Send + Sync>) -> Self {
let v = empty();
Self {
empty,
value: Arc::new(Mutex::new(v)),
}
}
pub(in crate::operator) fn take(&self) -> T {
replace(&mut *self.value.lock().unwrap(), (self.empty)())
}
pub(super) fn map<F, O: 'static>(&self, func: F) -> O
where
F: Fn(&T) -> O,
{
func(self.value.lock().unwrap().borrow())
}
fn update<F>(&self, f: F)
where
F: FnOnce(&mut T),
{
f(&mut *self.value.lock().unwrap());
}
pub(in crate::operator) fn set(&self, v: T) {
*self.value.lock().unwrap() = v;
}
pub(in crate::operator) fn clear(&self) {
*self.value.lock().unwrap() = (self.empty)();
}
}
pub(crate) struct InputHandleInternal<T> {
pub(crate) mailbox: Vec<Mailbox<T>>,
}
impl<T> InputHandleInternal<T>
where
T: Clone,
{
fn new(workers: Range<usize>, empty_val: Arc<dyn Fn() -> T + Send + Sync>) -> Self {
assert!(!workers.is_empty());
Self {
mailbox: workers
.clone()
.map(move |_| Mailbox::new(empty_val.clone()))
.collect(),
}
}
fn set_for_worker(&self, worker: usize, v: T) {
self.mailbox(worker).set(v);
}
fn update_for_worker<F>(&self, worker: usize, f: F)
where
F: FnOnce(&mut T),
{
self.mailbox(worker).update(f);
}
fn set_for_all(&self, v: T) {
for i in 0..self.mailbox.len() - 1 {
self.mailbox[i].set(v.clone());
}
self.mailbox[self.mailbox.len() - 1].set(v);
}
fn clear_for_all(&self) {
for mailbox in self.mailbox.iter() {
mailbox.clear();
}
}
fn mailbox(&self, worker: usize) -> &Mailbox<T> {
&self.mailbox[worker]
}
}
pub struct InputHandle<T>(pub(crate) Arc<InputHandleInternal<T>>);
impl<T> Clone for InputHandle<T> {
fn clone(&self) -> Self {
Self(self.0.clone())
}
}
impl<T> InputHandle<T>
where
T: Send + Clone + 'static,
{
fn new(empty_val: Arc<dyn Fn() -> T + Send + Sync>) -> Self {
match Runtime::runtime() {
None => Self(Arc::new(InputHandleInternal::new(0..1, empty_val))),
Some(runtime) => {
let input_id = runtime.sequence_next();
runtime
.local_store()
.entry(InputId::new(input_id))
.or_insert_with(|| {
Self(Arc::new(InputHandleInternal::new(
runtime.layout().local_workers(),
empty_val,
)))
})
.value()
.clone()
}
}
}
fn mailbox(&self, worker: usize) -> &Mailbox<T> {
self.0.mailbox(worker)
}
pub fn set_for_worker(&self, worker: usize, v: T) {
self.0.set_for_worker(worker, v);
}
pub fn update_for_worker<F>(&self, worker: usize, f: F)
where
F: FnOnce(&mut T),
{
self.0.update_for_worker(worker, f);
}
pub fn set_for_all(&self, v: T) {
self.0.set_for_all(v);
}
pub fn clear_for_all(&self) {
self.0.clear_for_all();
}
}
pub struct Input<IT, OT, F> {
location: &'static Location<'static>,
mailbox: Mailbox<IT>,
input_func: F,
phantom: PhantomData<OT>,
}
impl<IT, OT, F> Input<IT, OT, F>
where
IT: Clone + Send + 'static,
{
pub fn new(
location: &'static Location<'static>,
input_func: F,
default: Arc<dyn Fn() -> IT + Send + Sync>,
) -> (Self, InputHandle<IT>) {
let handle = InputHandle::new(default);
let mailbox = handle.mailbox(Runtime::local_worker_offset()).clone();
let input = Self {
location,
mailbox,
input_func,
phantom: PhantomData,
};
(input, handle)
}
}
impl<IT, OT, F> Operator for Input<IT, OT, F>
where
IT: 'static,
OT: 'static,
F: 'static,
{
fn name(&self) -> Cow<'static, str> {
Cow::from("Input")
}
fn is_input(&self) -> bool {
true
}
fn location(&self) -> OperatorLocation {
Some(self.location)
}
fn fixedpoint(&self, _scope: Scope) -> bool {
false
}
}
impl<IT, OT, F> SourceOperator<OT> for Input<IT, OT, F>
where
IT: Clone + Debug + 'static,
OT: 'static,
F: Fn(IT) -> OT + 'static,
{
async fn eval(&mut self) -> OT {
let v = self.mailbox.take();
(self.input_func)(v)
}
}