2 files changed, 61 insertions, 55 deletions

diff --git a/crates/salsa/src/lib.rs b/crates/salsa/src/lib.rs
index 69c7b35fa..a54f2a06f 100644
--- a/crates/salsa/src/lib.rs
+++ b/crates/salsa/src/lib.rs
@@ -3,53 +3,52 @@ 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);
+type GroundQueryFn<T, D> = fn(&T, &D) -> (D, OutputFingerprint);
+type QueryFn<T, D> = fn(&QueryCtx<T, D>, &D) -> (D, OutputFingerprint);
 
 #[derive(Debug)]
-pub struct Db<T> {
-    db: Arc<DbState<T>>,
-    query_config: Arc<QueryConfig<T>>,
+pub struct Db<T, D> {
+    db: Arc<DbState<T, D>>,
+    query_config: Arc<QueryConfig<T, D>>,
 }
 
-pub struct QueryConfig<T> {
-    ground_fn: HashMap<QueryTypeId, GroundQueryFn<T>>,
-    query_fn: HashMap<QueryTypeId, QueryFn<T>>,
+pub struct QueryConfig<T, D> {
+    ground_fn: HashMap<QueryTypeId, GroundQueryFn<T, D>>,
+    query_fn: HashMap<QueryTypeId, QueryFn<T, D>>,
 }
 
-impl<T> ::std::fmt::Debug for QueryConfig<T> {
+impl<T, D> ::std::fmt::Debug for QueryConfig<T, D> {
     fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
         ::std::fmt::Display::fmt("QueryConfig { ... }", f)
     }
 }
 
 #[derive(Debug)]
-struct DbState<T> {
+struct DbState<T, D> {
     ground_data: T,
     gen: Gen,
-    graph: Mutex<im::HashMap<QueryId, (Gen, Arc<QueryRecord>)>>,
+    graph: Mutex<im::HashMap<QueryId, (Gen, Arc<QueryRecord<D>>)>>,
 }
 
 #[derive(Debug)]
-struct QueryRecord {
-    params: Arc<Any + Send + Sync + 'static>,
-    output: Arc<Any + Send + Sync + 'static>,
+struct QueryRecord<D> {
+    params: D,
+    output: D,
     output_fingerprint: OutputFingerprint,
     deps: Vec<(QueryId, OutputFingerprint)>,
 }
 
-impl<T> DbState<T> {
+impl<T, D> DbState<T, D> {
     fn record(
         &self,
         query_id: QueryId,
-        params: Arc<Any + Send + Sync + 'static>,
-        output: Arc<Any + Send + Sync + 'static>,
+        params: D,
+        output: D,
         output_fingerprint: OutputFingerprint,
         deps: Vec<(QueryId, OutputFingerprint)>,
     ) {
@@ -64,7 +63,7 @@ impl<T> DbState<T> {
     }
 }
 
-impl<T> QueryConfig<T> {
+impl<T, D> QueryConfig<T, D> {
     pub fn new() -> Self {
         QueryConfig {
             ground_fn: HashMap::new(),
@@ -74,7 +73,7 @@ impl<T> QueryConfig<T> {
     pub fn with_ground_query(
         mut self,
         query_type: QueryTypeId,
-        query_fn: GroundQueryFn<T>
+        query_fn: GroundQueryFn<T, D>
     ) -> Self {
         let prev = self.ground_fn.insert(query_type, query_fn);
         assert!(prev.is_none());
@@ -83,7 +82,7 @@ impl<T> QueryConfig<T> {
     pub fn with_query(
         mut self,
         query_type: QueryTypeId,
-        query_fn: QueryFn<T>,
+        query_fn: QueryFn<T, D>,
     ) -> Self {
         let prev = self.query_fn.insert(query_type, query_fn);
         assert!(prev.is_none());
@@ -91,15 +90,18 @@ impl<T> QueryConfig<T> {
     }
 }
 
-pub struct QueryCtx<T> {
-    db: Arc<DbState<T>>,
-    query_config: Arc<QueryConfig<T>>,
+pub struct QueryCtx<T, D> {
+    db: Arc<DbState<T, D>>,
+    query_config: Arc<QueryConfig<T, D>>,
     stack: RefCell<Vec<Vec<(QueryId, OutputFingerprint)>>>,
     executed: RefCell<Vec<QueryTypeId>>,
 }
 
-impl<T> QueryCtx<T> {
-    fn new(db: &Db<T>) -> QueryCtx<T> {
+impl<T, D> QueryCtx<T, D>
+where
+    D: Clone
+{
+    fn new(db: &Db<T, D>) -> QueryCtx<T, D> {
         QueryCtx {
             db: Arc::clone(&db.db),
             query_config: Arc::clone(&db.query_config),
@@ -110,8 +112,8 @@ impl<T> QueryCtx<T> {
     pub fn get(
         &self,
         query_id: QueryId,
-        params: Arc<Any + Send + Sync + 'static>,
-    ) -> Arc<Any + Send + Sync + 'static> {
+        params: D,
+    ) -> D {
         let (res, output_fingerprint) = self.get_inner(query_id, params);
         self.record_dep(query_id, output_fingerprint);
         res
@@ -120,8 +122,8 @@ impl<T> QueryCtx<T> {
     pub fn get_inner(
         &self,
         query_id: QueryId,
-        params: Arc<Any + Send + Sync + 'static>,
-    ) -> (Arc<Any + Send + Sync + 'static>, OutputFingerprint) {
+        params: D,
+    ) -> (D, OutputFingerprint) {
         let (gen, record) = {
             let guard = self.db.graph.lock();
             match guard.get(&query_id).map(|it| it.clone()){
@@ -139,7 +141,7 @@ impl<T> QueryCtx<T> {
             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 dep_params: D = {
                 let guard = self.db.graph.lock();
                 guard[&dep_query_id]
                 .1
@@ -160,29 +162,29 @@ impl<T> QueryCtx<T> {
     fn force(
         &self,
         query_id: QueryId,
-        params: Arc<Any + Send + Sync + 'static>,
-    ) -> (Arc<Any + Send + Sync + 'static>, OutputFingerprint) {
+        params: D,
+    ) -> (D, 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)
+            f(&self.db.ground_data, &params)
         } else if let Some(f) = self.query_fn_by_type(query_id.0) {
-            f(self, &*params)
+            f(self, &params)
         } else {
             panic!("unknown query type: {:?}", query_id.0);
         };
 
-        let res: Arc<Any + Send + Sync + 'static> = res.into();
+        let res: D = 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>> {
+    fn ground_query_fn_by_type(&self, query_type: QueryTypeId) -> Option<GroundQueryFn<T, D>> {
         self.query_config.ground_fn.get(&query_type).map(|&it| it)
     }
-    fn query_fn_by_type(&self, query_type: QueryTypeId) -> Option<QueryFn<T>> {
+    fn query_fn_by_type(&self, query_type: QueryTypeId) -> Option<QueryFn<T, D>> {
         self.query_config.query_fn.get(&query_type).map(|&it| it)
     }
     fn record_dep(
@@ -196,15 +198,18 @@ impl<T> QueryCtx<T> {
     }
 }
 
-impl<T> Db<T> {
-    pub fn new(query_config: QueryConfig<T>, ground_data: T) -> Db<T> {
+impl<T, D> Db<T, D>
+where
+    D: Clone
+{
+    pub fn new(query_config: QueryConfig<T, D>, ground_data: T) -> Db<T, D> {
         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> {
+    pub fn with_ground_data(&self, ground_data: T) -> Db<T, D> {
         let gen = Gen(self.db.gen.0 + 1);
         let graph = self.db.graph.lock().clone();
         let graph = Mutex::new(graph);
@@ -216,8 +221,8 @@ impl<T> Db<T> {
     pub fn get(
         &self,
         query_id: QueryId,
-        params: Box<Any + Send + Sync + 'static>,
-    ) -> (Arc<Any + Send + Sync + 'static>, Vec<QueryTypeId>) {
+        params: D,
+    ) -> (D, 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());
diff --git a/crates/salsa/tests/integration.rs b/crates/salsa/tests/integration.rs
index 7241eca38..2872d3913 100644
--- a/crates/salsa/tests/integration.rs
+++ b/crates/salsa/tests/integration.rs
@@ -7,6 +7,7 @@ use std::{
 };
 
 type State = HashMap<u32, String>;
+type Data = Arc<Any + Send + Sync + 'static>;
 const GET_TEXT: salsa::QueryTypeId = salsa::QueryTypeId(1);
 const GET_FILES: salsa::QueryTypeId = salsa::QueryTypeId(2);
 const FILE_NEWLINES: salsa::QueryTypeId = salsa::QueryTypeId(3);
@@ -14,9 +15,9 @@ const TOTAL_NEWLINES: salsa::QueryTypeId = salsa::QueryTypeId(4);
 
 fn mk_ground_query<T, R>(
     state: &State,
-    params: &(Any + Send + Sync + 'static),
+    params: &Data,
     f: fn(&State, &T) -> R,
-) -> (Box<Any + Send + Sync + 'static>, salsa::OutputFingerprint)
+) -> (Data, salsa::OutputFingerprint)
 where
     T: 'static,
     R: Hash + Send + Sync + 'static,
@@ -24,21 +25,21 @@ where
     let params = params.downcast_ref().unwrap();
     let result = f(state, params);
     let fingerprint = o_print(&result);
-    (Box::new(result), fingerprint)
+    (Arc::new(result), fingerprint)
 }
 
-fn get<T, R>(db: &salsa::Db<State>, query_type: salsa::QueryTypeId, param: T) -> (Arc<R>, Vec<salsa::QueryTypeId>)
+fn get<T, R>(db: &salsa::Db<State, Data>, 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 param = Arc::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>);
+struct QueryCtx<'a>(&'a salsa::QueryCtx<State, Data>);
 
 impl<'a> QueryCtx<'a> {
     fn get_text(&self, id: u32) -> Arc<String> {
@@ -60,10 +61,10 @@ impl<'a> QueryCtx<'a> {
 }
 
 fn mk_query<T, R>(
-    query_ctx: &salsa::QueryCtx<State>,
-    params: &(Any + Send + Sync + 'static),
+    query_ctx: &salsa::QueryCtx<State, Data>,
+    params: &Data,
     f: fn(QueryCtx, &T) -> R,
-) -> (Box<Any + Send + Sync + 'static>, salsa::OutputFingerprint)
+) -> (Data, salsa::OutputFingerprint)
 where
     T: 'static,
     R: Hash + Send + Sync + 'static,
@@ -72,11 +73,11 @@ where
     let query_ctx = QueryCtx(query_ctx);
     let result = f(query_ctx, params);
     let fingerprint = o_print(&result);
-    (Box::new(result), fingerprint)
+    (Arc::new(result), fingerprint)
 }
 
-fn mk_queries() -> salsa::QueryConfig<State> {
-    salsa::QueryConfig::<State>::new()
+fn mk_queries() -> salsa::QueryConfig<State, Data> {
+    salsa::QueryConfig::<State, Data>::new()
         .with_ground_query(GET_TEXT, |state, id| {
             mk_ground_query::<u32, String>(state, id, |state, id| state[id].clone())
         })