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

//! A remote table: an in-memory stand-in for a query engine, read by scanning and counted.
//! It is not `Eager` and shares no evaluator with it; it answers queries over rows.
#![allow(dead_code)]

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

/// What the engine was asked, and how many rows it sent back.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Asked {
    pub lookups: u64,
    pub summaries: u64,
    pub counts: u64,
    pub stem_lists: u64,
    pub key_lists: u64,
    pub cell_lists: u64,
    /// Keys and cells sent back by `keys` and `cells`: rows, not aggregates.
    pub rows_sent: u64,
}

#[derive(Default)]
pub struct Remote {
    table: Mutex<TableSource>,
    asked: Mutex<Asked>,
}

impl Remote {
    pub fn of(rows: impl IntoIterator<Item = (CellKey, CellValue)>) -> Self {
        let table = TableSource {
            rows: rows.into_iter().collect(),
        };
        Remote {
            table: Mutex::new(table),
            asked: Mutex::default(),
        }
    }

    /// Writes as another client of the database would: in place, or at the end, or removed.
    pub fn write(&self, key: CellKey, value: Option<CellValue>) {
        let mut table = self.table.lock().unwrap();
        let at = table.rows.iter().position(|(k, _)| *k == key);
        match (at, value) {
            (Some(i), Some(v)) => table.rows[i].1 = v,
            (None, Some(v)) => table.rows.push((key, v)),
            (Some(i), None) => {
                table.rows.remove(i);
            }
            (None, None) => {}
        }
    }

    pub fn asked(&self) -> Asked {
        *self.asked.lock().unwrap()
    }

    fn ask(&self, f: impl FnOnce(&mut Asked)) {
        f(&mut self.asked.lock().unwrap());
    }

    fn table(&self) -> std::sync::MutexGuard<'_, TableSource> {
        self.table.lock().unwrap()
    }
}

impl CellSource for Remote {
    fn value_at(&self, key: &CellKey) -> Option<CellValue> {
        self.ask(|a| a.lookups += 1);
        self.table().value_at(key)
    }

    fn summarize(&self, select: &Select) -> Summary {
        self.ask(|a| a.summaries += 1);
        self.table().summarize(select)
    }

    fn keys<'a>(
        &'a self,
        partial: &'a [(String, String)],
    ) -> Box<dyn Iterator<Item = CellKey> + 'a> {
        let keys: Vec<CellKey> = self.table().keys(partial).collect();
        self.ask(|a| {
            a.key_lists += 1;
            a.rows_sent += keys.len() as u64;
        });
        Box::new(keys.into_iter())
    }

    fn cells<'a>(
        &'a self,
        partial: &'a [(String, String)],
    ) -> Box<dyn Iterator<Item = (CellKey, CellValue)> + 'a> {
        let cells: Vec<_> = self.table().cells(partial).collect();
        self.ask(|a| {
            a.cell_lists += 1;
            a.rows_sent += cells.len() as u64;
        });
        Box::new(cells.into_iter())
    }

    fn count(&self, select: &Select, excluding: &[(String, String)]) -> usize {
        self.ask(|a| a.counts += 1);
        self.table().count(select, excluding)
    }

    fn distinct_without(&self, without: &[String]) -> Vec<CellKey> {
        let stems = self.table().distinct_without(without);
        self.ask(|a| {
            a.stem_lists += 1;
            a.rows_sent += stems.len() as u64;
        });
        stems
    }
}

/// A source with a fault, for checking that a differential can see it: it forgets the hidden
/// items a count leaves out.
pub struct IgnoresHidden(pub Arc<Remote>);

impl CellSource for IgnoresHidden {
    fn value_at(&self, key: &CellKey) -> Option<CellValue> {
        self.0.value_at(key)
    }

    fn summarize(&self, select: &Select) -> Summary {
        self.0.summarize(select)
    }

    fn keys<'a>(
        &'a self,
        partial: &'a [(String, String)],
    ) -> Box<dyn Iterator<Item = CellKey> + 'a> {
        self.0.keys(partial)
    }

    fn cells<'a>(
        &'a self,
        partial: &'a [(String, String)],
    ) -> Box<dyn Iterator<Item = (CellKey, CellValue)> + 'a> {
        self.0.cells(partial)
    }

    fn count(&self, select: &Select, _excluding: &[(String, String)]) -> usize {
        self.0.count(select, &[])
    }

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

/// A `Pushdown` over a remote table that holds the model's cells; the model keeps no store.
/// A cell edit is a write to the table, as another client's would be, then a refresh.
pub struct Pushed {
    pub pd: Pushdown,
    pub remote: Arc<Remote>,
}

impl Pushed {
    pub fn new(model: &Model, none_cats: Vec<String>) -> Pushed {
        let remote = Arc::new(Remote::of(rows_of(model)));
        let pd = Pushdown::new(bare(model), remote.clone(), none_cats);
        Pushed { pd, remote }
    }

    /// The same, reading the table through `wrap`.
    pub fn through<Q: CellSource + Send + Sync + 'static>(
        model: &Model,
        none_cats: Vec<String>,
        wrap: impl FnOnce(Arc<Remote>) -> Q,
    ) -> Pushed {
        let remote = Arc::new(Remote::of(rows_of(model)));
        let pd = Pushdown::new(bare(model), Arc::new(wrap(remote.clone())), none_cats);
        Pushed { pd, remote }
    }
}

fn rows_of(model: &Model) -> Vec<(CellKey, CellValue)> {
    model
        .data
        .iter_cells()
        .map(|(k, v)| (k, v.clone()))
        .collect()
}

/// `model` with no cells of its own.
fn bare(model: &Model) -> Model {
    let mut bare = model.clone();
    bare.data = improvise_core::model::cell::DataStore::new();
    bare
}

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

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

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

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

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

impl Maintained for Pushed {
    fn apply(&mut self, edit: &Edit) {
        match edit {
            Edit::SetCell(key, value) => self.remote.write(key.clone(), Some(value.clone())),
            Edit::ClearCell(key) => self.remote.write(key.clone(), None),
            _ => return self.pd.define(edit),
        }
        self.pd.refresh();
    }
}