github.com/MetalBlockchain/metalgo@v1.11.9/vms/components/avax/utxo_fetching.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package avax
     5  
     6  import (
     7  	"bytes"
     8  	"fmt"
     9  	"math"
    10  
    11  	"github.com/MetalBlockchain/metalgo/ids"
    12  	"github.com/MetalBlockchain/metalgo/utils"
    13  	"github.com/MetalBlockchain/metalgo/utils/set"
    14  
    15  	safemath "github.com/MetalBlockchain/metalgo/utils/math"
    16  )
    17  
    18  // GetBalance returns the current balance of [addrs]
    19  func GetBalance(db UTXOReader, addrs set.Set[ids.ShortID]) (uint64, error) {
    20  	utxos, err := GetAllUTXOs(db, addrs)
    21  	if err != nil {
    22  		return 0, fmt.Errorf("couldn't get UTXOs: %w", err)
    23  	}
    24  	balance := uint64(0)
    25  	for _, utxo := range utxos {
    26  		if out, ok := utxo.Out.(Amounter); ok {
    27  			balance, err = safemath.Add64(out.Amount(), balance)
    28  			if err != nil {
    29  				return 0, err
    30  			}
    31  		}
    32  	}
    33  	return balance, nil
    34  }
    35  
    36  func GetAllUTXOs(db UTXOReader, addrs set.Set[ids.ShortID]) ([]*UTXO, error) {
    37  	utxos, _, _, err := GetPaginatedUTXOs(
    38  		db,
    39  		addrs,
    40  		ids.ShortEmpty,
    41  		ids.Empty,
    42  		math.MaxInt,
    43  	)
    44  	return utxos, err
    45  }
    46  
    47  // GetPaginatedUTXOs returns UTXOs such that at least one of the addresses in
    48  // [addrs] is referenced.
    49  //
    50  // Returns at most [limit] UTXOs.
    51  //
    52  // Only returns UTXOs associated with addresses >= [startAddr].
    53  //
    54  // For address [startAddr], only returns UTXOs whose IDs are greater than
    55  // [startUTXOID].
    56  //
    57  // Returns:
    58  // * The fetched UTXOs
    59  // * The address associated with the last UTXO fetched
    60  // * The ID of the last UTXO fetched
    61  func GetPaginatedUTXOs(
    62  	db UTXOReader,
    63  	addrs set.Set[ids.ShortID],
    64  	lastAddr ids.ShortID,
    65  	lastUTXOID ids.ID,
    66  	limit int,
    67  ) ([]*UTXO, ids.ShortID, ids.ID, error) {
    68  	var (
    69  		utxos      []*UTXO
    70  		seen       set.Set[ids.ID] // IDs of UTXOs already in the list
    71  		searchSize = limit         // the limit diminishes which can impact the expected return
    72  		addrsList  = addrs.List()
    73  	)
    74  	utils.Sort(addrsList) // enforces the same ordering for pagination
    75  	for _, addr := range addrsList {
    76  		start := ids.Empty
    77  		if comp := bytes.Compare(addr.Bytes(), lastAddr.Bytes()); comp == -1 { // Skip addresses before [startAddr]
    78  			continue
    79  		} else if comp == 0 {
    80  			start = lastUTXOID
    81  		}
    82  
    83  		lastAddr = addr // The last address searched
    84  
    85  		utxoIDs, err := db.UTXOIDs(addr.Bytes(), start, searchSize) // Get UTXOs associated with [addr]
    86  		if err != nil {
    87  			return nil, ids.ShortID{}, ids.Empty, fmt.Errorf("couldn't get UTXOs for address %s: %w", addr, err)
    88  		}
    89  		for _, utxoID := range utxoIDs {
    90  			lastUTXOID = utxoID // The last searched UTXO - not the last found
    91  
    92  			if seen.Contains(utxoID) { // Already have this UTXO in the list
    93  				continue
    94  			}
    95  
    96  			utxo, err := db.GetUTXO(utxoID)
    97  			if err != nil {
    98  				return nil, ids.ShortID{}, ids.Empty, fmt.Errorf("couldn't get UTXO %s: %w", utxoID, err)
    99  			}
   100  
   101  			utxos = append(utxos, utxo)
   102  			seen.Add(utxoID)
   103  			limit--
   104  			if limit <= 0 {
   105  				return utxos, lastAddr, lastUTXOID, nil // Found [limit] utxos; stop.
   106  			}
   107  		}
   108  	}
   109  	return utxos, lastAddr, lastUTXOID, nil // Didn't reach the [limit] utxos; no more were found
   110  }