generic salsa algo

3 files changed, 399 insertions, 0 deletions

diff --git a/crates/salsa/Cargo.toml b/crates/salsa/Cargo.toml
new file mode 100644
index 000000000..9eb83234f
--- /dev/null
+++ b/crates/salsa/Cargo.toml
@@ -0,0 +1,8 @@
+[package]
+name = "salsa"
+version = "0.1.0"
+authors = ["Aleksey Kladov <[email protected]>"]
+
+[dependencies]
+parking_lot = "0.6.3"
+im = "12.0.0"
diff --git a/crates/salsa/src/lib.rs b/crates/salsa/src/lib.rs
new file mode 100644
index 000000000..69c7b35fa
--- /dev/null
+++ b/crates/salsa/src/lib.rs
@@ -0,0 +1,238 @@
+extern crate im;
+extern crate parking_lot;
+
+use std::{
+    sync::Arc,
+    any::Any,
+    collections::HashMap,
+    cell::RefCell,
+};
+use parking_lot::Mutex;
+
+type GroundQueryFn<T> = fn(&T, &(Any + Send + Sync + 'static)) -> (Box<Any + Send + Sync + 'static>, OutputFingerprint);
+type QueryFn<T> = fn(&QueryCtx<T>, &(Any + Send + Sync + 'static)) -> (Box<Any + Send + Sync + 'static>, OutputFingerprint);
+
+#[derive(Debug)]
+pub struct Db<T> {
+    db: Arc<DbState<T>>,
+    query_config: Arc<QueryConfig<T>>,
+}
+
+pub struct QueryConfig<T> {
+    ground_fn: HashMap<QueryTypeId, GroundQueryFn<T>>,
+    query_fn: HashMap<QueryTypeId, QueryFn<T>>,
+}
+
+impl<T> ::std::fmt::Debug for QueryConfig<T> {
+    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
+        ::std::fmt::Display::fmt("QueryConfig { ... }", f)
+    }
+}
+
+#[derive(Debug)]
+struct DbState<T> {
+    ground_data: T,
+    gen: Gen,
+    graph: Mutex<im::HashMap<QueryId, (Gen, Arc<QueryRecord>)>>,
+}
+
+#[derive(Debug)]
+struct QueryRecord {
+    params: Arc<Any + Send + Sync + 'static>,
+    output: Arc<Any + Send + Sync + 'static>,
+    output_fingerprint: OutputFingerprint,
+    deps: Vec<(QueryId, OutputFingerprint)>,
+}
+
+impl<T> DbState<T> {
+    fn record(
+        &self,
+        query_id: QueryId,
+        params: Arc<Any + Send + Sync + 'static>,
+        output: Arc<Any + Send + Sync + 'static>,
+        output_fingerprint: OutputFingerprint,
+        deps: Vec<(QueryId, OutputFingerprint)>,
+    ) {
+        let gen = self.gen;
+        let record = QueryRecord {
+            params,
+            output,
+            output_fingerprint,
+            deps,
+        };
+        self.graph.lock().insert(query_id, (gen, Arc::new(record)));
+    }
+}
+
+impl<T> QueryConfig<T> {
+    pub fn new() -> Self {
+        QueryConfig {
+            ground_fn: HashMap::new(),
+            query_fn: HashMap::new(),
+        }
+    }
+    pub fn with_ground_query(
+        mut self,
+        query_type: QueryTypeId,
+        query_fn: GroundQueryFn<T>
+    ) -> Self {
+        let prev = self.ground_fn.insert(query_type, query_fn);
+        assert!(prev.is_none());
+        self
+    }
+    pub fn with_query(
+        mut self,
+        query_type: QueryTypeId,
+        query_fn: QueryFn<T>,
+    ) -> Self {
+        let prev = self.query_fn.insert(query_type, query_fn);
+        assert!(prev.is_none());
+        self
+    }
+}
+
+pub struct QueryCtx<T> {
+    db: Arc<DbState<T>>,
+    query_config: Arc<QueryConfig<T>>,
+    stack: RefCell<Vec<Vec<(QueryId, OutputFingerprint)>>>,
+    executed: RefCell<Vec<QueryTypeId>>,
+}
+
+impl<T> QueryCtx<T> {
+    fn new(db: &Db<T>) -> QueryCtx<T> {
+        QueryCtx {
+            db: Arc::clone(&db.db),
+            query_config: Arc::clone(&db.query_config),
+            stack: RefCell::new(vec![Vec::new()]),
+            executed: RefCell::new(Vec::new()),
+        }
+    }
+    pub fn get(
+        &self,
+        query_id: QueryId,
+        params: Arc<Any + Send + Sync + 'static>,
+    ) -> Arc<Any + Send + Sync + 'static> {
+        let (res, output_fingerprint) = self.get_inner(query_id, params);
+        self.record_dep(query_id, output_fingerprint);
+        res
+    }
+
+    pub fn get_inner(
+        &self,
+        query_id: QueryId,
+        params: Arc<Any + Send + Sync + 'static>,
+    ) -> (Arc<Any + Send + Sync + 'static>, OutputFingerprint) {
+        let (gen, record) = {
+            let guard = self.db.graph.lock();
+            match guard.get(&query_id).map(|it| it.clone()){
+                None => {
+                    drop(guard);
+                    return self.force(query_id, params);
+                },
+                Some(it) => it,
+            }
+        };
+        if gen == self.db.gen {
+            return (record.output.clone(), record.output_fingerprint)
+        }
+        if self.query_config.ground_fn.contains_key(&query_id.0) {
+            return self.force(query_id, params);
+        }
+        for (dep_query_id, prev_fingerprint) in record.deps.iter().cloned() {
+            let dep_params: Arc<Any + Send + Sync + 'static> = {
+                let guard = self.db.graph.lock();
+                guard[&dep_query_id]
+                .1
+                .params
+                .clone()
+            };
+            if prev_fingerprint != self.get_inner(dep_query_id, dep_params).1 {
+                return self.force(query_id, params)
+            }
+        }
+        let gen = self.db.gen;
+        {
+            let mut guard = self.db.graph.lock();
+            guard[&query_id].0 = gen;
+        }
+        (record.output.clone(), record.output_fingerprint)
+    }
+    fn force(
+        &self,
+        query_id: QueryId,
+        params: Arc<Any + Send + Sync + 'static>,
+    ) -> (Arc<Any + Send + Sync + 'static>, OutputFingerprint) {
+        self.executed.borrow_mut().push(query_id.0);
+        self.stack.borrow_mut().push(Vec::new());
+
+        let (res, output_fingerprint) = if let Some(f) = self.ground_query_fn_by_type(query_id.0) {
+            f(&self.db.ground_data, &*params)
+        } else if let Some(f) = self.query_fn_by_type(query_id.0) {
+            f(self, &*params)
+        } else {
+            panic!("unknown query type: {:?}", query_id.0);
+        };
+
+        let res: Arc<Any + Send + Sync + 'static> = res.into();
+
+        let deps = self.stack.borrow_mut().pop().unwrap();
+        self.db.record(query_id, params, res.clone(), output_fingerprint, deps);
+        (res, output_fingerprint)
+    }
+    fn ground_query_fn_by_type(&self, query_type: QueryTypeId) -> Option<GroundQueryFn<T>> {
+        self.query_config.ground_fn.get(&query_type).map(|&it| it)
+    }
+    fn query_fn_by_type(&self, query_type: QueryTypeId) -> Option<QueryFn<T>> {
+        self.query_config.query_fn.get(&query_type).map(|&it| it)
+    }
+    fn record_dep(
+        &self,
+        query_id: QueryId,
+        output_fingerprint: OutputFingerprint,
+    ) -> () {
+        let mut stack = self.stack.borrow_mut();
+        let deps = stack.last_mut().unwrap();
+        deps.push((query_id, output_fingerprint))
+    }
+}
+
+impl<T> Db<T> {
+    pub fn new(query_config: QueryConfig<T>, ground_data: T) -> Db<T> {
+        Db {
+            db: Arc::new(DbState { ground_data, gen: Gen(0), graph: Default::default() }),
+            query_config: Arc::new(query_config),
+        }
+    }
+
+    pub fn with_ground_data(&self, ground_data: T) -> Db<T> {
+        let gen = Gen(self.db.gen.0 + 1);
+        let graph = self.db.graph.lock().clone();
+        let graph = Mutex::new(graph);
+        Db {
+            db: Arc::new(DbState { ground_data, gen, graph }),
+            query_config: Arc::clone(&self.query_config)
+        }
+    }
+    pub fn get(
+        &self,
+        query_id: QueryId,
+        params: Box<Any + Send + Sync + 'static>,
+    ) -> (Arc<Any + Send + Sync + 'static>, Vec<QueryTypeId>) {
+        let ctx = QueryCtx::new(self);
+        let res = ctx.get(query_id, params.into());
+        let executed = ::std::mem::replace(&mut *ctx.executed.borrow_mut(), Vec::new());
+        (res, executed)
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+struct Gen(u64);
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct InputFingerprint(pub u64);
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct OutputFingerprint(pub u64);
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct QueryTypeId(pub u16);
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct QueryId(pub QueryTypeId, pub InputFingerprint);
+
diff --git a/crates/salsa/tests/integration.rs b/crates/salsa/tests/integration.rs
new file mode 100644
index 000000000..7241eca38
--- /dev/null
+++ b/crates/salsa/tests/integration.rs
@@ -0,0 +1,153 @@
+extern crate salsa;
+use std::{
+    sync::Arc,
+    collections::hash_map::{HashMap, DefaultHasher},
+    any::Any,
+    hash::{Hash, Hasher},
+};
+
+type State = HashMap<u32, String>;
+const GET_TEXT: salsa::QueryTypeId = salsa::QueryTypeId(1);
+const GET_FILES: salsa::QueryTypeId = salsa::QueryTypeId(2);
+const FILE_NEWLINES: salsa::QueryTypeId = salsa::QueryTypeId(3);
+const TOTAL_NEWLINES: salsa::QueryTypeId = salsa::QueryTypeId(4);
+
+fn mk_ground_query<T, R>(
+    state: &State,
+    params: &(Any + Send + Sync + 'static),
+    f: fn(&State, &T) -> R,
+) -> (Box<Any + Send + Sync + 'static>, salsa::OutputFingerprint)
+where
+    T: 'static,
+    R: Hash + Send + Sync + 'static,
+{
+    let params = params.downcast_ref().unwrap();
+    let result = f(state, params);
+    let fingerprint = o_print(&result);
+    (Box::new(result), fingerprint)
+}
+
+fn get<T, R>(db: &salsa::Db<State>, query_type: salsa::QueryTypeId, param: T) -> (Arc<R>, Vec<salsa::QueryTypeId>)
+where
+    T: Hash + Send + Sync + 'static,
+    R: Send + Sync + 'static,
+{
+    let i_print = i_print(&param);
+    let param = Box::new(param);
+    let (res, trace) = db.get(salsa::QueryId(query_type, i_print), param);
+    (res.downcast().unwrap(), trace)
+}
+
+struct QueryCtx<'a>(&'a salsa::QueryCtx<State>);
+
+impl<'a> QueryCtx<'a> {
+    fn get_text(&self, id: u32) -> Arc<String> {
+        let i_print = i_print(&id);
+        let text = self.0.get(salsa::QueryId(GET_TEXT, i_print), Arc::new(id));
+        text.downcast().unwrap()
+    }
+    fn get_files(&self) -> Arc<Vec<u32>> {
+        let i_print = i_print(&());
+        let files = self.0.get(salsa::QueryId(GET_FILES, i_print), Arc::new(()));
+        let res = files.downcast().unwrap();
+        res
+    }
+    fn get_n_lines(&self, id: u32) -> usize {
+        let i_print = i_print(&id);
+        let n_lines = self.0.get(salsa::QueryId(FILE_NEWLINES, i_print), Arc::new(id));
+        *n_lines.downcast().unwrap()
+    }
+}
+
+fn mk_query<T, R>(
+    query_ctx: &salsa::QueryCtx<State>,
+    params: &(Any + Send + Sync + 'static),
+    f: fn(QueryCtx, &T) -> R,
+) -> (Box<Any + Send + Sync + 'static>, salsa::OutputFingerprint)
+where
+    T: 'static,
+    R: Hash + Send + Sync + 'static,
+{
+    let params: &T = params.downcast_ref().unwrap();
+    let query_ctx = QueryCtx(query_ctx);
+    let result = f(query_ctx, params);
+    let fingerprint = o_print(&result);
+    (Box::new(result), fingerprint)
+}
+
+fn mk_queries() -> salsa::QueryConfig<State> {
+    salsa::QueryConfig::<State>::new()
+        .with_ground_query(GET_TEXT, |state, id| {
+            mk_ground_query::<u32, String>(state, id, |state, id| state[id].clone())
+        })
+        .with_ground_query(GET_FILES, |state, id| {
+            mk_ground_query::<(), Vec<u32>>(state, id, |state, &()| state.keys().cloned().collect())
+        })
+        .with_query(FILE_NEWLINES, |query_ctx, id| {
+            mk_query(query_ctx, id, |query_ctx, &id| {
+                let text = query_ctx.get_text(id);
+                text.lines().count()
+            })
+        })
+        .with_query(TOTAL_NEWLINES, |query_ctx, id| {
+            mk_query(query_ctx, id, |query_ctx, &()| {
+                let mut total = 0;
+                for &id in query_ctx.get_files().iter() {
+                    total += query_ctx.get_n_lines(id)
+                }
+                total
+            })
+        })
+}
+
+#[test]
+fn test_number_of_lines() {
+    let mut state = State::new();
+    let db = salsa::Db::new(mk_queries(), state.clone());
+    let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
+    assert_eq!(*newlines, 0);
+    assert_eq!(trace.len(), 2);
+    let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
+    assert_eq!(*newlines, 0);
+    assert_eq!(trace.len(), 0);
+
+    state.insert(1, "hello\nworld".to_string());
+    let db = db.with_ground_data(state.clone());
+    let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
+    assert_eq!(*newlines, 2);
+    assert_eq!(trace.len(), 4);
+
+    state.insert(2, "spam\neggs".to_string());
+    let db = db.with_ground_data(state.clone());
+    let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
+    assert_eq!(*newlines, 4);
+    assert_eq!(trace.len(), 5);
+
+    for i in 0..10 {
+        state.insert(i + 10, "spam".to_string());
+    }
+    let db = db.with_ground_data(state.clone());
+    let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
+    assert_eq!(*newlines, 14);
+    assert_eq!(trace.len(), 24);
+
+    state.insert(15, String::new());
+    let db = db.with_ground_data(state.clone());
+    let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
+    assert_eq!(*newlines, 13);
+    assert_eq!(trace.len(), 15);
+}
+
+fn o_print<T: Hash>(x: &T) -> salsa::OutputFingerprint {
+    let mut hasher = DefaultHasher::new();
+    x.hash(&mut hasher);
+    let hash = hasher.finish();
+    salsa::OutputFingerprint(hash)
+}
+
+fn i_print<T: Hash>(x: &T) -> salsa::InputFingerprint {
+    let mut hasher = DefaultHasher::new();
+    x.hash(&mut hasher);
+    let hash = hasher.finish();
+    salsa::InputFingerprint(hash)
+}