eliminate 'cloned' call on hydration; recover parallel hydration on L…

…EM store
argumentcomputer · Sep 26, 2023 · 9c416ad · 9c416ad
1 parent e85a6fd
commit 9c416ad
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 38 deletions.
diff --git a/src/lem/store.rs b/src/lem/store.rs
@@ -3,6 +3,7 @@ use arc_swap::ArcSwap;
 use elsa::sync::{FrozenMap, FrozenVec};
 use elsa::sync_index_set::FrozenIndexSet;
 use nom::{sequence::preceded, Parser};
+use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 use std::{cell::RefCell, rc::Rc, sync::Arc};
 
 use crate::{
@@ -51,9 +52,8 @@ pub struct Store<F: LurkField> {
 
     pub poseidon_cache: PoseidonCache<F>,
 
-    dehydrated: FrozenVec<Box<Ptr<F>>>,
+    dehydrated: ArcSwap<FrozenVec<Box<Ptr<F>>>>,
     z_cache: FrozenMap<Ptr<F>, Box<ZPtr<F>>>,
-    next_hydration_idx: ArcSwap<usize>,
 
     comms: FrozenMap<FWrap<F>, Box<(F, Ptr<F>)>>, // hash -> (secret, src)
 }
@@ -65,7 +65,7 @@ impl<F: LurkField> Store<F> {
         let ptr = Ptr::Tuple2(tag, idx);
         if inserted {
             // this is for `hydrate_z_cache`
-            self.dehydrated.push(Box::new(ptr));
+            self.dehydrated.load().push(Box::new(ptr));
         }
         ptr
     }
@@ -85,7 +85,7 @@ impl<F: LurkField> Store<F> {
         let ptr = Ptr::Tuple3(tag, idx);
         if inserted {
             // this is for `hydrate_z_cache`
-            self.dehydrated.push(Box::new(ptr));
+            self.dehydrated.load().push(Box::new(ptr));
         }
         ptr
     }
@@ -112,7 +112,7 @@ impl<F: LurkField> Store<F> {
         let ptr = Ptr::Tuple4(tag, idx);
         if inserted {
             // this is for `hydrate_z_cache`
-            self.dehydrated.push(Box::new(ptr));
+            self.dehydrated.load().push(Box::new(ptr));
         }
         ptr
     }
@@ -536,26 +536,15 @@ impl<F: LurkField> Store<F> {
     /// Hashes `Ptr` trees from the bottom to the top, avoiding deep recursions
     /// in `hash_ptr`.
     pub fn hydrate_z_cache(&self) {
-        let dehydrated_len = self.dehydrated.len();
-        // TODO: can we recover the parallel hydration?
-        for i in **self.next_hydration_idx.load()..dehydrated_len {
-            let ptr = self
-                .dehydrated
-                .get(i)
-                .expect("Out of bounds. Implementation broken");
-            self.hash_ptr(ptr).expect("failed to hash_expr");
-        }
-        self.next_hydration_idx.swap(Arc::new(dehydrated_len));
-    }
-
-    pub fn clear_dehydrated_queue(&mut self) {
-        self.dehydrated = FrozenVec::new();
-        self.next_hydration_idx.swap(Arc::new(0));
-    }
-
-    pub fn hydrate_z_cache_and_clear_queue(&mut self) {
-        self.hydrate_z_cache();
-        self.clear_dehydrated_queue();
+        self.dehydrated
+            .load()
+            .iter()
+            .collect::<Vec<_>>()
+            .par_iter()
+            .for_each(|ptr| {
+                self.hash_ptr(ptr).expect("failed to hash_ptr");
+            });
+        self.dehydrated.swap(Arc::new(FrozenVec::default()));
     }
 
     pub fn to_vector(&self, ptrs: &[Ptr<F>]) -> Result<Vec<F>> {

diff --git a/src/store.rs b/src/store.rs
@@ -1363,24 +1363,25 @@ impl<F: LurkField> Store<F> {
     pub fn hydrate_scalar_cache(&self) {
         self.ensure_constants();
 
-        let dehydrated = self.dehydrated.load().iter().cloned().collect::<Vec<_>>();
-
-        dehydrated.par_iter().for_each(|ptr| {
-            self.hash_expr(ptr).expect("failed to hash_expr");
-        });
+        self.dehydrated
+            .load()
+            .iter()
+            .collect::<Vec<_>>()
+            .par_iter()
+            .for_each(|ptr| {
+                self.hash_expr(ptr).expect("failed to hash_expr");
+            });
 
         self.dehydrated.swap(Arc::new(FrozenVec::default()));
 
-        let dehydrated_cont = self
-            .dehydrated_cont
+        self.dehydrated_cont
             .load()
             .iter()
-            .cloned()
-            .collect::<Vec<_>>();
-
-        dehydrated_cont.par_iter().for_each(|ptr| {
-            self.hash_cont(ptr).expect("failed to hash_cont");
-        });
+            .collect::<Vec<_>>()
+            .par_iter()
+            .for_each(|ptr| {
+                self.hash_cont(ptr).expect("failed to hash_cont");
+            });
 
         self.dehydrated_cont.swap(Arc::new(FrozenVec::default()));
     }