github.com/MetalBlockchain/metalgo@v1.11.9/wallet/subnet/primary/common/utxos.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package common
     5  
     6  import (
     7  	"context"
     8  	"sync"
     9  
    10  	"golang.org/x/exp/maps"
    11  
    12  	"github.com/MetalBlockchain/metalgo/database"
    13  	"github.com/MetalBlockchain/metalgo/ids"
    14  	"github.com/MetalBlockchain/metalgo/vms/components/avax"
    15  )
    16  
    17  var (
    18  	_ UTXOs      = (*utxos)(nil)
    19  	_ ChainUTXOs = (*chainUTXOs)(nil)
    20  )
    21  
    22  type UTXOs interface {
    23  	AddUTXO(ctx context.Context, sourceChainID, destinationChainID ids.ID, utxo *avax.UTXO) error
    24  	RemoveUTXO(ctx context.Context, sourceChainID, destinationChainID, utxoID ids.ID) error
    25  
    26  	UTXOs(ctx context.Context, sourceChainID, destinationChainID ids.ID) ([]*avax.UTXO, error)
    27  	GetUTXO(ctx context.Context, sourceChainID, destinationChainID, utxoID ids.ID) (*avax.UTXO, error)
    28  }
    29  
    30  type ChainUTXOs interface {
    31  	AddUTXO(ctx context.Context, destinationChainID ids.ID, utxo *avax.UTXO) error
    32  	RemoveUTXO(ctx context.Context, sourceChainID, utxoID ids.ID) error
    33  
    34  	UTXOs(ctx context.Context, sourceChainID ids.ID) ([]*avax.UTXO, error)
    35  	GetUTXO(ctx context.Context, sourceChainID, utxoID ids.ID) (*avax.UTXO, error)
    36  }
    37  
    38  func NewUTXOs() UTXOs {
    39  	return &utxos{
    40  		sourceToDestToUTXOIDToUTXO: make(map[ids.ID]map[ids.ID]map[ids.ID]*avax.UTXO),
    41  	}
    42  }
    43  
    44  func NewChainUTXOs(chainID ids.ID, utxos UTXOs) ChainUTXOs {
    45  	return &chainUTXOs{
    46  		utxos:   utxos,
    47  		chainID: chainID,
    48  	}
    49  }
    50  
    51  type utxos struct {
    52  	lock sync.RWMutex
    53  	// sourceChainID -> destinationChainID -> utxoID -> utxo
    54  	sourceToDestToUTXOIDToUTXO map[ids.ID]map[ids.ID]map[ids.ID]*avax.UTXO
    55  }
    56  
    57  func (u *utxos) AddUTXO(_ context.Context, sourceChainID, destinationChainID ids.ID, utxo *avax.UTXO) error {
    58  	u.lock.Lock()
    59  	defer u.lock.Unlock()
    60  
    61  	destToUTXOIDToUTXO, ok := u.sourceToDestToUTXOIDToUTXO[sourceChainID]
    62  	if !ok {
    63  		destToUTXOIDToUTXO = make(map[ids.ID]map[ids.ID]*avax.UTXO)
    64  		u.sourceToDestToUTXOIDToUTXO[sourceChainID] = destToUTXOIDToUTXO
    65  	}
    66  
    67  	utxoIDToUTXO, ok := destToUTXOIDToUTXO[destinationChainID]
    68  	if !ok {
    69  		utxoIDToUTXO = make(map[ids.ID]*avax.UTXO)
    70  		destToUTXOIDToUTXO[destinationChainID] = utxoIDToUTXO
    71  	}
    72  
    73  	utxoIDToUTXO[utxo.InputID()] = utxo
    74  	return nil
    75  }
    76  
    77  func (u *utxos) RemoveUTXO(_ context.Context, sourceChainID, destinationChainID, utxoID ids.ID) error {
    78  	u.lock.Lock()
    79  	defer u.lock.Unlock()
    80  
    81  	destToUTXOIDToUTXO := u.sourceToDestToUTXOIDToUTXO[sourceChainID]
    82  	utxoIDToUTXO := destToUTXOIDToUTXO[destinationChainID]
    83  	_, ok := utxoIDToUTXO[utxoID]
    84  	if !ok {
    85  		return nil
    86  	}
    87  
    88  	delete(utxoIDToUTXO, utxoID)
    89  	if len(utxoIDToUTXO) != 0 {
    90  		return nil
    91  	}
    92  
    93  	delete(destToUTXOIDToUTXO, destinationChainID)
    94  	if len(destToUTXOIDToUTXO) != 0 {
    95  		return nil
    96  	}
    97  
    98  	delete(u.sourceToDestToUTXOIDToUTXO, sourceChainID)
    99  	return nil
   100  }
   101  
   102  func (u *utxos) UTXOs(_ context.Context, sourceChainID, destinationChainID ids.ID) ([]*avax.UTXO, error) {
   103  	u.lock.RLock()
   104  	defer u.lock.RUnlock()
   105  
   106  	destToUTXOIDToUTXO := u.sourceToDestToUTXOIDToUTXO[sourceChainID]
   107  	utxoIDToUTXO := destToUTXOIDToUTXO[destinationChainID]
   108  	return maps.Values(utxoIDToUTXO), nil
   109  }
   110  
   111  func (u *utxos) GetUTXO(_ context.Context, sourceChainID, destinationChainID, utxoID ids.ID) (*avax.UTXO, error) {
   112  	u.lock.RLock()
   113  	defer u.lock.RUnlock()
   114  
   115  	destToUTXOIDToUTXO := u.sourceToDestToUTXOIDToUTXO[sourceChainID]
   116  	utxoIDToUTXO := destToUTXOIDToUTXO[destinationChainID]
   117  	utxo, ok := utxoIDToUTXO[utxoID]
   118  	if !ok {
   119  		return nil, database.ErrNotFound
   120  	}
   121  	return utxo, nil
   122  }
   123  
   124  type chainUTXOs struct {
   125  	utxos   UTXOs
   126  	chainID ids.ID
   127  }
   128  
   129  func (c *chainUTXOs) AddUTXO(ctx context.Context, destinationChainID ids.ID, utxo *avax.UTXO) error {
   130  	return c.utxos.AddUTXO(ctx, c.chainID, destinationChainID, utxo)
   131  }
   132  
   133  func (c *chainUTXOs) RemoveUTXO(ctx context.Context, sourceChainID, utxoID ids.ID) error {
   134  	return c.utxos.RemoveUTXO(ctx, sourceChainID, c.chainID, utxoID)
   135  }
   136  
   137  func (c *chainUTXOs) UTXOs(ctx context.Context, sourceChainID ids.ID) ([]*avax.UTXO, error) {
   138  	return c.utxos.UTXOs(ctx, sourceChainID, c.chainID)
   139  }
   140  
   141  func (c *chainUTXOs) GetUTXO(ctx context.Context, sourceChainID, utxoID ids.ID) (*avax.UTXO, error) {
   142  	return c.utxos.GetUTXO(ctx, sourceChainID, c.chainID, utxoID)
   143  }