Fold key into qtree

Author: laurmaedje

src/cache.rs

@@ -159,7 +159,7 @@ impl<C: 'static, Out: 'static> Cache<C, Out> {
 /// The internal data for a cache.
 pub struct CacheData<C, Out> {
     /// Maps from hashes to memoized results.
-    entries: HashMap<u128, QuestionTree<C, u128, (Out, Vec<C>)>>,
+    tree: QuestionTree<C, u128, (Out, Vec<C>)>,
 }
 
 impl<C: PartialEq, Out: 'static> CacheData<C, Out> {
@@ -169,7 +169,7 @@ impl<C: PartialEq, Out: 'static> CacheData<C, Out> {
         In: Input<Call = C>,
         C: Clone + Hash,
     {
-        self.entries.get(&key)?.get(|c| input.call(c.clone()))
+        self.tree.get(key, |c| input.call(c.clone()))
     }
 
     /// Insert an entry into the cache.
@@ -184,15 +184,13 @@ impl<C: PartialEq, Out: 'static> CacheData<C, Out> {
         In: Input<Call = C>,
         C: Clone + Hash,
     {
-        self.entries
-            .entry(key)
-            .or_default()
-            .insert(recording.immutable, (output, recording.mutable))
+        self.tree
+            .insert(key, recording.immutable, (output, recording.mutable))
     }
 }
 
 impl<C, Out> Default for CacheData<C, Out> {
     fn default() -> Self {
-        Self { entries: HashMap::new() }
+        Self { tree: QuestionTree::new() }
     }
 }

src/qtree.rs

@@ -83,7 +83,7 @@ pub struct QuestionTree<Q, A, T> {
     questions: Slab<Q>,
     results: Slab<T>,
     links: HashMap<(usize, A), State>,
-    start: Option<State>,
+    start: HashMap<u128, State>,
 }
 
 impl<Q, A, T> QuestionTree<Q, A, T> {
@@ -92,7 +92,7 @@ impl<Q, A, T> QuestionTree<Q, A, T> {
             questions: Slab::new(),
             results: Slab::new(),
             links: HashMap::new(),
-            start: None,
+            start: HashMap::new(),
         }
     }
 }
@@ -102,8 +102,8 @@ where
     Q: Hash + Clone,
     A: Hash + Eq + Clone,
 {
-    pub fn get(&self, mut oracle: impl FnMut(&Q) -> A) -> Option<&T> {
-        let mut state = self.start?;
+    pub fn get(&self, key: u128, mut oracle: impl FnMut(&Q) -> A) -> Option<&T> {
+        let mut state = *self.start.get(&key)?;
         loop {
             match state.kind() {
                 StateKind::Result(r) => return Some(self.results.get(r).unwrap()),
@@ -118,18 +118,22 @@ where
 
     pub fn insert(
         &mut self,
+        key: u128,
         mut sequence: LookaheadSequence<Q, A>,
         value: T,
     ) -> Result<(), InsertError> {
-        let mut state = self.start;
+        let mut state = self.start.get(&key).copied();
         let mut predecessor = None;
 
         loop {
             let pair = if state.is_none() || predecessor.is_some() {
                 let Some((q, a)) = sequence.next() else { break };
                 let qi = self.questions.alloc(q);
                 let new = State::question(qi);
-                self.link(predecessor.take(), new);
+                if state.is_none() {
+                    self.start.insert(key, new);
+                }
+                self.link(&state, predecessor.take(), new);
                 state = Some(new);
                 (qi, a)
             } else {
@@ -150,21 +154,21 @@ where
             }
         }
 
-        if predecessor.is_none() && self.start.is_some() {
+        if predecessor.is_none() && state.is_some() {
             return Err(InsertError::AlreadyExists);
         }
 
         let ri = self.results.alloc(value);
-        self.link(predecessor, State::result(ri));
+        let target = State::result(ri);
+        if state.is_none() {
+            self.start.insert(key, target);
+        }
+        self.link(&state, predecessor, target);
 
         Ok(())
     }
 
-    fn link(&mut self, from: Option<(usize, A)>, target: State) {
-        if self.start.is_none() {
-            self.start = Some(target);
-        }
-
+    fn link(&mut self, state: &Option<State>, from: Option<(usize, A)>, target: State) {
         if let Some(pair) = from {
             self.links.insert(pair, target);
         }
@@ -263,36 +267,37 @@ mod tests {
     #[test]
     fn test_question_tree() {
         let mut tree = QuestionTree::<char, u128, &'static str>::new();
-        tree.insert(s([('a', 10), ('b', 15)]), "first").unwrap();
-        tree.insert(s([('a', 10), ('b', 20)]), "second").unwrap();
-        tree.insert(s([('a', 15), ('c', 15)]), "third").unwrap();
+        tree.insert(0, s([('a', 10), ('b', 15)]), "first").unwrap();
+        tree.insert(0, s([('a', 10), ('b', 20)]), "second").unwrap();
+        tree.insert(0, s([('a', 15), ('c', 15)]), "third").unwrap();
         assert_eq!(
-            tree.get(|&c| match c {
+            tree.get(0, |&c| match c {
                 'a' => 10,
                 'b' => 15,
                 _ => 20,
             }),
             Some(&"first")
         );
         assert_eq!(
-            tree.get(|&c| match c {
+            tree.get(0, |&c| match c {
                 'a' => 10,
                 _ => 20,
             }),
             Some(&"second")
         );
-        assert_eq!(tree.get(|_| 15), Some(&"third"));
-        assert_eq!(tree.get(|_| 10), None);
+        assert_eq!(tree.get(0, |_| 15), Some(&"third"));
+        assert_eq!(tree.get(0, |_| 10), None);
     }
 
     #[test]
     fn test_question_tree_pull_forward() {
         let mut tree = QuestionTree::<char, u128, &'static str>::new();
-        tree.insert(s([('a', 10), ('b', 15)]), "first").unwrap();
-        tree.insert(s([('a', 10), ('c', 15), ('b', 20)]), "second").unwrap();
-        tree.insert(s([('a', 15), ('b', 30), ('c', 15)]), "third").unwrap();
+        tree.insert(0, s([('a', 10), ('b', 15)]), "first").unwrap();
+        tree.insert(0, s([('a', 10), ('c', 15), ('b', 20)]), "second")
+            .unwrap();
+        tree.insert(0, s([('a', 15), ('b', 30), ('c', 15)]), "third").unwrap();
         assert_eq!(
-            tree.get(|&c| match c {
+            tree.get(0, |&c| match c {
                 'a' => 10,
                 'b' => 20,
                 'c' => 15,
@@ -301,7 +306,7 @@ mod tests {
             Some(&"second")
         );
         assert_eq!(
-            tree.get(|&c| match c {
+            tree.get(0, |&c| match c {
                 'a' => 15,
                 'b' => 30,
                 'c' => 15,
@@ -327,15 +332,15 @@ mod tests {
         let mut kept = Vec::new();
         for case in cases.iter() {
             let &(ref numbers, value) = case;
-            match tree.insert(s(sequence(numbers)), value) {
+            match tree.insert(0, s(sequence(numbers)), value) {
                 Ok(()) => kept.push(case),
                 Err(InsertError::AlreadyExists) => {}
                 Err(InsertError::WrongQuestion) => {} // Err(error) => panic!("{error:?}"),
             }
         }
         for (numbers, value) in kept {
             let map: HashMap<u64, u16> = sequence(numbers).collect();
-            assert_eq!(tree.get(|s| map[s]), Some(value));
+            assert_eq!(tree.get(0, |s| map[s]), Some(value));
         }
     }