use std::collections::{BinaryHeap, HashMap};
use timely::order::{PartialOrder, TotalOrder};
use timely::dataflow::{Scope, Stream};
use timely::dataflow::operators::generic::Operator;
use timely::dataflow::channels::pact::Exchange;
use timely::progress::Timestamp;
use timely::progress::Antichain;
use timely::dataflow::operators::Capability;
use crate::lattice::Lattice;
use crate::operators::arrange::arrangement::Arranged;
use crate::trace::Builder;
use crate::trace::{self, Trace, TraceReader, Batch, Cursor};
use crate::{ExchangeData, Hashable};
use super::TraceAgent;
pub fn arrange_from_upsert<G, Tr>(
stream: &Stream<G, (Tr::KeyOwned, Option<Tr::ValOwned>, G::Timestamp)>,
name: &str,
) -> Arranged<G, TraceAgent<Tr>>
where
G: Scope,
G::Timestamp: Lattice+Ord+TotalOrder+ExchangeData,
Tr::KeyOwned: ExchangeData+Hashable+std::hash::Hash,
Tr::ValOwned: ExchangeData,
Tr: Trace+TraceReader<Time=G::Timestamp,Diff=isize>+'static,
Tr::Batch: Batch,
Tr::Builder: Builder<Item = ((Tr::KeyOwned, Tr::ValOwned), Tr::Time, Tr::Diff)>,
{
let mut reader: Option<TraceAgent<Tr>> = None;
let stream = {
let reader = &mut reader;
let exchange = Exchange::new(move |update: &(Tr::KeyOwned,Option<Tr::ValOwned>,G::Timestamp)| (update.0).hashed().into());
stream.unary_frontier(exchange, name, move |_capability, info| {
let logger = {
let scope = stream.scope();
let register = scope.log_register();
register.get::<crate::logging::DifferentialEvent>("differential/arrange")
};
let mut capabilities = Antichain::<Capability<G::Timestamp>>::new();
let mut buffer = Vec::new();
let activator = Some(stream.scope().activator_for(&info.address[..]));
let mut empty_trace = Tr::new(info.clone(), logger.clone(), activator);
if let Some(exert_logic) = stream.scope().config().get::<trace::ExertionLogic>("differential/default_exert_logic").cloned() {
empty_trace.set_exert_logic(exert_logic);
}
let (mut reader_local, mut writer) = TraceAgent::new(empty_trace, info, logger);
*reader = Some(reader_local.clone());
let mut prev_frontier = Antichain::from_elem(<G::Timestamp as Timestamp>::minimum());
let mut priority_queue = BinaryHeap::<std::cmp::Reverse<(G::Timestamp, Tr::KeyOwned, Option<Tr::ValOwned>)>>::new();
let mut updates = Vec::new();
move |input, output| {
input.for_each(|cap, data| {
capabilities.insert(cap.retain());
data.swap(&mut buffer);
for (key, val, time) in buffer.drain(..) {
priority_queue.push(std::cmp::Reverse((time, key, val)))
}
});
assert!(PartialOrder::less_equal(&prev_frontier.borrow(), &input.frontier().frontier()));
if prev_frontier.borrow() != input.frontier().frontier() {
if capabilities.elements().iter().any(|c| !input.frontier().less_equal(c.time())) {
let mut upper = Antichain::new();
for (index, capability) in capabilities.elements().iter().enumerate() {
if !input.frontier().less_equal(capability.time()) {
upper.clear();
for time in input.frontier().frontier().iter() {
upper.insert(time.clone());
}
for other_capability in &capabilities.elements()[(index + 1) .. ] {
upper.insert(other_capability.time().clone());
}
let mut to_process = HashMap::new();
while priority_queue.peek().map(|std::cmp::Reverse((t,_k,_v))| !upper.less_equal(t)).unwrap_or(false) {
let std::cmp::Reverse((time, key, val)) = priority_queue.pop().expect("Priority queue just ensured non-empty");
to_process.entry(key).or_insert(Vec::new()).push((time, std::cmp::Reverse(val)));
}
if priority_queue.capacity() > 4 * priority_queue.len() {
priority_queue.shrink_to_fit();
}
let mut to_process = to_process.into_iter().collect::<Vec<_>>();
to_process.sort();
let (mut trace_cursor, trace_storage) = reader_local.cursor();
let mut builder = Tr::Builder::new();
for (key, mut list) in to_process.drain(..) {
use trace::cursor::MyTrait;
let mut prev_value: Option<Tr::ValOwned> = None;
trace_cursor.seek_key_owned(&trace_storage, &key);
if trace_cursor.get_key(&trace_storage).map(|k| k.equals(&key)).unwrap_or(false) {
while let Some(val) = trace_cursor.get_val(&trace_storage) {
let mut count = 0;
trace_cursor.map_times(&trace_storage, |_time, diff| count += *diff);
assert!(count == 0 || count == 1);
if count == 1 {
assert!(prev_value.is_none());
prev_value = Some(val.into_owned());
}
trace_cursor.step_val(&trace_storage);
}
trace_cursor.step_key(&trace_storage);
}
list.sort();
list.dedup_by(|(t1,_), (t2,_)| t1 == t2);
for (time, std::cmp::Reverse(next)) in list {
if prev_value != next {
if let Some(prev) = prev_value {
updates.push(((key.clone(), prev), time.clone(), -1));
}
if let Some(next) = next.as_ref() {
updates.push(((key.clone(), next.clone()), time.clone(), 1));
}
prev_value = next;
}
}
updates.sort();
for update in updates.drain(..) {
builder.push(update);
}
}
let batch = builder.done(prev_frontier.clone(), upper.clone(), Antichain::from_elem(G::Timestamp::minimum()));
prev_frontier.clone_from(&upper);
writer.insert(batch.clone(), Some(capability.time().clone()));
output.session(&capabilities.elements()[index]).give(batch);
}
}
let mut new_capabilities = Antichain::new();
if let Some(std::cmp::Reverse((time, _, _))) = priority_queue.peek() {
if let Some(capability) = capabilities.elements().iter().find(|c| c.time().less_equal(time)) {
new_capabilities.insert(capability.delayed(time));
}
else {
panic!("failed to find capability");
}
}
capabilities = new_capabilities;
}
else {
writer.seal(input.frontier().frontier().to_owned());
}
prev_frontier.clear();
prev_frontier.extend(input.frontier().frontier().iter().cloned());
reader_local.set_logical_compaction(prev_frontier.borrow());
reader_local.set_physical_compaction(prev_frontier.borrow());
}
writer.exert();
}
})
};
Arranged { stream, trace: reader.unwrap() }
}