Merge pull request AleoNet#1466 from ljedrz/lexicographic_memorymap

Use a BTreeMap instead of IndexMap in MemoryMap

synthesizer/Cargo.toml

@@ -88,6 +88,9 @@ default-features = false
 [dependencies.anyhow]
 version = "1.0.70"
 
+[dependencies.bincode]
+version = "1"
+
 [dependencies.blake2]
 version = "0.10"
 default-features = false

synthesizer/src/store/helpers/memory_map.rs

@@ -19,10 +19,10 @@ use console::network::prelude::*;
 use indexmap::IndexMap;
 
 use core::{borrow::Borrow, hash::Hash};
-use indexmap::map;
 use parking_lot::{Mutex, RwLock};
 use std::{
     borrow::Cow,
+    collections::{btree_map, BTreeMap},
     sync::{
         atomic::{AtomicBool, Ordering},
         Arc,
@@ -34,7 +34,10 @@ pub struct MemoryMap<
     K: Copy + Clone + PartialEq + Eq + Hash + Serialize + for<'de> Deserialize<'de> + Send + Sync,
     V: Clone + PartialEq + Eq + Serialize + for<'de> Deserialize<'de> + Send + Sync,
 > {
-    map: Arc<RwLock<IndexMap<K, V>>>,
+    // The reason for using BTreeMap with binary keys is for the order of items to be the same as
+    // the one in the RocksDB-backed DataMap; if not for that, it could be any map
+    // with fast lookups and the keys could be typed (i.e. just `K` instead of `Vec<u8>`).
+    map: Arc<RwLock<BTreeMap<Vec<u8>, V>>>,
     batch_in_progress: Arc<AtomicBool>,
     atomic_batch: Arc<Mutex<IndexMap<K, Option<V>>>>,
 }
@@ -56,8 +59,12 @@ impl<
 {
     /// Initializes a new `MemoryMap` from the given iterator.
     fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
+        // Serialize each key in the iterator, and collect them into a map.
+        // Note: The 'unwrap' is safe here, because the keys are defined by us.
+        let map = iter.into_iter().map(|(k, v)| (bincode::serialize(&k).unwrap(), v)).collect();
+        // Return the new map.
         Self {
-            map: Arc::new(RwLock::new(IndexMap::from_iter(iter))),
+            map: Arc::new(RwLock::new(map)),
             batch_in_progress: Default::default(),
             atomic_batch: Default::default(),
         }
@@ -84,9 +91,10 @@ impl<
             }
             // Otherwise, insert the key-value pair directly into the map.
             false => {
-                self.map.write().insert(key, value);
+                self.map.write().insert(bincode::serialize(&key)?, value);
             }
         }
+
         Ok(())
     }
 
@@ -104,9 +112,10 @@ impl<
             }
             // Otherwise, remove the key-value pair directly from the map.
             false => {
-                self.map.write().shift_remove(key);
+                self.map.write().remove(&bincode::serialize(&key)?);
             }
         }
+
         Ok(())
     }
 
@@ -150,11 +159,20 @@ impl<
         if !operations.is_empty() {
             // Acquire a write lock on the map.
             let mut locked_map = self.map.write();
+
+            // Prepare the key for each queued operation.
+            // Note: This step is taken to ensure (with 100% certainty) that there will be
+            // no chance to fail partway through committing the queued operations.
+            let prepared_operations = operations
+                .into_iter()
+                .map(|(key, value)| Ok((bincode::serialize(&key)?, value)))
+                .collect::<Result<Vec<_>>>()?;
+
             // Perform all the queued operations.
-            for operation in operations {
-                match operation {
-                    (key, Some(value)) => locked_map.insert(key, value),
-                    (key, None) => locked_map.remove(&key),
+            for (key, value) in prepared_operations {
+                match value {
+                    Some(value) => locked_map.insert(key, value),
+                    None => locked_map.remove(&key),
                 };
             }
         }
@@ -172,9 +190,9 @@ impl<
     V: 'a + Clone + PartialEq + Eq + Serialize + for<'de> Deserialize<'de> + Send + Sync,
 > MapRead<'a, K, V> for MemoryMap<K, V>
 {
-    type Iterator = core::iter::Map<map::IntoIter<K, V>, fn((K, V)) -> (Cow<'a, K>, Cow<'a, V>)>;
-    type Keys = core::iter::Map<map::IntoKeys<K, V>, fn(K) -> Cow<'a, K>>;
-    type Values = core::iter::Map<map::IntoValues<K, V>, fn(V) -> Cow<'a, V>>;
+    type Iterator = core::iter::Map<btree_map::IntoIter<Vec<u8>, V>, fn((Vec<u8>, V)) -> (Cow<'a, K>, Cow<'a, V>)>;
+    type Keys = core::iter::Map<btree_map::IntoKeys<Vec<u8>, V>, fn(Vec<u8>) -> Cow<'a, K>>;
+    type Values = core::iter::Map<btree_map::IntoValues<Vec<u8>, V>, fn(V) -> Cow<'a, V>>;
 
     ///
     /// Returns `true` if the given key exists in the map.
@@ -184,7 +202,7 @@ impl<
         K: Borrow<Q>,
         Q: PartialEq + Eq + Hash + Serialize + ?Sized,
     {
-        Ok(self.map.read().contains_key(key))
+        Ok(self.map.read().contains_key(&bincode::serialize(key)?))
     }
 
     ///
@@ -195,7 +213,7 @@ impl<
         K: Borrow<Q>,
         Q: PartialEq + Eq + Hash + Serialize + ?Sized,
     {
-        Ok(self.map.read().get(key).cloned().map(Cow::Owned))
+        Ok(self.map.read().get(&bincode::serialize(key)?).cloned().map(Cow::Owned))
     }
 
     ///
@@ -219,14 +237,16 @@ impl<
     /// Returns an iterator visiting each key-value pair in the map.
     ///
     fn iter(&'a self) -> Self::Iterator {
-        self.map.read().clone().into_iter().map(|(k, v)| (Cow::Owned(k), Cow::Owned(v)))
+        // Note: The 'unwrap' is safe here, because the keys are defined by us.
+        self.map.read().clone().into_iter().map(|(k, v)| (Cow::Owned(bincode::deserialize(&k).unwrap()), Cow::Owned(v)))
     }
 
     ///
     /// Returns an iterator over each key in the map.
     ///
     fn keys(&'a self) -> Self::Keys {
-        self.map.read().clone().into_keys().map(Cow::Owned)
+        // Note: The 'unwrap' is safe here, because the keys are defined by us.
+        self.map.read().clone().into_keys().map(|k| Cow::Owned(bincode::deserialize(&k).unwrap()))
     }
 
     ///
@@ -242,7 +262,7 @@ impl<
     V: Clone + PartialEq + Eq + Serialize + for<'de> Deserialize<'de> + Send + Sync,
 > Deref for MemoryMap<K, V>
 {
-    type Target = Arc<RwLock<IndexMap<K, V>>>;
+    type Target = Arc<RwLock<BTreeMap<Vec<u8>, V>>>;
 
     fn deref(&self) -> &Self::Target {
         &self.map