use std::rc::Rc;
use std::cell::RefCell;
use crate::progress::Antichain;
use crate::progress::Timestamp;
use crate::progress::ChangeBatch;
use crate::progress::frontier::MutableAntichain;
use crate::dataflow::channels::pullers::Counter as PullCounter;
use crate::dataflow::channels::pushers::CounterCore as PushCounter;
use crate::dataflow::channels::pushers::buffer::{BufferCore, Session};
use crate::dataflow::channels::BundleCore;
use crate::communication::{Push, Pull, message::RefOrMut};
use crate::Container;
use crate::logging::TimelyLogger as Logger;
use crate::dataflow::operators::InputCapability;
use crate::dataflow::operators::capability::CapabilityTrait;
pub struct InputHandleCore<T: Timestamp, D: Container, P: Pull<BundleCore<T, D>>> {
pull_counter: PullCounter<T, D, P>,
internal: Rc<RefCell<Vec<Rc<RefCell<ChangeBatch<T>>>>>>,
summaries: Rc<RefCell<Vec<Antichain<T::Summary>>>>,
logging: Option<Logger>,
}
pub type InputHandle<T, D, P> = InputHandleCore<T, Vec<D>, P>;
pub struct FrontieredInputHandleCore<'a, T: Timestamp, D: Container+'a, P: Pull<BundleCore<T, D>>+'a> {
pub handle: &'a mut InputHandleCore<T, D, P>,
pub frontier: &'a MutableAntichain<T>,
}
pub type FrontieredInputHandle<'a, T, D, P> = FrontieredInputHandleCore<'a, T, Vec<D>, P>;
impl<'a, T: Timestamp, D: Container, P: Pull<BundleCore<T, D>>> InputHandleCore<T, D, P> {
#[inline]
pub fn next(&mut self) -> Option<(InputCapability<T>, RefOrMut<D>)> {
let internal = &self.internal;
let summaries = &self.summaries;
self.pull_counter.next_guarded().map(|(guard, bundle)| {
match bundle.as_ref_or_mut() {
RefOrMut::Ref(bundle) => {
(InputCapability::new(internal.clone(), summaries.clone(), guard), RefOrMut::Ref(&bundle.data))
},
RefOrMut::Mut(bundle) => {
(InputCapability::new(internal.clone(), summaries.clone(), guard), RefOrMut::Mut(&mut bundle.data))
},
}
})
}
#[inline]
pub fn for_each<F: FnMut(InputCapability<T>, RefOrMut<D>)>(&mut self, mut logic: F) {
let mut logging = self.logging.take();
while let Some((cap, data)) = self.next() {
logging.as_mut().map(|l| l.log(crate::logging::GuardedMessageEvent { is_start: true }));
logic(cap, data);
logging.as_mut().map(|l| l.log(crate::logging::GuardedMessageEvent { is_start: false }));
}
self.logging = logging;
}
}
impl<'a, T: Timestamp, D: Container, P: Pull<BundleCore<T, D>>+'a> FrontieredInputHandleCore<'a, T, D, P> {
pub fn new(handle: &'a mut InputHandleCore<T, D, P>, frontier: &'a MutableAntichain<T>) -> Self {
FrontieredInputHandleCore {
handle,
frontier,
}
}
#[inline]
pub fn next(&mut self) -> Option<(InputCapability<T>, RefOrMut<D>)> {
self.handle.next()
}
#[inline]
pub fn for_each<F: FnMut(InputCapability<T>, RefOrMut<D>)>(&mut self, logic: F) {
self.handle.for_each(logic)
}
#[inline]
pub fn frontier(&self) -> &'a MutableAntichain<T> {
self.frontier
}
}
pub fn _access_pull_counter<T: Timestamp, D: Container, P: Pull<BundleCore<T, D>>>(input: &mut InputHandleCore<T, D, P>) -> &mut PullCounter<T, D, P> {
&mut input.pull_counter
}
pub fn new_input_handle<T: Timestamp, D: Container, P: Pull<BundleCore<T, D>>>(
pull_counter: PullCounter<T, D, P>,
internal: Rc<RefCell<Vec<Rc<RefCell<ChangeBatch<T>>>>>>,
summaries: Rc<RefCell<Vec<Antichain<T::Summary>>>>,
logging: Option<Logger>
) -> InputHandleCore<T, D, P> {
InputHandleCore {
pull_counter,
internal,
summaries,
logging,
}
}
#[derive(Debug)]
pub struct OutputWrapper<T: Timestamp, D: Container, P: Push<BundleCore<T, D>>> {
push_buffer: BufferCore<T, D, PushCounter<T, D, P>>,
internal_buffer: Rc<RefCell<ChangeBatch<T>>>,
}
impl<T: Timestamp, D: Container, P: Push<BundleCore<T, D>>> OutputWrapper<T, D, P> {
pub fn new(push_buffer: BufferCore<T, D, PushCounter<T, D, P>>, internal_buffer: Rc<RefCell<ChangeBatch<T>>>) -> Self {
OutputWrapper {
push_buffer,
internal_buffer,
}
}
pub fn activate(&mut self) -> OutputHandleCore<T, D, P> {
OutputHandleCore {
push_buffer: &mut self.push_buffer,
internal_buffer: &self.internal_buffer,
}
}
}
pub struct OutputHandleCore<'a, T: Timestamp, C: Container+'a, P: Push<BundleCore<T, C>>+'a> {
push_buffer: &'a mut BufferCore<T, C, PushCounter<T, C, P>>,
internal_buffer: &'a Rc<RefCell<ChangeBatch<T>>>,
}
pub type OutputHandle<'a, T, D, P> = OutputHandleCore<'a, T, Vec<D>, P>;
impl<'a, T: Timestamp, C: Container, P: Push<BundleCore<T, C>>> OutputHandleCore<'a, T, C, P> {
pub fn session<'b, CT: CapabilityTrait<T>>(&'b mut self, cap: &'b CT) -> Session<'b, T, C, PushCounter<T, C, P>> where 'a: 'b {
assert!(cap.valid_for_output(&self.internal_buffer), "Attempted to open output session with invalid capability");
self.push_buffer.session(cap.time())
}
pub fn cease(&mut self) {
self.push_buffer.cease();
}
}
impl<'a, T: Timestamp, C: Container, P: Push<BundleCore<T, C>>> Drop for OutputHandleCore<'a, T, C, P> {
fn drop(&mut self) {
self.push_buffer.cease();
}
}