refactor(storage-proofs): move cache into global hash map

This allows the cache to be reused accross multiple graphs

Author: dignifiedquire

storage-proofs/src/zigzag_graph.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::marker::PhantomData;
 use std::sync::{Arc, RwLock};
 
@@ -11,75 +12,73 @@ use crate::settings;
 /// The expansion degree used for ZigZag Graphs.
 pub const EXP_DEGREE: usize = 8;
 
-// Cache of node's parents.
-pub type ParentCache = Vec<Option<Vec<u32>>>;
+lazy_static! {
+    // This parents cache is currently used for the *expanded parents only*, generated
+    // by the expensive Feistel operations in the ZigZag, it doesn't contain the
+    // "base" (in the `Graph` terminology) parents, which are cheaper to compute.
+    // It is indexed by the `Graph.identifier`, to ensure that the right cache is used.
+    static ref PARENT_CACHE: Arc<RwLock<HashMap<String, ParentCache>>> = Arc::new(RwLock::new(HashMap::new()));
+}
 
 // ZigZagGraph will hold two different (but related) `ParentCache`,
-// the first one for the `forward` direction and the second one
-// for the `reversed`.
+// the first one for the `forward` direction and the second one for the `reversed`.
 #[derive(Debug, Clone)]
-pub struct ShareableParentCache {
-    forward: Arc<RwLock<ParentCache>>,
-    reverse: Arc<RwLock<ParentCache>>,
+pub struct ParentCache {
+    forward: Vec<Option<Vec<u32>>>,
+    reverse: Vec<Option<Vec<u32>>>,
     // Keep the size of the cache outside the lock to be easily accessible.
     cache_entries: u32,
 }
 
-impl ShareableParentCache {
+impl ParentCache {
     pub fn new(cache_entries: u32) -> Self {
-        ShareableParentCache {
-            forward: Arc::new(RwLock::new(vec![None; cache_entries as usize])),
-            reverse: Arc::new(RwLock::new(vec![None; cache_entries as usize])),
+        ParentCache {
+            forward: vec![None; cache_entries as usize],
+            reverse: vec![None; cache_entries as usize],
             cache_entries,
         }
     }
 
     pub fn contains_forward(&self, node: u32) -> bool {
         assert!(node < self.cache_entries);
-        self.forward.read().unwrap()[node as usize].is_some()
+        self.forward[node as usize].is_some()
     }
 
     pub fn contains_reverse(&self, node: u32) -> bool {
         assert!(node < self.cache_entries);
-        self.reverse.read().unwrap()[node as usize].is_some()
+        self.reverse[node as usize].is_some()
     }
 
     pub fn read_forward<F, T>(&self, node: u32, mut cb: F) -> T
     where
         F: FnMut(Option<&Vec<u32>>) -> T,
     {
         assert!(node < self.cache_entries);
-        let lock = self.forward.read().unwrap();
-        cb(lock[node as usize].as_ref())
+        cb(self.forward[node as usize].as_ref())
     }
 
     pub fn read_reverse<F, T>(&self, node: u32, mut cb: F) -> T
     where
         F: FnMut(Option<&Vec<u32>>) -> T,
     {
         assert!(node < self.cache_entries);
-        let lock = self.reverse.read().unwrap();
-        cb(lock[node as usize].as_ref())
+        cb(self.reverse[node as usize].as_ref())
     }
 
-    pub fn write_forward(&self, node: u32, parents: Vec<u32>) {
+    pub fn write_forward(&mut self, node: u32, parents: Vec<u32>) {
         assert!(node < self.cache_entries);
 
-        let mut write_lock = self.forward.write().unwrap();
-
-        let old_value = std::mem::replace(&mut write_lock[node as usize], Some(parents));
+        let old_value = std::mem::replace(&mut self.forward[node as usize], Some(parents));
 
         debug_assert_eq!(old_value, None);
         // We shouldn't be rewriting entries (with most likely the same values),
         // this would be a clear indication of a bug.
     }
 
-    pub fn write_reverse(&self, node: u32, parents: Vec<u32>) {
+    pub fn write_reverse(&mut self, node: u32, parents: Vec<u32>) {
         assert!(node < self.cache_entries);
 
-        let mut write_lock = self.reverse.write().unwrap();
-
-        let old_value = std::mem::replace(&mut write_lock[node as usize], Some(parents));
+        let old_value = std::mem::replace(&mut self.reverse[node as usize], Some(parents));
 
         debug_assert_eq!(old_value, None);
         // We shouldn't be rewriting entries (with most likely the same values),
@@ -97,16 +96,8 @@ where
     base_graph: G,
     pub reversed: bool,
     feistel_precomputed: FeistelPrecomputed,
-
-    // This parents cache is currently used for the *expanded parents only*, generated
-    // by the expensive Feistel operations in the ZigZag, it doesn't contain the
-    // "base" (in the `Graph` terminology) parents, which are cheaper to compute.
-    // This is not an LRU cache, it holds the first `cache_entries` of the total
-    // possible `base_graph.size()` (the assumption here is that we either request
-    // all entries sequentially when encoding or any random entry once when proving
-    // or verifying, but there's no locality to take advantage of so keep the logic
-    // as simple as possible).
-    parents_cache: Option<ShareableParentCache>,
+    id: String,
+    use_cache: bool,
     _h: PhantomData<H>,
 }
 
@@ -115,14 +106,14 @@ pub type ZigZagBucketGraph<H> = ZigZagGraph<H, BucketGraph<H>>;
 impl<'a, H, G> Layerable<H> for ZigZagGraph<H, G>
 where
     H: Hasher,
-    G: Graph<H> + 'static,
+    G: Graph<H> + ParameterSetMetadata + 'static,
 {
 }
 
 impl<H, G> ZigZagGraph<H, G>
 where
     H: Hasher,
-    G: Graph<H>,
+    G: Graph<H> + ParameterSetMetadata,
 {
     pub fn new(
         base_graph: Option<G>,
@@ -136,25 +127,40 @@ where
             assert_eq!(expansion_degree, EXP_DEGREE);
         }
 
-        let parents_cache = if settings::SETTINGS.lock().unwrap().maximize_caching {
-            info!("using parents cache of unlimited size",);
-            assert!(nodes <= std::u32::MAX as usize);
-            Some(ShareableParentCache::new(nodes as u32))
-        } else {
-            None
-        };
+        let use_cache = settings::SETTINGS.lock().unwrap().maximize_caching;
 
-        ZigZagGraph {
-            base_graph: match base_graph {
-                Some(graph) => graph,
-                None => G::new(nodes, base_degree, 0, seed),
-            },
+        let base_graph = match base_graph {
+            Some(graph) => graph,
+            None => G::new(nodes, base_degree, 0, seed),
+        };
+        let bg_id = base_graph.identifier();
+
+        let res = ZigZagGraph {
+            base_graph,
+            id: format!(
+                "zigzag_graph::ZigZagGraph{{expansion_degree: {} base_graph: {} }}",
+                expansion_degree, bg_id,
+            ),
             expansion_degree,
+            use_cache,
             reversed: false,
             feistel_precomputed: feistel::precompute((expansion_degree * nodes) as feistel::Index),
-            parents_cache,
             _h: PhantomData,
+        };
+
+        if use_cache {
+            info!("using parents cache of unlimited size",);
+            assert!(nodes <= std::u32::MAX as usize);
+
+            if !PARENT_CACHE.read().unwrap().contains_key(&res.id) {
+                PARENT_CACHE
+                    .write()
+                    .unwrap()
+                    .insert(res.id.clone(), ParentCache::new(nodes as u32));
+            }
         }
+
+        res
     }
 }
 
@@ -164,11 +170,7 @@ where
     G: Graph<H> + ParameterSetMetadata,
 {
     fn identifier(&self) -> String {
-        format!(
-            "zigzag_graph::ZigZagGraph{{expansion_degree: {} base_graph: {} }}",
-            self.expansion_degree,
-            self.base_graph.identifier()
-        )
+        self.id.clone()
     }
 
     fn sector_size(&self) -> u64 {
@@ -263,7 +265,7 @@ impl<Z: ZigZag> Graph<Z::BaseHasher> for Z {
 impl<'a, H, G> ZigZagGraph<H, G>
 where
     H: Hasher,
-    G: Graph<H>,
+    G: Graph<H> + ParameterSetMetadata,
 {
     // Assign `expansion_degree` parents to `node` using an invertible function. That
     // means we can't just generate random values between `[0, size())`, we need to
@@ -329,22 +331,45 @@ where
     // the current direction set in the graph and return a copy of it (or
     // `None` to signal a cache miss).
     fn contains_parents_cache(&self, node: usize) -> bool {
-        if let Some(ref cache) = self.parents_cache {
-            if self.forward() {
-                cache.contains_forward(node as u32)
+        if self.use_cache {
+            if let Some(ref cache) = PARENT_CACHE.read().unwrap().get(&self.id) {
+                if self.forward() {
+                    cache.contains_forward(node as u32)
+                } else {
+                    cache.contains_reverse(node as u32)
+                }
             } else {
-                cache.contains_reverse(node as u32)
+                false
             }
         } else {
             false
         }
     }
+
+    fn generate_expanded_parents(&self, node: usize) -> Vec<u32> {
+        (0..self.expansion_degree)
+            .filter_map(|i| {
+                let other = self.correspondent(node, i);
+                if self.reversed {
+                    if other > node {
+                        Some(other as u32)
+                    } else {
+                        None
+                    }
+                } else if other < node {
+                    Some(other as u32)
+                } else {
+                    None
+                }
+            })
+            .collect()
+    }
 }
 
 impl<'a, H, G> ZigZag for ZigZagGraph<H, G>
 where
     H: Hasher,
-    G: Graph<H>,
+    G: Graph<H> + ParameterSetMetadata,
 {
     type BaseHasher = H;
     type BaseGraph = G;
@@ -395,43 +420,37 @@ where
     where
         F: FnMut(&Vec<u32>) -> T,
     {
-        if !self.contains_parents_cache(node) {
-            let parents: Vec<u32> = (0..self.expansion_degree)
-                .filter_map(|i| {
-                    let other = self.correspondent(node, i);
-                    if self.reversed {
-                        if other > node {
-                            Some(other as u32)
-                        } else {
-                            None
-                        }
-                    } else if other < node {
-                        Some(other as u32)
-                    } else {
-                        None
-                    }
-                })
-                .collect();
-
-            if let Some(ref cache) = self.parents_cache {
-                if self.forward() {
-                    cache.write_forward(node as u32, parents);
-                } else {
-                    cache.write_reverse(node as u32, parents);
-                }
-            } else {
-                return cb(&parents);
-            }
+        if !self.use_cache {
+            // No cache usage, generate on demand.
+            return cb(&self.generate_expanded_parents(node));
         }
 
-        if let Some(ref cache) = self.parents_cache {
+        // Check if we need to fill the cache.
+        if !self.contains_parents_cache(node) {
+            // Cache is empty so we need to generate the parents.
+            let parents = self.generate_expanded_parents(node);
+
+            // Store the newly generated cached value.
+            let mut cache_lock = PARENT_CACHE.write().unwrap();
+            let cache = cache_lock
+                .get_mut(&self.id)
+                .expect("Invalid cache construction");
             if self.forward() {
-                cache.read_forward(node as u32, |parents| cb(parents.unwrap()))
+                cache.write_forward(node as u32, parents);
             } else {
-                cache.read_reverse(node as u32, |parents| cb(parents.unwrap()))
+                cache.write_reverse(node as u32, parents);
             }
+        }
+
+        // We made sure the cache is filled above, now we can return the value.
+        let cache_lock = PARENT_CACHE.read().unwrap();
+        let cache = cache_lock
+            .get(&self.id)
+            .expect("Invalid cache construction");
+        if self.forward() {
+            cache.read_forward(node as u32, |parents| cb(parents.unwrap()))
         } else {
-            unreachable!()
+            cache.read_reverse(node as u32, |parents| cb(parents.unwrap()))
         }
     }