[commonware-storage/mutable] Make mutable MMR use the memory efficient index (#680)

Author: roberto-bayardo

storage/src/index/mod.rs

@@ -205,4 +205,24 @@ mod tests {
         values.sort();
         assert_eq!(values, vec![0, 2]);
     }
+
+    #[test_traced]
+    fn test_index_mutate_through_iterator() {
+        let (_, context, _) = Executor::default();
+        let mut index = Index::init(context.clone(), TwoCap);
+
+        index.insert(b"key", 1);
+        index.insert(b"key", 2);
+        index.insert(b"key", 3);
+
+        for value in index.get_mut(b"key") {
+            // Mutate the value
+            *value += 10;
+        }
+
+        assert_eq!(
+            index.get(b"key").copied().collect::<Vec<_>>(),
+            vec![11, 13, 12]
+        );
+    }
 }

storage/src/index/storage.rs

@@ -18,13 +18,25 @@ pub struct ValueIterator<'a, V> {
     next: Option<&'a Record<V>>,
 }
 
+// An iterator over all values associated with a translated key, allowing for mutation.
+pub struct MutableValueIterator<'a, V> {
+    next: Option<&'a mut Record<V>>,
+}
+
 impl<V> ValueIterator<'_, V> {
     /// Create a `ValueIterator` that returns no items.
     fn empty() -> Self {
         ValueIterator { next: None }
     }
 }
 
+impl<V> MutableValueIterator<'_, V> {
+    /// Create a `MutableValueIterator` that returns no items.
+    fn empty() -> Self {
+        MutableValueIterator { next: None }
+    }
+}
+
 impl<'a, V> Iterator for ValueIterator<'a, V> {
     type Item = &'a V;
 
@@ -40,10 +52,29 @@ impl<'a, V> Iterator for ValueIterator<'a, V> {
     }
 }
 
+impl<'a, V> Iterator for MutableValueIterator<'a, V> {
+    type Item = &'a mut V;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.next.take() {
+            Some(next) => {
+                let loc = &mut next.value;
+                self.next = next.next.as_deref_mut();
+                Some(loc)
+            }
+            None => None,
+        }
+    }
+}
+
 impl<V> Record<V> {
     fn iter(&self) -> ValueIterator<V> {
         ValueIterator { next: Some(self) }
     }
+
+    fn iter_mut(&mut self) -> MutableValueIterator<V> {
+        MutableValueIterator { next: Some(self) }
+    }
 }
 
 /// An index that maps translated keys to values.
@@ -113,6 +144,15 @@ impl<T: Translator, V> Index<T, V> {
         }
     }
 
+    /// Retrieve all values associated with a translated key, allowing for mutation.
+    pub fn get_mut(&mut self, key: &[u8]) -> MutableValueIterator<V> {
+        let translated_key = self.translator.transform(key);
+        match self.map.get_mut(&translated_key) {
+            Some(head) => head.iter_mut(),
+            None => MutableValueIterator::empty(),
+        }
+    }
+
     /// Remove values associated with the key that match the `prune` predicate.
     ///
     /// If this function is never called, the amount of memory used by old values will grow