Use Map2 instead of Map and raw HAMTs in market actor (#1358)

Author: anorth

actors/init/src/state.rs

@@ -7,7 +7,9 @@ use fvm_ipld_encoding::tuple::*;
 use fvm_shared::address::{Address, Protocol};
 use fvm_shared::ActorID;
 
-use fil_actors_runtime::{actor_error, ActorError, Map2, DEFAULT_CONF, FIRST_NON_SINGLETON_ADDR};
+use fil_actors_runtime::{
+    actor_error, ActorError, Map2, DEFAULT_HAMT_CONFIG, FIRST_NON_SINGLETON_ADDR,
+};
 
 #[derive(Serialize_tuple, Deserialize_tuple, Clone, Debug)]
 pub struct State {
@@ -21,7 +23,7 @@ pub type AddressMap<BS> = Map2<BS, Address, ActorID>;
 
 impl State {
     pub fn new<BS: Blockstore>(store: &BS, network_name: String) -> Result<Self, ActorError> {
-        let empty = AddressMap::flush_empty(store, DEFAULT_CONF)?;
+        let empty = AddressMap::flush_empty(store, DEFAULT_HAMT_CONFIG)?;
         Ok(Self { address_map: empty, next_id: FIRST_NON_SINGLETON_ADDR, network_name })
     }
 
@@ -37,7 +39,7 @@ impl State {
         robust_addr: &Address,
         delegated_addr: Option<&Address>,
     ) -> Result<(ActorID, bool), ActorError> {
-        let mut map = AddressMap::load(store, &self.address_map, DEFAULT_CONF, "addresses")?;
+        let mut map = AddressMap::load(store, &self.address_map, DEFAULT_HAMT_CONFIG, "addresses")?;
         let (id, existing) = if let Some(delegated_addr) = delegated_addr {
             // If there's a delegated address, either recall the already-mapped actor ID or
             // create and map a new one.
@@ -87,7 +89,7 @@ impl State {
         if addr.protocol() == Protocol::ID {
             return Ok(Some(*addr));
         }
-        let map = AddressMap::load(store, &self.address_map, DEFAULT_CONF, "addresses")?;
+        let map = AddressMap::load(store, &self.address_map, DEFAULT_HAMT_CONFIG, "addresses")?;
         let found = map.get(addr)?;
         Ok(found.copied().map(Address::new_id))
     }

actors/init/src/testing.rs

@@ -6,7 +6,7 @@ use fvm_shared::{
     ActorID,
 };
 
-use fil_actors_runtime::{MessageAccumulator, DEFAULT_CONF, FIRST_NON_SINGLETON_ADDR};
+use fil_actors_runtime::{MessageAccumulator, DEFAULT_HAMT_CONFIG, FIRST_NON_SINGLETON_ADDR};
 
 use crate::state::AddressMap;
 use crate::State;
@@ -34,7 +34,7 @@ pub fn check_state_invariants<BS: Blockstore>(
     let mut stable_address_by_id = HashMap::<ActorID, Address>::new();
     let mut delegated_address_by_id = HashMap::<ActorID, Address>::new();
 
-    match AddressMap::load(store, &state.address_map, DEFAULT_CONF, "addresses") {
+    match AddressMap::load(store, &state.address_map, DEFAULT_HAMT_CONFIG, "addresses") {
         Ok(address_map) => {
             let ret = address_map.for_each(|key, actor_id| {
                 acc.require(key.protocol() != Protocol::ID, format!("key {key} is an ID address"));

actors/market/src/balance_table.rs

@@ -3,57 +3,64 @@
 
 use cid::Cid;
 use fvm_ipld_blockstore::Blockstore;
-use fvm_ipld_hamt::Error as HamtError;
 use fvm_shared::address::Address;
 use fvm_shared::econ::TokenAmount;
 use num_traits::Zero;
 
-use fil_actors_runtime::{make_empty_map, make_map_with_root_and_bitwidth, Map};
-
-pub const BALANCE_TABLE_BITWIDTH: u32 = 6;
+use fil_actors_runtime::{
+    actor_error, ActorContext, ActorError, Config, Map2, DEFAULT_HAMT_CONFIG,
+};
 
 /// Balance table which handles getting and updating token balances specifically
-pub struct BalanceTable<'a, BS>(pub Map<'a, BS, TokenAmount>);
+pub struct BalanceTable<BS: Blockstore>(pub Map2<BS, Address, TokenAmount>);
+
+const CONF: Config = Config { bit_width: 6, ..DEFAULT_HAMT_CONFIG };
 
-impl<'a, BS> BalanceTable<'a, BS>
+impl<BS> BalanceTable<BS>
 where
     BS: Blockstore,
 {
     /// Initializes a new empty balance table
-    pub fn new(bs: &'a BS) -> Self {
-        Self(make_empty_map(bs, BALANCE_TABLE_BITWIDTH))
+    pub fn new(bs: BS, name: &'static str) -> Self {
+        Self(Map2::empty(bs, CONF, name))
     }
 
     /// Initializes a balance table from a root Cid
-    pub fn from_root(bs: &'a BS, cid: &Cid) -> Result<Self, HamtError> {
-        Ok(Self(make_map_with_root_and_bitwidth(cid, bs, BALANCE_TABLE_BITWIDTH)?))
+    pub fn from_root(bs: BS, cid: &Cid, name: &'static str) -> Result<Self, ActorError> {
+        Ok(Self(Map2::load(bs, cid, CONF, name)?))
     }
 
     /// Retrieve root from balance table
-    pub fn root(&mut self) -> Result<Cid, HamtError> {
+    pub fn root(&mut self) -> Result<Cid, ActorError> {
         self.0.flush()
     }
 
     /// Gets token amount for given address in balance table
-    pub fn get(&self, key: &Address) -> Result<TokenAmount, HamtError> {
-        if let Some(v) = self.0.get(&key.to_bytes())? {
+    pub fn get(&self, key: &Address) -> Result<TokenAmount, ActorError> {
+        if let Some(v) = self.0.get(key)? {
             Ok(v.clone())
         } else {
             Ok(TokenAmount::zero())
         }
     }
 
     /// Adds token amount to previously initialized account.
-    pub fn add(&mut self, key: &Address, value: &TokenAmount) -> Result<(), HamtError> {
+    pub fn add(&mut self, key: &Address, value: &TokenAmount) -> Result<(), ActorError> {
         let prev = self.get(key)?;
         let sum = &prev + value;
         if sum.is_negative() {
-            Err(format!("New balance in table cannot be negative: {}", sum).into())
+            Err(actor_error!(
+                illegal_argument,
+                "negative balance for {} adding {} to {}",
+                key,
+                value,
+                prev
+            ))
         } else if sum.is_zero() && !prev.is_zero() {
-            self.0.delete(&key.to_bytes())?;
+            self.0.delete(key).context("adding balance")?;
             Ok(())
         } else {
-            self.0.set(key.to_bytes().into(), sum)?;
+            self.0.set(key, sum).context("adding balance")?;
             Ok(())
         }
     }
@@ -66,34 +73,39 @@ where
         key: &Address,
         req: &TokenAmount,
         floor: &TokenAmount,
-    ) -> Result<TokenAmount, HamtError> {
+    ) -> Result<TokenAmount, ActorError> {
         let prev = self.get(key)?;
         let available = std::cmp::max(TokenAmount::zero(), prev - floor);
         let sub: TokenAmount = std::cmp::min(&available, req).clone();
 
         if sub.is_positive() {
-            self.add(key, &-sub.clone())?;
+            self.add(key, &-sub.clone()).context("subtracting balance")?;
         }
 
         Ok(sub)
     }
 
     /// Subtracts value from a balance, and errors if full amount was not substracted.
-    pub fn must_subtract(&mut self, key: &Address, req: &TokenAmount) -> Result<(), HamtError> {
+    pub fn must_subtract(&mut self, key: &Address, req: &TokenAmount) -> Result<(), ActorError> {
         let prev = self.get(key)?;
 
         if req > &prev {
-            Err("couldn't subtract the requested amount".into())
+            Err(actor_error!(
+                illegal_argument,
+                "negative balance for {} subtracting {} from {}",
+                key,
+                req,
+                prev
+            ))
         } else {
             self.add(key, &-req)
         }
     }
 
     /// Returns total balance held by this balance table
     #[allow(dead_code)]
-    pub fn total(&self) -> Result<TokenAmount, HamtError> {
+    pub fn total(&self) -> Result<TokenAmount, ActorError> {
         let mut total = TokenAmount::zero();
-
         self.0.for_each(|_, v: &TokenAmount| {
             total += v;
             Ok(())
@@ -116,7 +128,7 @@ mod tests {
         let addr1 = Address::new_id(100);
         let addr2 = Address::new_id(101);
         let store = MemoryBlockstore::default();
-        let mut bt = BalanceTable::new(&store);
+        let mut bt = BalanceTable::new(&store, "test");
 
         assert!(bt.total().unwrap().is_zero());
 
@@ -143,7 +155,7 @@ mod tests {
     fn balance_subtracts() {
         let addr = Address::new_id(100);
         let store = MemoryBlockstore::default();
-        let mut bt = BalanceTable::new(&store);
+        let mut bt = BalanceTable::new(&store, "test");
 
         bt.add(&addr, &TokenAmount::from_atto(80u8)).unwrap();
         assert_eq!(bt.get(&addr).unwrap(), TokenAmount::from_atto(80u8));

actors/market/src/lib.rs

@@ -176,16 +176,10 @@ impl Actor {
 
         let store = rt.store();
         let st: State = rt.state()?;
-        let balances = BalanceTable::from_root(store, &st.escrow_table)
-            .context_code(ExitCode::USR_ILLEGAL_STATE, "failed to load escrow table")?;
-        let locks = BalanceTable::from_root(store, &st.locked_table)
-            .context_code(ExitCode::USR_ILLEGAL_STATE, "failed to load locked table")?;
-        let balance = balances
-            .get(&account)
-            .context_code(ExitCode::USR_ILLEGAL_STATE, "failed to get escrow balance")?;
-        let locked = locks
-            .get(&account)
-            .context_code(ExitCode::USR_ILLEGAL_STATE, "failed to get locked balance")?;
+        let balances = BalanceTable::from_root(store, &st.escrow_table, "escrow table")?;
+        let locks = BalanceTable::from_root(store, &st.locked_table, "locled table")?;
+        let balance = balances.get(&account)?;
+        let locked = locks.get(&account)?;
 
         Ok(GetBalanceReturn { balance, locked })
     }
@@ -427,7 +421,7 @@ impl Actor {
             let mut pending_deals: Vec<Cid> = vec![];
             let mut deal_proposals: Vec<(DealID, DealProposal)> = vec![];
             let mut deals_by_epoch: Vec<(ChainEpoch, DealID)> = vec![];
-            let mut pending_deal_allocation_ids: Vec<(BytesKey, AllocationID)> = vec![];
+            let mut pending_deal_allocation_ids: Vec<(DealID, AllocationID)> = vec![];
 
             // All storage dealProposals will be added in an atomic transaction; this operation will be unrolled if any of them fails.
             // This should only fail on programmer error because all expected invalid conditions should be filtered in the first set of checks.
@@ -444,7 +438,7 @@ impl Actor {
                 // Store verified allocation (if any) in the pending allocation IDs map.
                 // It will be removed when the deal is activated or expires.
                 if let Some(alloc_id) = deal_allocation_ids.get(&valid_deal.cid) {
-                    pending_deal_allocation_ids.push((deal_id_key(deal_id), *alloc_id));
+                    pending_deal_allocation_ids.push((deal_id, *alloc_id));
                 }
 
                 // Randomize the first epoch for when the deal will be processed so an attacker isn't able to
@@ -619,13 +613,12 @@ impl Actor {
 
                         // Extract and remove any verified allocation ID for the pending deal.
                         let allocation = st
-                            .remove_pending_deal_allocation_id(rt.store(), &deal_id_key(deal_id))
+                            .remove_pending_deal_allocation_id(rt.store(), deal_id)
                             .context(format!(
                                 "failed to remove pending deal allocation id {}",
                                 deal_id
                             ))?
-                            .unwrap_or((BytesKey(vec![]), NO_ALLOCATION_ID))
-                            .1;
+                            .unwrap_or(NO_ALLOCATION_ID);
 
                         if allocation != NO_ALLOCATION_ID {
                             verified_infos.push(VerifiedDealInfo {
@@ -795,7 +788,7 @@ impl Actor {
                         })?;
 
                         // Delete pending deal allocation id (if present).
-                        st.remove_pending_deal_allocation_id(rt.store(), &deal_id_key(deal_id))?;
+                        st.remove_pending_deal_allocation_id(rt.store(), deal_id)?;
 
                         continue;
                     }