actors_runtime/util/
balance_table.rs

1// Copyright 2019-2022 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use cid::Cid;
5use fvm_ipld_hamt::Error;
6use fvm_shared::address::Address;
7use fvm_shared::bigint::bigint_ser::BigIntDe;
8use fvm_shared::blockstore::Blockstore;
9use fvm_shared::econ::TokenAmount;
10use num_traits::{Signed, Zero};
11
12use crate::{make_empty_map, make_map_with_root_and_bitwidth, Map};
13
14pub const BALANCE_TABLE_BITWIDTH: u32 = 6;
15
16/// Balance table which handles getting and updating token balances specifically
17pub struct BalanceTable<'a, BS>(Map<'a, BS, BigIntDe>);
18impl<'a, BS> BalanceTable<'a, BS>
19where
20    BS: Blockstore,
21{
22    /// Initializes a new empty balance table
23    pub fn new(bs: &'a BS) -> Self {
24        Self(make_empty_map(bs, BALANCE_TABLE_BITWIDTH))
25    }
26
27    /// Initializes a balance table from a root Cid
28    pub fn from_root(bs: &'a BS, cid: &Cid) -> Result<Self, Error> {
29        Ok(Self(make_map_with_root_and_bitwidth(
30            cid,
31            bs,
32            BALANCE_TABLE_BITWIDTH,
33        )?))
34    }
35
36    /// Retrieve root from balance table
37    pub fn root(&mut self) -> Result<Cid, Error> {
38        self.0.flush()
39    }
40
41    /// Gets token amount for given address in balance table
42    pub fn get(&self, key: &Address) -> Result<TokenAmount, Error> {
43        if let Some(v) = self.0.get(&key.to_bytes())? {
44            Ok(v.0.clone())
45        } else {
46            Ok(0.into())
47        }
48    }
49
50    /// Adds token amount to previously initialized account.
51    pub fn add(&mut self, key: &Address, value: &TokenAmount) -> Result<(), Error> {
52        let prev = self.get(key)?;
53        let sum = &prev + value;
54        if sum.is_negative() {
55            Err(format!("New balance in table cannot be negative: {}", sum).into())
56        } else if sum.is_zero() && !prev.is_zero() {
57            self.0.delete(&key.to_bytes())?;
58            Ok(())
59        } else {
60            self.0.set(key.to_bytes().into(), BigIntDe(sum))?;
61            Ok(())
62        }
63    }
64
65    /// Subtracts up to the specified amount from a balance, without reducing the balance
66    /// below some minimum.
67    /// Returns the amount subtracted (always positive or zero).
68    pub fn subtract_with_minimum(
69        &mut self,
70        key: &Address,
71        req: &TokenAmount,
72        floor: &TokenAmount,
73    ) -> Result<TokenAmount, Error> {
74        let prev = self.get(key)?;
75        let available = std::cmp::max(TokenAmount::zero(), prev - floor);
76        let sub: TokenAmount = std::cmp::min(&available, req).clone();
77
78        if sub.is_positive() {
79            self.add(key, &-sub.clone())?;
80        }
81
82        Ok(sub)
83    }
84
85    /// Subtracts value from a balance, and errors if full amount was not substracted.
86    pub fn must_subtract(&mut self, key: &Address, req: &TokenAmount) -> Result<(), Error> {
87        let prev = self.get(key)?;
88
89        if req > &prev {
90            Err("couldn't subtract the requested amount".into())
91        } else {
92            self.add(key, &-req)
93        }
94    }
95
96    /// Returns total balance held by this balance table
97    pub fn total(&self) -> Result<TokenAmount, Error> {
98        let mut total = TokenAmount::default();
99
100        self.0.for_each(|_, v: &BigIntDe| {
101            total += &v.0;
102            Ok(())
103        })?;
104
105        Ok(total)
106    }
107}
actors_runtime/util/balance_table.rs

actors_runtime/util/
balance_table.rs
