use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use core::cell::{Cell, RefCell};
use core::fmt;
use crate::error::QueryError;
use crate::revision::Revision;
use crate::stats::Stats;
use crate::system::System;
struct Input<V> {
value: V,
changed_at: Revision,
}
type MemoMap<S> = BTreeMap<<S as System>::Key, Memo<<S as System>::Key, <S as System>::Value>>;
struct Memo<K, V> {
value: V,
deps: Vec<K>,
changed_at: Revision,
verified_at: Revision,
}
pub struct Database<S: System> {
system: S,
revision: Revision,
inputs: BTreeMap<S::Key, Input<S::Value>>,
memos: RefCell<MemoMap<S>>,
frames: RefCell<Vec<Vec<S::Key>>>,
active: RefCell<Vec<S::Key>>,
computed: Cell<u64>,
validated: Cell<u64>,
hits: Cell<u64>,
}
impl<S: System> Database<S> {
#[must_use]
pub fn new(system: S) -> Self {
Self {
system,
revision: Revision::START,
inputs: BTreeMap::new(),
memos: RefCell::new(BTreeMap::new()),
frames: RefCell::new(Vec::new()),
active: RefCell::new(Vec::new()),
computed: Cell::new(0),
validated: Cell::new(0),
hits: Cell::new(0),
}
}
pub fn set(&mut self, key: S::Key, value: S::Value) {
self.memos.get_mut().remove(&key);
match self.inputs.get_mut(&key) {
Some(existing) if existing.value == value => {
}
Some(existing) => {
self.revision = self.revision.next();
existing.value = value;
existing.changed_at = self.revision;
}
None => {
self.revision = self.revision.next();
let changed_at = self.revision;
self.inputs.insert(key, Input { value, changed_at });
}
}
}
pub fn get(&self, key: &S::Key) -> Result<S::Value, QueryError> {
if let Some(frame) = self.frames.borrow_mut().last_mut() {
frame.push(key.clone());
}
let (value, _changed_at) = self.eval(key)?;
Ok(value)
}
#[must_use]
#[inline]
pub const fn revision(&self) -> Revision {
self.revision
}
#[must_use]
pub fn stats(&self) -> Stats {
Stats {
computed: self.computed.get(),
validated: self.validated.get(),
hits: self.hits.get(),
}
}
#[must_use]
#[inline]
pub const fn system(&self) -> &S {
&self.system
}
fn eval(&self, key: &S::Key) -> Result<(S::Value, Revision), QueryError> {
if let Some(input) = self.inputs.get(key) {
return Ok((input.value.clone(), input.changed_at));
}
{
let memos = self.memos.borrow();
if let Some(memo) = memos.get(key) {
if memo.verified_at == self.revision {
self.hits.set(self.hits.get().saturating_add(1));
return Ok((memo.value.clone(), memo.changed_at));
}
}
}
let snapshot = self
.memos
.borrow()
.get(key)
.map(|m| (m.deps.clone(), m.verified_at, m.value.clone(), m.changed_at));
if let Some((deps, verified_at, value, changed_at)) = snapshot {
let mut inputs_changed = false;
for dep in &deps {
let (_dep_value, dep_changed_at) = self.eval(dep)?;
if dep_changed_at > verified_at {
inputs_changed = true;
break;
}
}
if !inputs_changed {
if let Some(memo) = self.memos.borrow_mut().get_mut(key) {
memo.verified_at = self.revision;
}
self.validated.set(self.validated.get().saturating_add(1));
return Ok((value, changed_at));
}
return self.recompute(key, Some((value, changed_at)));
}
self.recompute(key, None)
}
fn recompute(
&self,
key: &S::Key,
previous: Option<(S::Value, Revision)>,
) -> Result<(S::Value, Revision), QueryError> {
if self.active.borrow().iter().any(|active| active == key) {
return Err(QueryError::Cycle);
}
self.active.borrow_mut().push(key.clone());
self.frames.borrow_mut().push(Vec::new());
let result = self.system.compute(self, key);
let deps = self.frames.borrow_mut().pop().unwrap_or_default();
self.active.borrow_mut().pop();
let value = result?;
let changed_at = match previous {
Some((old_value, old_changed_at)) if old_value == value => old_changed_at,
_ => self.revision,
};
self.computed.set(self.computed.get().saturating_add(1));
let memo = Memo {
value: value.clone(),
deps,
changed_at,
verified_at: self.revision,
};
self.memos.borrow_mut().insert(key.clone(), memo);
Ok((value, changed_at))
}
}
impl<S: System> fmt::Debug for Database<S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Database")
.field("revision", &self.revision)
.field("inputs", &self.inputs.len())
.field("memos", &self.memos.borrow().len())
.field("stats", &self.stats())
.finish()
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
use alloc::string::{String, ToString};
use core::cell::Cell as StdCell;
struct Counting {
runs: StdCell<[u32; 8]>,
}
impl Counting {
fn new() -> Self {
Self {
runs: StdCell::new([0; 8]),
}
}
fn runs_of(&self, key: usize) -> u32 {
self.runs.get()[key]
}
fn bump(&self, key: usize) {
let mut r = self.runs.get();
r[key] += 1;
self.runs.set(r);
}
}
impl System for Counting {
type Key = usize;
type Value = i64;
fn compute(&self, db: &Database<Self>, key: &usize) -> Result<i64, QueryError> {
self.bump(*key);
match key {
0 | 1 => Ok(0), 2 => Ok(db.get(&0)? * 10),
3 => Ok(db.get(&2)? + db.get(&1)?),
4 => Ok(db.get(&2)?.signum()),
_ => Ok(0),
}
}
}
#[test]
fn test_input_get_returns_set_value() {
let mut db = Database::new(Counting::new());
db.set(0, 7);
assert_eq!(db.get(&0).unwrap(), 7);
assert_eq!(db.system().runs_of(0), 0);
}
#[test]
fn test_first_get_computes_once() {
let mut db = Database::new(Counting::new());
db.set(0, 3);
assert_eq!(db.get(&2).unwrap(), 30);
assert_eq!(db.system().runs_of(2), 1);
assert_eq!(db.stats().computed, 1);
}
#[test]
fn test_second_get_is_a_hit() {
let mut db = Database::new(Counting::new());
db.set(0, 3);
assert_eq!(db.get(&2).unwrap(), 30);
let before = db.stats();
assert_eq!(db.get(&2).unwrap(), 30);
let after = db.stats();
assert_eq!(after.hits - before.hits, 1);
assert_eq!(after.computed, before.computed); }
#[test]
fn test_input_change_recomputes_dependents() {
let mut db = Database::new(Counting::new());
db.set(0, 3);
assert_eq!(db.get(&2).unwrap(), 30);
db.set(0, 4);
assert_eq!(db.get(&2).unwrap(), 40);
assert_eq!(db.system().runs_of(2), 2);
}
#[test]
fn test_unchanged_input_set_does_not_recompute() {
let mut db = Database::new(Counting::new());
db.set(0, 3);
assert_eq!(db.get(&2).unwrap(), 30);
db.set(0, 3); assert_eq!(db.get(&2).unwrap(), 30);
assert_eq!(db.system().runs_of(2), 1);
}
#[test]
fn test_multilayer_hit_after_no_change() {
let mut db = Database::new(Counting::new());
db.set(0, 3);
assert_eq!(db.get(&4).unwrap(), 1);
assert_eq!(db.system().runs_of(4), 1);
assert_eq!(db.system().runs_of(2), 1);
db.set(0, 3);
assert_eq!(db.get(&4).unwrap(), 1);
assert_eq!(db.system().runs_of(4), 1);
assert_eq!(db.system().runs_of(2), 1);
}
#[test]
fn test_early_cutoff_when_intermediate_value_is_stable() {
struct AbsChain {
top_runs: StdCell<u32>,
}
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
enum K {
In,
Abs,
Plus1,
}
impl System for AbsChain {
type Key = K;
type Value = i64;
fn compute(&self, db: &Database<Self>, key: &K) -> Result<i64, QueryError> {
match key {
K::In => Ok(0),
K::Abs => Ok(db.get(&K::In)?.abs()),
K::Plus1 => {
self.top_runs.set(self.top_runs.get() + 1);
Ok(db.get(&K::Abs)? + 1)
}
}
}
}
let mut db = Database::new(AbsChain {
top_runs: StdCell::new(0),
});
db.set(K::In, 5);
assert_eq!(db.get(&K::Plus1).unwrap(), 6);
assert_eq!(db.system().top_runs.get(), 1);
db.set(K::In, -5);
assert_eq!(db.get(&K::Plus1).unwrap(), 6);
assert_eq!(db.system().top_runs.get(), 1);
assert!(db.stats().validated >= 1);
}
#[test]
fn test_default_value_when_input_unset() {
let db = Database::new(Counting::new());
assert_eq!(db.get(&2).unwrap(), 0);
}
#[test]
fn test_direct_self_cycle_is_error() {
struct SelfDep;
impl System for SelfDep {
type Key = u32;
type Value = u32;
fn compute(&self, db: &Database<Self>, k: &u32) -> Result<u32, QueryError> {
db.get(k)
}
}
let db = Database::new(SelfDep);
assert_eq!(db.get(&1), Err(QueryError::Cycle));
}
#[test]
fn test_indirect_cycle_is_error() {
struct PingPong;
impl System for PingPong {
type Key = u32;
type Value = u32;
fn compute(&self, db: &Database<Self>, k: &u32) -> Result<u32, QueryError> {
match k {
0 => db.get(&1),
_ => db.get(&0),
}
}
}
let db = Database::new(PingPong);
assert_eq!(db.get(&0), Err(QueryError::Cycle));
}
#[test]
fn test_state_is_clean_after_cycle_error() {
struct SelfDep;
impl System for SelfDep {
type Key = u32;
type Value = u32;
fn compute(&self, db: &Database<Self>, k: &u32) -> Result<u32, QueryError> {
if *k == 0 { db.get(&0) } else { Ok(*k) }
}
}
let db = Database::new(SelfDep);
assert_eq!(db.get(&0), Err(QueryError::Cycle));
assert_eq!(db.get(&9).unwrap(), 9);
}
#[test]
fn test_set_promotes_derived_key_to_input() {
let mut db = Database::new(Counting::new());
db.set(0, 3);
assert_eq!(db.get(&2).unwrap(), 30); db.set(2, 999); assert_eq!(db.get(&2).unwrap(), 999);
let runs_before = db.system().runs_of(2);
assert_eq!(db.get(&2).unwrap(), 999);
assert_eq!(db.system().runs_of(2), runs_before);
}
#[test]
fn test_branching_dependencies_tracked_precisely() {
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
enum K {
Select,
Left,
Right,
Picked,
}
struct Switch {
runs: StdCell<u32>,
}
impl System for Switch {
type Key = K;
type Value = i64;
fn compute(&self, db: &Database<Self>, key: &K) -> Result<i64, QueryError> {
match key {
K::Select | K::Left | K::Right => Ok(0),
K::Picked => {
self.runs.set(self.runs.get() + 1);
if db.get(&K::Select)? == 0 {
db.get(&K::Left)
} else {
db.get(&K::Right)
}
}
}
}
}
let mut db = Database::new(Switch {
runs: StdCell::new(0),
});
db.set(K::Select, 0);
db.set(K::Left, 100);
db.set(K::Right, 200);
assert_eq!(db.get(&K::Picked).unwrap(), 100);
assert_eq!(db.system().runs.get(), 1);
db.set(K::Right, 999);
assert_eq!(db.get(&K::Picked).unwrap(), 100);
assert_eq!(db.system().runs.get(), 1);
db.set(K::Left, 111);
assert_eq!(db.get(&K::Picked).unwrap(), 111);
assert_eq!(db.system().runs.get(), 2);
}
#[test]
fn test_string_values_work() {
struct Greeter;
impl System for Greeter {
type Key = u32;
type Value = String;
fn compute(&self, db: &Database<Self>, k: &u32) -> Result<String, QueryError> {
if *k == 0 {
Ok("world".to_string())
} else {
Ok(alloc::format!("hello, {}", db.get(&0)?))
}
}
}
let db = Database::new(Greeter);
assert_eq!(db.get(&1).unwrap(), "hello, world");
}
#[test]
fn test_debug_reports_shape() {
let mut db = Database::new(Counting::new());
db.set(0, 1);
let _ = db.get(&2).unwrap();
let rendered = alloc::format!("{db:?}");
assert!(rendered.contains("revision"));
assert!(rendered.contains("memos"));
}
#[test]
fn test_diamond_recomputes_shared_dep_once() {
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
enum K {
In,
Left,
Right,
Top,
}
struct Diamond {
in_runs: StdCell<u32>,
}
impl System for Diamond {
type Key = K;
type Value = i64;
fn compute(&self, db: &Database<Self>, key: &K) -> Result<i64, QueryError> {
match key {
K::In => Ok(0),
K::Left => Ok(db.get(&K::In)? + 1),
K::Right => Ok(db.get(&K::In)? + 2),
K::Top => {
self.in_runs.set(self.in_runs.get() + 1);
Ok(db.get(&K::Left)? + db.get(&K::Right)?)
}
}
}
}
let mut db = Database::new(Diamond {
in_runs: StdCell::new(0),
});
db.set(K::In, 10);
assert_eq!(db.get(&K::Top).unwrap(), (10 + 1) + (10 + 2));
assert_eq!(db.system().in_runs.get(), 1);
db.set(K::In, 20);
assert_eq!(db.get(&K::Top).unwrap(), (20 + 1) + (20 + 2));
assert_eq!(db.system().in_runs.get(), 2);
}
#[test]
fn test_multiple_deps_each_invalidate() {
let mut db = Database::new(Counting::new());
db.set(0, 1);
db.set(1, 2);
assert_eq!(db.get(&3).unwrap(), 12);
db.set(1, 5);
assert_eq!(db.get(&3).unwrap(), 15);
db.set(0, 2); assert_eq!(db.get(&3).unwrap(), 25);
}
}