Merge pull request #49 from iotaledger/feat/use-shrinking-maps

Use shrinking maps

pkg/evilwallet/aliasmanager.go

@@ -4,16 +4,18 @@ import (
 	"go.uber.org/atomic"
 
 	"github.com/iotaledger/evil-tools/pkg/models"
+	"github.com/iotaledger/hive.go/ds/shrinkingmap"
 	"github.com/iotaledger/hive.go/ierrors"
+	"github.com/iotaledger/hive.go/lo"
 	"github.com/iotaledger/hive.go/runtime/syncutils"
 )
 
 // region AliasManager /////////////////////////////////////////////////////////////////////////////////////////////////
 
 // AliasManager is the manager for output aliases.
 type AliasManager struct {
-	outputMap map[string]*models.OutputData
-	inputMap  map[string]*models.OutputData
+	outputMap *shrinkingmap.ShrinkingMap[string, *models.OutputData]
+	inputMap  *shrinkingmap.ShrinkingMap[string, *models.OutputData]
 
 	outputAliasCount *atomic.Uint64
 	mu               syncutils.RWMutex
@@ -22,8 +24,8 @@ type AliasManager struct {
 // NewAliasManager creates and returns a new AliasManager.
 func NewAliasManager() *AliasManager {
 	return &AliasManager{
-		outputMap:        make(map[string]*models.OutputData),
-		inputMap:         make(map[string]*models.OutputData),
+		outputMap:        shrinkingmap.New[string, *models.OutputData](),
+		inputMap:         shrinkingmap.New[string, *models.OutputData](),
 		outputAliasCount: atomic.NewUint64(0),
 	}
 }
@@ -33,22 +35,23 @@ func (a *AliasManager) AddOutputAlias(output *models.OutputData, aliasName strin
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
-	a.outputMap[aliasName] = output
+	a.outputMap.Set(aliasName, output)
 }
 
 // AddInputAlias adds an input alias.
 func (a *AliasManager) AddInputAlias(input *models.OutputData, aliasName string) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
-	a.inputMap[aliasName] = input
+	a.inputMap.Set(aliasName, input)
 }
 
 // GetInput returns the input for the alias specified.
 func (a *AliasManager) GetInput(aliasName string) (*models.OutputData, bool) {
 	a.mu.RLock()
 	defer a.mu.RUnlock()
-	in, ok := a.inputMap[aliasName]
+
+	in, ok := a.inputMap.Get(aliasName)
 
 	return in, ok
 }
@@ -58,16 +61,16 @@ func (a *AliasManager) GetOutput(aliasName string) *models.OutputData {
 	a.mu.RLock()
 	defer a.mu.RUnlock()
 
-	return a.outputMap[aliasName]
+	return lo.Return1(a.outputMap.Get(aliasName))
 }
 
 // ClearAllAliases clears all aliases.
 func (a *AliasManager) ClearAllAliases() {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
-	a.inputMap = make(map[string]*models.OutputData)
-	a.outputMap = make(map[string]*models.OutputData)
+	a.inputMap.Clear()
+	a.outputMap.Clear()
 }
 
 // ClearAliases clears provided aliases.
@@ -76,10 +79,10 @@ func (a *AliasManager) ClearAliases(aliases ScenarioAlias) {
 	defer a.mu.Unlock()
 
 	for _, in := range aliases.Inputs {
-		delete(a.inputMap, in)
+		a.inputMap.Delete(in)
 	}
 	for _, out := range aliases.Outputs {
-		delete(a.outputMap, out)
+		a.outputMap.Delete(out)
 	}
 }
 

pkg/evilwallet/evilwallet.go

@@ -641,13 +641,6 @@ func (e *EvilWallet) prepareFlatOptionsForAccountScenario(scenario *EvilScenario
 	}, allAliases
 }
 
-// SetTxOutputsSolid marks all outputs as solid in OutputManager for clientID.
-func (e *EvilWallet) SetTxOutputsSolid(outputs iotago.OutputIDs, clientID string) {
-	for _, out := range outputs {
-		e.outputManager.SetOutputIDSolidForIssuer(out, clientID)
-	}
-}
-
 // endregion ///////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 func WithClients(urls ...string) options.Option[EvilWallet] {

pkg/evilwallet/options.go

@@ -68,15 +68,13 @@ func NewOptions(options ...Option) (option *Options, err error) {
 type Option func(*Options)
 
 func (o *Options) isBalanceProvided() bool {
-	provided := false
-
 	for _, output := range o.aliasOutputs {
 		if output.BaseTokenAmount() > 0 {
-			provided = true
+			return true
 		}
 	}
 
-	return provided
+	return false
 }
 
 func (o *Options) isWalletProvidedForInputsOutputs() error {

pkg/evilwallet/output_manager.go

@@ -8,7 +8,7 @@ import (
 
 	"github.com/iotaledger/evil-tools/pkg/models"
 	"github.com/iotaledger/evil-tools/pkg/utils"
-	"github.com/iotaledger/hive.go/ds/types"
+	"github.com/iotaledger/hive.go/ds/shrinkingmap"
 	"github.com/iotaledger/hive.go/lo"
 	"github.com/iotaledger/hive.go/log"
 	"github.com/iotaledger/hive.go/runtime/syncutils"
@@ -21,22 +21,18 @@ type OutputManager struct {
 	connector models.Connector
 
 	wallets         *Wallets
-	tempIDWalletMap map[models.TempOutputID]*Wallet
-
-	// stores solid outputs per node
-	issuerSolidOutIDMap map[string]map[iotago.OutputID]types.Empty
+	tempIDWalletMap *shrinkingmap.ShrinkingMap[models.TempOutputID, *Wallet]
 
 	syncutils.RWMutex
 }
 
 // NewOutputManager creates an OutputManager instance. All outputs are mapped based on their address, so address should never be reused.
 func NewOutputManager(connector models.Connector, wallets *Wallets, logger log.Logger) *OutputManager {
 	return &OutputManager{
-		Logger:              logger,
-		connector:           connector,
-		wallets:             wallets,
-		tempIDWalletMap:     make(map[models.TempOutputID]*Wallet),
-		issuerSolidOutIDMap: make(map[string]map[iotago.OutputID]types.Empty),
+		Logger:          logger,
+		connector:       connector,
+		wallets:         wallets,
+		tempIDWalletMap: shrinkingmap.New[models.TempOutputID, *Wallet](),
 	}
 }
 
@@ -45,40 +41,15 @@ func (o *OutputManager) setTempIDWalletMap(id models.TempOutputID, wallet *Walle
 	o.Lock()
 	defer o.Unlock()
 
-	o.tempIDWalletMap[id] = wallet
+	o.tempIDWalletMap.Set(id, wallet)
 }
 
 // TempIDWalletMap returns wallet corresponding to the address stored in OutputManager.
 func (o *OutputManager) TempIDWalletMap(outputID models.TempOutputID) *Wallet {
 	o.RLock()
 	defer o.RUnlock()
 
-	return o.tempIDWalletMap[outputID]
-}
-
-// SetOutputIDSolidForIssuer sets solid flag for the provided outputID and issuer.
-func (o *OutputManager) SetOutputIDSolidForIssuer(outputID iotago.OutputID, issuer string) {
-	o.Lock()
-	defer o.Unlock()
-
-	if _, ok := o.issuerSolidOutIDMap[issuer]; !ok {
-		o.issuerSolidOutIDMap[issuer] = make(map[iotago.OutputID]types.Empty)
-	}
-	o.issuerSolidOutIDMap[issuer][outputID] = types.Void
-}
-
-// IssuerSolidOutIDMap checks whether output was marked as solid for a given node.
-func (o *OutputManager) IssuerSolidOutIDMap(issuer string, outputID iotago.OutputID) (isSolid bool) {
-	o.RLock()
-	defer o.RUnlock()
-
-	if solidOutputs, ok := o.issuerSolidOutIDMap[issuer]; ok {
-		if _, isSolid = solidOutputs[outputID]; isSolid {
-			return
-		}
-	}
-
-	return
+	return lo.Return1(o.tempIDWalletMap.Get(outputID))
 }
 
 // Track the confirmed statuses of the given outputIDs, it returns true if all of them are confirmed.
@@ -164,7 +135,7 @@ func (o *OutputManager) getOutputFromWallet(id models.TempOutputID) (output *mod
 	o.RLock()
 	defer o.RUnlock()
 
-	w, ok := o.tempIDWalletMap[id]
+	w, ok := o.tempIDWalletMap.Get(id)
 	if ok {
 		output = w.UnspentOutput(id)
 	}

pkg/evilwallet/wallet.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/iotaledger/evil-tools/pkg/models"
 	"github.com/iotaledger/evil-tools/pkg/walletmanager"
+	"github.com/iotaledger/hive.go/ds/shrinkingmap"
 	"github.com/iotaledger/hive.go/ds/types"
 	"github.com/iotaledger/hive.go/lo"
 	"github.com/iotaledger/hive.go/runtime/syncutils"
@@ -21,11 +22,11 @@ import (
 type Wallet struct {
 	ID                walletID
 	walletType        WalletType
-	unspentOutputs    map[models.TempOutputID]*models.OutputData // maps addr to its unspentOutput
-	indexTempIDMap    map[uint32]models.TempOutputID
-	addrIndexMap      map[string]uint32
-	inputTransactions map[string]types.Empty
-	reuseTempIDPool   map[models.TempOutputID]types.Empty
+	unspentOutputs    *shrinkingmap.ShrinkingMap[models.TempOutputID, *models.OutputData] // maps addr to its unspentOutput
+	indexTempIDMap    *shrinkingmap.ShrinkingMap[uint32, models.TempOutputID]
+	addrIndexMap      *shrinkingmap.ShrinkingMap[string, uint32]
+	inputTransactions *shrinkingmap.ShrinkingMap[string, types.Empty]
+	reuseTempIDPool   *shrinkingmap.ShrinkingMap[models.TempOutputID, types.Empty]
 	seed              [32]byte
 
 	nextAddrIdxToUse atomic.Uint32 // used during filling in wallet with new outputs
@@ -47,19 +48,16 @@ func NewWallet(wType ...WalletType) *Wallet {
 		walletType:        walletType,
 		ID:                -1,
 		seed:              tpkg.RandEd25519Seed(),
-		unspentOutputs:    make(map[models.TempOutputID]*models.OutputData),
-		indexTempIDMap:    make(map[uint32]models.TempOutputID),
-		addrIndexMap:      make(map[string]uint32),
-		inputTransactions: make(map[string]types.Empty),
+		unspentOutputs:    shrinkingmap.New[models.TempOutputID, *models.OutputData](),
+		indexTempIDMap:    shrinkingmap.New[uint32, models.TempOutputID](),
+		addrIndexMap:      shrinkingmap.New[string, uint32](),
+		inputTransactions: shrinkingmap.New[string, types.Empty](),
+		reuseTempIDPool:   shrinkingmap.New[models.TempOutputID, types.Empty](),
 		nextAddrToSpent:   *idxSpent,
 		nextAddrIdxToUse:  *addrUsed,
 		RWMutex:           &syncutils.RWMutex{},
 	}
 
-	if walletType == Reuse {
-		w.reuseTempIDPool = make(map[models.TempOutputID]types.Empty)
-	}
-
 	return w
 }
 
@@ -78,7 +76,7 @@ func (w *Wallet) Address() *iotago.Ed25519Address {
 	keyManager := lo.PanicOnErr(wallet.NewKeyManager(w.seed[:], walletmanager.BIP32PathForIndex(index)))
 	//nolint:forcetypeassert
 	addr := keyManager.Address(iotago.AddressEd25519).(*iotago.Ed25519Address)
-	w.addrIndexMap[addr.String()] = index
+	w.addrIndexMap.Set(addr.String(), index)
 
 	return addr
 }
@@ -100,17 +98,19 @@ func (w *Wallet) UnspentOutput(id models.TempOutputID) *models.OutputData {
 	w.RLock()
 	defer w.RUnlock()
 
-	return w.unspentOutputs[id]
+	return lo.Return1(w.unspentOutputs.Get(id))
 }
 
 // UnspentOutputs returns all unspent outputs on the wallet.
 func (w *Wallet) UnspentOutputs() (outputs map[models.TempOutputID]*models.OutputData) {
 	w.RLock()
 	defer w.RUnlock()
 	outputs = make(map[models.TempOutputID]*models.OutputData)
-	for addr, outs := range w.unspentOutputs {
-		outputs[addr] = outs
-	}
+
+	w.unspentOutputs.ForEach(func(key models.TempOutputID, value *models.OutputData) bool {
+		outputs[key] = value
+		return true
+	})
 
 	return outputs
 }
@@ -120,27 +120,27 @@ func (w *Wallet) IndexTempIDMap(outIndex uint32) models.TempOutputID {
 	w.RLock()
 	defer w.RUnlock()
 
-	return w.indexTempIDMap[outIndex]
+	return lo.Return1(w.indexTempIDMap.Get(outIndex))
 }
 
 // AddrIndexMap returns the index for the address specified.
 func (w *Wallet) AddrIndexMap(address string) uint32 {
 	w.RLock()
 	defer w.RUnlock()
 
-	return w.addrIndexMap[address]
+	return lo.Return1(w.addrIndexMap.Get(address))
 }
 
 // AddUnspentOutput adds an unspentOutput of a given wallet.
 func (w *Wallet) AddUnspentOutput(id models.TempOutputID, output *models.OutputData) {
 	w.Lock()
 	defer w.Unlock()
 
-	w.unspentOutputs[id] = output
-	w.indexTempIDMap[output.AddressIndex] = id
+	w.unspentOutputs.Set(id, output)
+	w.indexTempIDMap.Set(output.AddressIndex, id)
 
 	if w.walletType == Reuse {
-		w.reuseTempIDPool[id] = types.Void
+		w.reuseTempIDPool.Set(id, types.Void)
 	}
 }
 
@@ -150,7 +150,7 @@ func (w *Wallet) UnspentOutputBalance(id models.TempOutputID) iotago.BaseToken {
 	defer w.RUnlock()
 
 	total := iotago.BaseToken(0)
-	if out, ok := w.unspentOutputs[id]; ok {
+	if out, ok := w.unspentOutputs.Get(id); ok {
 		total += out.OutputStruct.BaseTokenAmount()
 	}
 
@@ -169,7 +169,7 @@ func (w *Wallet) UnspentOutputsLeft() (left int) {
 
 	switch w.walletType {
 	case Reuse:
-		left = len(w.reuseTempIDPool)
+		left = w.reuseTempIDPool.Size()
 	default:
 		left = int(w.nextAddrIdxToUse.Load() - w.nextAddrToSpent.Load())
 	}
@@ -183,7 +183,7 @@ func (w *Wallet) AddReuseAddress(id models.TempOutputID) {
 	defer w.Unlock()
 
 	if w.walletType == Reuse {
-		w.reuseTempIDPool[id] = types.Void
+		w.reuseTempIDPool.Set(id, types.Void)
 	}
 }
 
@@ -193,11 +193,8 @@ func (w *Wallet) GetReuseAddress() models.TempOutputID {
 	defer w.Unlock()
 
 	if w.walletType == Reuse {
-		if len(w.reuseTempIDPool) > 0 {
-			for id := range w.reuseTempIDPool {
-				delete(w.reuseTempIDPool, id)
-				return id
-			}
+		if id, _, exists := w.reuseTempIDPool.Pop(); exists {
+			return id
 		}
 	}
 
@@ -248,7 +245,7 @@ func (w *Wallet) KeyPair(index uint32) (ed25519.PrivateKey, ed25519.PublicKey) {
 
 // UnspentOutputsLength returns the number of unspent outputs on the wallet.
 func (w *Wallet) UnspentOutputsLength() int {
-	return len(w.unspentOutputs)
+	return w.unspentOutputs.Size()
 }
 
 func (w *Wallet) getAndIncNextAddrToSpent() uint32 {