github.com/filecoin-project/specs-actors/v4@v4.0.2/actors/util/adt/balancetable.go (about)

     1  package adt
     2  
     3  import (
     4  	addr "github.com/filecoin-project/go-address"
     5  	"github.com/filecoin-project/go-state-types/abi"
     6  	"github.com/filecoin-project/go-state-types/big"
     7  	cid "github.com/ipfs/go-cid"
     8  	"golang.org/x/xerrors"
     9  )
    10  
    11  // Bitwidth of balance table HAMTs, determined empirically from mutation
    12  // patterns and projections of mainnet data
    13  const BalanceTableBitwidth = 6
    14  
    15  // A specialization of a map of addresses to (positive) token amounts.
    16  // Absent keys implicitly have a balance of zero.
    17  type BalanceTable Map
    18  
    19  // Interprets a store as balance table with root `r`.
    20  func AsBalanceTable(s Store, r cid.Cid) (*BalanceTable, error) {
    21  	m, err := AsMap(s, r, BalanceTableBitwidth)
    22  	if err != nil {
    23  		return nil, err
    24  	}
    25  
    26  	return &BalanceTable{
    27  		root:  m.root,
    28  		store: s,
    29  	}, nil
    30  }
    31  
    32  // Returns the root cid of underlying HAMT.
    33  func (t *BalanceTable) Root() (cid.Cid, error) {
    34  	return (*Map)(t).Root()
    35  }
    36  
    37  // Gets the balance for a key, which is zero if they key has never been added to.
    38  func (t *BalanceTable) Get(key addr.Address) (abi.TokenAmount, error) {
    39  	var value abi.TokenAmount
    40  	found, err := (*Map)(t).Get(abi.AddrKey(key), &value)
    41  	if !found || err != nil {
    42  		value = big.Zero()
    43  	}
    44  
    45  	return value, err
    46  }
    47  
    48  // Adds an amount to a balance, requiring the resulting balance to be non-negative.
    49  func (t *BalanceTable) Add(key addr.Address, value abi.TokenAmount) error {
    50  	prev, err := t.Get(key)
    51  	if err != nil {
    52  		return err
    53  	}
    54  	sum := big.Add(prev, value)
    55  	sign := sum.Sign()
    56  	if sign < 0 {
    57  		return xerrors.Errorf("adding %v to balance %v would give negative: %v", value, prev, sum)
    58  	} else if sign == 0 && !prev.IsZero() {
    59  		return (*Map)(t).Delete(abi.AddrKey(key))
    60  	}
    61  	return (*Map)(t).Put(abi.AddrKey(key), &sum)
    62  }
    63  
    64  // Subtracts up to the specified amount from a balance, without reducing the balance below some minimum.
    65  // Returns the amount subtracted.
    66  func (t *BalanceTable) SubtractWithMinimum(key addr.Address, req abi.TokenAmount, floor abi.TokenAmount) (abi.TokenAmount, error) {
    67  	prev, err := t.Get(key)
    68  	if err != nil {
    69  		return big.Zero(), err
    70  	}
    71  
    72  	available := big.Max(big.Zero(), big.Sub(prev, floor))
    73  	sub := big.Min(available, req)
    74  	if sub.Sign() > 0 {
    75  		err = t.Add(key, sub.Neg())
    76  		if err != nil {
    77  			return big.Zero(), err
    78  		}
    79  	}
    80  	return sub, nil
    81  }
    82  
    83  // MustSubtract subtracts the given amount from the account's balance.
    84  // Returns an error if the account has insufficient balance
    85  func (t *BalanceTable) MustSubtract(key addr.Address, req abi.TokenAmount) error {
    86  	prev, err := t.Get(key)
    87  	if err != nil {
    88  		return err
    89  	}
    90  	if req.GreaterThan(prev) {
    91  		return xerrors.New("couldn't subtract the requested amount")
    92  	}
    93  	return t.Add(key, req.Neg())
    94  }
    95  
    96  // Returns the total balance held by this BalanceTable
    97  func (t *BalanceTable) Total() (abi.TokenAmount, error) {
    98  	total := big.Zero()
    99  	var cur abi.TokenAmount
   100  	err := (*Map)(t).ForEach(&cur, func(key string) error {
   101  		total = big.Add(total, cur)
   102  		return nil
   103  	})
   104  	return total, err
   105  }