SetMultimap implementation

vm/actor/src/lib.rs

@@ -18,6 +18,9 @@ use unsigned_varint::decode::Error as UVarintError;
 
 const HAMT_BIT_WIDTH: u8 = 5;
 
+type EmptyType = [u8; 0];
+const EMPTY_VALUE: EmptyType = [];
+
 /// Used when invocation requires parameters to be an empty array of bytes
 #[inline]
 fn check_empty_params(params: &Serialized) -> Result<(), EncodingError> {

vm/actor/src/util/mod.rs

@@ -3,6 +3,10 @@
 
 mod balance_table;
 mod multimap;
+mod set;
+mod set_multimap;
 
 pub use self::balance_table::BalanceTable;
 pub use self::multimap::*;
+pub use self::set::Set;
+pub use self::set_multimap::SetMultimap;

vm/actor/src/util/set.rs

@@ -0,0 +1,80 @@
+// Copyright 2020 ChainSafe Systems
+// SPDX-License-Identifier: Apache-2.0, MIT
+
+use crate::{EmptyType, EMPTY_VALUE, HAMT_BIT_WIDTH};
+use cid::Cid;
+use ipld_blockstore::BlockStore;
+use ipld_hamt::{Error, Hamt};
+
+/// Set is a Hamt with empty values for the purpose of acting as a hash set.
+#[derive(Debug)]
+pub struct Set<'a, BS>(Hamt<'a, String, BS>);
+
+impl<'a, BS: BlockStore> PartialEq for Set<'a, BS> {
+    fn eq(&self, other: &Self) -> bool {
+        self.0 == other.0
+    }
+}
+
+impl<'a, BS> Set<'a, BS>
+where
+    BS: BlockStore,
+{
+    /// Initializes a new empty Set.
+    pub fn new(bs: &'a BS) -> Self {
+        Self(Hamt::new_with_bit_width(bs, HAMT_BIT_WIDTH))
+    }
+
+    /// Initializes a Set from a root Cid.
+    pub fn from_root(bs: &'a BS, cid: &Cid) -> Result<Self, Error> {
+        Ok(Self(Hamt::load_with_bit_width(cid, bs, HAMT_BIT_WIDTH)?))
+    }
+
+    /// Retrieve root from the Set.
+    #[inline]
+    pub fn root(&mut self) -> Result<Cid, Error> {
+        self.0.flush()
+    }
+
+    /// Adds key to the set.
+    #[inline]
+    pub fn put(&mut self, key: String) -> Result<(), String> {
+        // Set hamt node to array root
+        Ok(self.0.set(key, EMPTY_VALUE)?)
+    }
+
+    /// Checks if key exists in the set.
+    #[inline]
+    pub fn has(&self, key: &str) -> Result<bool, String> {
+        Ok(self.0.get::<_, EmptyType>(key)?.is_some())
+    }
+
+    /// Deletes key from set.
+    #[inline]
+    pub fn delete(&mut self, key: &str) -> Result<(), String> {
+        self.0.delete(key)?;
+
+        Ok(())
+    }
+
+    /// Iterates through all keys in the set.
+    pub fn for_each<F>(&self, mut f: F) -> Result<(), String>
+    where
+        F: FnMut(&String) -> Result<(), String>,
+    {
+        // Calls the for each function on the hamt with ignoring the value
+        self.0.for_each(&mut |s, _: EmptyType| f(s))
+    }
+
+    /// Collects all keys from the set into a vector.
+    pub fn collect_keys(&self) -> Result<Vec<String>, String> {
+        let mut ret_keys = Vec::new();
+
+        self.for_each(|k| {
+            ret_keys.push(k.clone());
+            Ok(())
+        })?;
+
+        Ok(ret_keys)
+    }
+}

vm/actor/src/util/set_multimap.rs

@@ -0,0 +1,102 @@
+// Copyright 2020 ChainSafe Systems
+// SPDX-License-Identifier: Apache-2.0, MIT
+
+use super::Set;
+use crate::{deal_key, parse_uint_key, DealID, HAMT_BIT_WIDTH};
+use address::Address;
+use cid::Cid;
+use ipld_blockstore::BlockStore;
+use ipld_hamt::{Error, Hamt};
+
+/// SetMultimap is a hamt with values that are also a hamt but are of the set variant.
+/// This allows hash sets to be indexable by an address.
+pub struct SetMultimap<'a, BS>(Hamt<'a, String, BS>);
+impl<'a, BS> SetMultimap<'a, BS>
+where
+    BS: BlockStore,
+{
+    /// Initializes a new empty SetMultimap.
+    pub fn new(bs: &'a BS) -> Self {
+        Self(Hamt::new_with_bit_width(bs, HAMT_BIT_WIDTH))
+    }
+
+    /// Initializes a SetMultimap from a root Cid.
+    pub fn from_root(bs: &'a BS, cid: &Cid) -> Result<Self, Error> {
+        Ok(Self(Hamt::load_with_bit_width(cid, bs, HAMT_BIT_WIDTH)?))
+    }
+
+    /// Retrieve root from the SetMultimap.
+    #[inline]
+    pub fn root(&mut self) -> Result<Cid, Error> {
+        self.0.flush()
+    }
+
+    /// Puts the DealID in the hash set of the key.
+    pub fn put(&mut self, key: &Address, value: DealID) -> Result<(), String> {
+        // Get construct amt from retrieved cid or create new
+        let mut set = self.get(key)?.unwrap_or_else(|| Set::new(self.0.store()));
+
+        set.put(deal_key(value))?;
+
+        // Save and calculate new root
+        let new_root = set.root()?;
+
+        // Set hamt node to set new root
+        Ok(self.0.set(key.hash_key(), &new_root)?)
+    }
+
+    /// Gets the set at the given index of the `SetMultimap`
+    #[inline]
+    pub fn get(&self, key: &Address) -> Result<Option<Set<'a, BS>>, String> {
+        match self.0.get(&key.hash_key())? {
+            Some(cid) => Ok(Some(Set::from_root(self.0.store(), &cid)?)),
+            None => Ok(None),
+        }
+    }
+
+    /// Removes a DealID from a key hash set.
+    #[inline]
+    pub fn remove(&mut self, key: &Address, v: DealID) -> Result<(), String> {
+        // Get construct amt from retrieved cid and return if no set exists
+        let mut set = match self.get(key)? {
+            Some(s) => s,
+            None => return Ok(()),
+        };
+
+        set.delete(&deal_key(v))?;
+
+        // Save and calculate new root
+        let new_root = set.root()?;
+
+        Ok(self.0.set(key.hash_key(), &new_root)?)
+    }
+
+    /// Removes set at index.
+    #[inline]
+    pub fn remove_all(&mut self, key: &Address) -> Result<(), String> {
+        // Remove entry from table
+        self.0.delete(&key.hash_key())?;
+
+        Ok(())
+    }
+
+    /// Iterates through keys and converts them to a DealID to call a function on each.
+    pub fn for_each<F>(&self, key: &Address, mut f: F) -> Result<(), String>
+    where
+        F: FnMut(DealID) -> Result<(), String>,
+    {
+        // Get construct amt from retrieved cid and return if no set exists
+        let set = match self.get(key)? {
+            Some(s) => s,
+            None => return Ok(()),
+        };
+
+        set.for_each(|k| {
+            let v =
+                parse_uint_key(k).map_err(|e| format!("Could not parse key: {}, ({})", k, e))?;
+
+            // Run function on all parsed keys
+            f(v)
+        })
+    }
+}

vm/actor/tests/set_multimap_test.rs

@@ -0,0 +1,48 @@
+// Copyright 2020 ChainSafe Systems
+// SPDX-License-Identifier: Apache-2.0, MIT
+
+use actor::{deal_key, SetMultimap};
+use address::Address;
+
+#[test]
+fn put_remove() {
+    let store = db::MemoryDB::default();
+    let mut smm = SetMultimap::new(&store);
+
+    let addr = Address::new_id(100).unwrap();
+    assert_eq!(smm.get(&addr), Ok(None));
+
+    smm.put(&addr, 8).unwrap();
+    smm.put(&addr, 2).unwrap();
+    smm.remove(&addr, 2).unwrap();
+
+    let set = smm.get(&addr).unwrap().unwrap();
+    assert_eq!(set.has(&deal_key(8)), Ok(true));
+    assert_eq!(set.has(&deal_key(2)), Ok(false));
+
+    smm.remove_all(&addr).unwrap();
+    assert_eq!(smm.get(&addr), Ok(None));
+}
+
+#[test]
+fn for_each() {
+    let store = db::MemoryDB::default();
+    let mut smm = SetMultimap::new(&store);
+
+    let addr = Address::new_id(100).unwrap();
+    assert_eq!(smm.get(&addr), Ok(None));
+
+    smm.put(&addr, 8).unwrap();
+    smm.put(&addr, 3).unwrap();
+    smm.put(&addr, 2).unwrap();
+    smm.put(&addr, 8).unwrap();
+
+    let mut vals: Vec<u64> = Vec::new();
+    smm.for_each(&addr, |i| {
+        vals.push(i);
+        Ok(())
+    })
+    .unwrap();
+
+    assert_eq!(vals.len(), 3);
+}

vm/actor/tests/set_test.rs

@@ -0,0 +1,47 @@
+// Copyright 2020 ChainSafe Systems
+// SPDX-License-Identifier: Apache-2.0, MIT
+
+use actor::Set;
+
+#[test]
+fn put() {
+    let store = db::MemoryDB::default();
+    let mut set = Set::new(&store);
+
+    let key = "test";
+    assert_eq!(set.has(&key), Ok(false));
+
+    set.put(key.to_owned()).unwrap();
+    assert_eq!(set.has(&key), Ok(true));
+}
+
+#[test]
+fn collect_keys() {
+    let store = db::MemoryDB::default();
+    let mut set = Set::new(&store);
+
+    set.put("0".to_owned()).unwrap();
+
+    assert_eq!(set.collect_keys().unwrap(), ["0".to_owned()]);
+
+    set.put("1".to_owned()).unwrap();
+    set.put("2".to_owned()).unwrap();
+    set.put("3".to_owned()).unwrap();
+
+    assert_eq!(set.collect_keys().unwrap().len(), 4);
+}
+
+#[test]
+fn delete() {
+    let store = db::MemoryDB::default();
+    let mut set = Set::new(&store);
+
+    assert_eq!(set.has(&"0"), Ok(false));
+    set.put("0".to_owned()).unwrap();
+    assert_eq!(set.has(&"0"), Ok(true));
+    set.delete(&"0").unwrap();
+    assert_eq!(set.has(&"0"), Ok(false));
+
+    // Test delete when doesn't exist doesn't error
+    set.delete(&"0").unwrap();
+}