improvise/crates/improvise-incremental/tests/support/live.rs

//! A live table: a scripted source that reports its changes, the way a database would.
//! Its rows are those as of the last drain; a write is pending until then.
#![allow(dead_code)]

use improvise_core::model::cell::{CellChange, CellKey, CellValue, DataStore};
use improvise_core::model::source::{CellSource, ChangeSource, Select, Summary};
use improvise_core::model::{Aggregation, Model};
use improvise_incremental::{Derived, Edit, Incremental, Maintained, Reading};
use std::sync::{Arc, Mutex};

#[derive(Default)]
struct State {
    /// What a read sees: the rows as of the last drain.
    visible: DataStore,
    /// The rows with every pending write applied, to say what a write changed.
    latest: DataStore,
    pending: Vec<CellChange>,
    /// Rows handed out by a snapshot read (`cells` with no partial).
    snapshot_rows: u64,
}

#[derive(Default)]
pub struct LiveTable {
    state: Mutex<State>,
}

impl LiveTable {
    /// A table already holding `rows`, with nothing pending.
    pub fn of(rows: impl IntoIterator<Item = (CellKey, CellValue)>) -> Self {
        let mut state = State::default();
        for (key, value) in rows {
            state.visible.set(key.clone(), value.clone());
            state.latest.set(key, value);
        }
        LiveTable {
            state: Mutex::new(state),
        }
    }

    /// Writes `value` at `key` (or removes it), as a database user would: visible after the next drain.
    pub fn write(&self, key: CellKey, value: Option<CellValue>) {
        let mut state = self.state.lock().unwrap();
        let old = state.latest.get(&key).cloned();
        match &value {
            Some(v) => state.latest.set(key.clone(), v.clone()),
            None if old.is_none() => return,
            None => state.latest.remove(&key),
        }
        state.pending.push(CellChange {
            key,
            old,
            new: value,
        });
    }

    /// How many rows snapshot reads have handed out so far.
    pub fn snapshot_rows(&self) -> u64 {
        self.state.lock().unwrap().snapshot_rows
    }
}

impl CellSource for LiveTable {
    fn value_at(&self, key: &CellKey) -> Option<CellValue> {
        self.state.lock().unwrap().visible.value_at(key)
    }

    fn summarize(&self, select: &Select) -> Summary {
        self.state.lock().unwrap().visible.summarize(select)
    }

    fn keys<'a>(
        &'a self,
        partial: &'a [(String, String)],
    ) -> Box<dyn Iterator<Item = CellKey> + 'a> {
        let state = self.state.lock().unwrap();
        let keys: Vec<CellKey> = state.visible.keys(partial).collect();
        Box::new(keys.into_iter())
    }

    fn cells<'a>(
        &'a self,
        partial: &'a [(String, String)],
    ) -> Box<dyn Iterator<Item = (CellKey, CellValue)> + 'a> {
        let mut state = self.state.lock().unwrap();
        let cells: Vec<_> = state.visible.cells(partial).collect();
        if partial.is_empty() {
            state.snapshot_rows += cells.len() as u64;
        }
        Box::new(cells.into_iter())
    }

    fn count(&self, select: &Select, excluding: &[(String, String)]) -> usize {
        self.state.lock().unwrap().visible.count(select, excluding)
    }

    fn distinct_without(&self, without: &[String]) -> Vec<CellKey> {
        self.state.lock().unwrap().visible.distinct_without(without)
    }
}

impl ChangeSource for LiveTable {
    fn drain_changes(&self) -> Vec<CellChange> {
        let mut state = self.state.lock().unwrap();
        let changes = std::mem::take(&mut state.pending);
        for change in &changes {
            match &change.new {
                Some(value) => state.visible.set(change.key.clone(), value.clone()),
                None => state.visible.remove(&change.key),
            }
        }
        changes
    }
}

/// An `Incremental` over a live table that holds the model's cells; the model keeps no store.
/// A cell edit is a write to the table, the way another client's would be, followed by a drain.
pub struct Fed {
    pub inc: Incremental,
    pub live: Arc<LiveTable>,
}

impl Fed {
    pub fn new(model: &Model, none_cats: Vec<String>) -> Fed {
        let rows = model.data.iter_cells().map(|(k, v)| (k, v.clone()));
        let live = Arc::new(LiveTable::of(rows));
        let mut bare = model.clone();
        bare.data = DataStore::new();
        let inc = Incremental::over(bare, live.clone(), none_cats);
        Fed { inc, live }
    }
}

impl Derived for Fed {
    fn value(&self, key: &CellKey, agg: Aggregation) -> Option<CellValue> {
        self.inc.value(key, agg)
    }

    fn raw_fold(&self, partial: &CellKey) -> Option<CellValue> {
        self.inc.raw_fold(partial)
    }

    fn covered(&self, key: &CellKey, reading: &Reading) -> Option<CellValue> {
        self.inc.covered(key, reading)
    }

    fn total(&self, key: &CellKey, over: &[String], reading: &Reading) -> Option<CellValue> {
        self.inc.total(key, over, reading)
    }

    fn occupancy(&self, partial: &CellKey, reading: &Reading) -> usize {
        self.inc.occupancy(partial, reading)
    }
}

impl Maintained for Fed {
    fn apply(&mut self, edit: &Edit) {
        match edit {
            Edit::SetCell(key, value) => self.live.write(key.clone(), Some(value.clone())),
            Edit::ClearCell(key) => self.live.write(key.clone(), None),
            _ => return self.inc.apply(edit),
        }
        self.inc.sync();
    }
}