improvise/crates/zset/src/trace.rs

//! The trace interface: operator state behind look up, read a range, apply a delta.

use crate::ZSet;
use std::collections::BTreeMap;
use std::ops::Bound;

/// An integrated `ZSet<(K, V)>`, indexed by `K`. Reads return owned values, so an
/// implementation may page or compress what it holds. See `context/code/zset.md`.
pub trait Trace<K, V> {
    /// Adds `w` to the weight of `(k, v)`. A panic must leave the trace as it was: a
    /// transaction (`Circuit::begin`) logs an apply only once it has returned.
    fn apply(&mut self, k: &K, v: &V, w: i64);
    /// The values at `k` with their non-zero weights, ascending by value.
    fn lookup(&self, k: &K) -> Vec<(V, i64)>;
    /// Every `(key, value, weight)` with a key in the range, ascending.
    fn range(&self, lo: Bound<&K>, hi: Bound<&K>) -> Vec<(K, V, i64)>;
}

/// The in-memory trace: a sorted map from key to a Z-set of values.
#[derive(Clone, Debug)]
pub struct MemTrace<K, V> {
    groups: BTreeMap<K, ZSet<V>>,
}

impl<K, V> Default for MemTrace<K, V> {
    fn default() -> Self {
        MemTrace {
            groups: BTreeMap::new(),
        }
    }
}

impl<K: Ord, V: Ord> MemTrace<K, V> {
    pub fn new() -> Self {
        Self::default()
    }

    /// The number of keys with at least one value.
    pub fn keys(&self) -> usize {
        self.groups.len()
    }
}

fn is_empty_range<K: Ord>(lo: Bound<&K>, hi: Bound<&K>) -> bool {
    match (lo, hi) {
        (Bound::Included(a), Bound::Included(b)) => a > b,
        (Bound::Included(a) | Bound::Excluded(a), Bound::Included(b) | Bound::Excluded(b)) => {
            a >= b
        }
        _ => false,
    }
}

impl<K: Ord + Clone, V: Ord + Clone> Trace<K, V> for MemTrace<K, V> {
    fn apply(&mut self, k: &K, v: &V, w: i64) {
        // Every clone is made before the map changes, so a panicking one changes nothing.
        match self.groups.get_mut(k) {
            Some(group) => {
                group.insert(v.clone(), w);
                if group.is_zero() {
                    self.groups.remove(k);
                }
            }
            None if w != 0 => {
                let group = ZSet::singleton(v.clone(), w);
                self.groups.insert(k.clone(), group);
            }
            None => {}
        }
    }

    fn lookup(&self, k: &K) -> Vec<(V, i64)> {
        self.groups
            .get(k)
            .map(|g| g.iter().map(|(v, w)| (v.clone(), w)).collect())
            .unwrap_or_default()
    }

    fn range(&self, lo: Bound<&K>, hi: Bound<&K>) -> Vec<(K, V, i64)> {
        if is_empty_range(lo, hi) {
            return Vec::new();
        }
        self.groups
            .range::<K, _>((lo, hi))
            .flat_map(|(k, g)| g.iter().map(move |(v, w)| (k.clone(), v.clone(), w)))
            .collect()
    }
}