actors_runtime/util/
set_multimap.rs

1// Copyright 2019-2022 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use std::borrow::Borrow;
5
6use cid::Cid;
7use fvm_ipld_hamt::Error;
8use fvm_shared::blockstore::Blockstore;
9use fvm_shared::clock::ChainEpoch;
10use fvm_shared::deal::DealID;
11use fvm_shared::HAMT_BIT_WIDTH;
12
13use super::Set;
14use crate::{make_empty_map, make_map_with_root, parse_uint_key, u64_key, Map};
15
16/// SetMultimap is a hamt with values that are also a hamt but are of the set variant.
17/// This allows hash sets to be indexable by an address.
18pub struct SetMultimap<'a, BS>(Map<'a, BS, Cid>);
19impl<'a, BS> SetMultimap<'a, BS>
20where
21    BS: Blockstore,
22{
23    /// Initializes a new empty SetMultimap.
24    pub fn new(bs: &'a BS) -> Self {
25        Self(make_empty_map(bs, HAMT_BIT_WIDTH))
26    }
27
28    /// Initializes a SetMultimap from a root Cid.
29    pub fn from_root(bs: &'a BS, cid: &Cid) -> Result<Self, Error> {
30        Ok(Self(make_map_with_root(cid, bs)?))
31    }
32
33    /// Retrieve root from the SetMultimap.
34    #[inline]
35    pub fn root(&mut self) -> Result<Cid, Error> {
36        self.0.flush()
37    }
38
39    /// Puts the DealID in the hash set of the key.
40    pub fn put(&mut self, key: ChainEpoch, value: DealID) -> Result<(), Error> {
41        // Get construct amt from retrieved cid or create new
42        let mut set = self.get(key)?.unwrap_or_else(|| Set::new(self.0.store()));
43
44        set.put(u64_key(value))?;
45
46        // Save and calculate new root
47        let new_root = set.root()?;
48
49        // Set hamt node to set new root
50        self.0.set(u64_key(key as u64), new_root)?;
51        Ok(())
52    }
53
54    /// Puts slice of DealIDs in the hash set of the key.
55    pub fn put_many(&mut self, key: ChainEpoch, values: &[DealID]) -> Result<(), Error> {
56        // Get construct amt from retrieved cid or create new
57        let mut set = self.get(key)?.unwrap_or_else(|| Set::new(self.0.store()));
58
59        for &v in values {
60            set.put(u64_key(v))?;
61        }
62
63        // Save and calculate new root
64        let new_root = set.root()?;
65
66        // Set hamt node to set new root
67        self.0.set(u64_key(key as u64), new_root)?;
68        Ok(())
69    }
70
71    /// Gets the set at the given index of the `SetMultimap`
72    #[inline]
73    pub fn get(&self, key: ChainEpoch) -> Result<Option<Set<'a, BS>>, Error> {
74        match self.0.get(&u64_key(key as u64))? {
75            Some(cid) => Ok(Some(Set::from_root(*self.0.store(), cid)?)),
76            None => Ok(None),
77        }
78    }
79
80    /// Removes a DealID from a key hash set.
81    #[inline]
82    pub fn remove(&mut self, key: ChainEpoch, v: DealID) -> Result<(), Error> {
83        // Get construct amt from retrieved cid and return if no set exists
84        let mut set = match self.get(key)? {
85            Some(s) => s,
86            None => return Ok(()),
87        };
88
89        set.delete(u64_key(v).borrow())?;
90
91        // Save and calculate new root
92        let new_root = set.root()?;
93        self.0.set(u64_key(key as u64), new_root)?;
94        Ok(())
95    }
96
97    /// Removes set at index.
98    #[inline]
99    pub fn remove_all(&mut self, key: ChainEpoch) -> Result<(), Error> {
100        // Remove entry from table
101        self.0.delete(&u64_key(key as u64))?;
102
103        Ok(())
104    }
105
106    /// Iterates through keys and converts them to a DealID to call a function on each.
107    pub fn for_each<F>(&self, key: ChainEpoch, mut f: F) -> Result<(), Error>
108    where
109        F: FnMut(DealID) -> Result<(), Error>,
110    {
111        // Get construct amt from retrieved cid and return if no set exists
112        let set = match self.get(key)? {
113            Some(s) => s,
114            None => return Ok(()),
115        };
116
117        set.for_each(|k| {
118            let v = parse_uint_key(k)
119                .map_err(|e| anyhow::anyhow!("Could not parse key: {:?}, ({})", &k.0, e))?;
120
121            // Run function on all parsed keys
122            Ok(f(v)?)
123        })
124    }
125}