Add wallet apply and revert logic to coreutils

.github/workflows/test.yml

@@ -11,7 +11,7 @@ jobs:
     strategy:
       matrix:
         os: [ ubuntu-latest, macos-latest, windows-latest ]
-        go-version: [ '1.20', '1.21' ]
+        go-version: [ '1.21', '1.22' ]
     runs-on: ${{ matrix.os }}
     timeout-minutes: 10
     steps:

go.mod

@@ -1,6 +1,6 @@
 module go.sia.tech/coreutils
 
-go 1.21
+go 1.21.6
 
 require (
 	go.etcd.io/bbolt v1.3.8

testutil/wallet.go

@@ -2,6 +2,7 @@ package testutil
 
 import (
 	"fmt"
+	"slices"
 	"sync"
 
 	"go.sia.tech/core/types"
@@ -11,192 +12,83 @@ import (
 
 // An EphemeralWalletStore is a Store that does not persist its state to disk. It is
 // primarily useful for testing or as a reference implementation.
-type EphemeralWalletStore struct {
-	privateKey types.PrivateKey
+type (
+	EphemeralWalletStore struct {
+		privateKey types.PrivateKey
 
-	mu          sync.Mutex
-	uncommitted []*chain.ApplyUpdate
-	tip         types.ChainIndex
-	utxos       map[types.Hash256]types.SiacoinElement
+		mu          sync.Mutex
+		uncommitted []*chain.ApplyUpdate
 
-	immaturePayoutTransactions map[uint64][]wallet.Transaction // transactions are never removed to simplify reorg handling
-	transactions               []wallet.Transaction
-}
-
-// ProcessChainRevertUpdate implements chain.Subscriber
-func (es *EphemeralWalletStore) ProcessChainRevertUpdate(cru *chain.RevertUpdate) error {
-	es.mu.Lock()
-	defer es.mu.Unlock()
-
-	// if the update is uncommitted, remove it from the uncommitted list
-	if len(es.uncommitted) > 0 && es.uncommitted[len(es.uncommitted)-1].State.Index == cru.State.Index {
-		es.uncommitted = es.uncommitted[:len(es.uncommitted)-1]
-		return nil
+		tip          types.ChainIndex
+		utxos        map[types.SiacoinOutputID]wallet.SiacoinElement
+		transactions []wallet.Transaction
 	}
 
-	walletAddress := types.StandardUnlockHash(es.privateKey.PublicKey())
-
-	// revert any siacoin element changes
-	cru.ForEachSiacoinElement(func(se types.SiacoinElement, spent bool) {
-		if se.SiacoinOutput.Address != walletAddress {
-			return
-		}
-
-		if spent {
-			es.utxos[se.ID] = se
-		} else {
-			delete(es.utxos, se.ID)
-		}
-	})
-
-	// remove any transactions that were added in the reverted block
-	filtered := es.transactions[:0]
-	for _, txn := range es.transactions {
-		if txn.Index == cru.State.Index {
-			continue
-		}
-		filtered = append(filtered, txn)
+	ephemeralWalletUpdateTxn struct {
+		store *EphemeralWalletStore
 	}
+)
 
-	// update element proofs
-	for id, se := range es.utxos {
-		cru.UpdateElementProof(&se.StateElement)
-		es.utxos[id] = se
+func (et *ephemeralWalletUpdateTxn) WalletStateElements() (elements []types.StateElement, _ error) {
+	for _, se := range et.store.utxos {
+		elements = append(elements, se.StateElement)
 	}
-	return nil
+	return
 }
 
-// ProcessChainApplyUpdate implements chain.Subscriber
-func (es *EphemeralWalletStore) ProcessChainApplyUpdate(cau *chain.ApplyUpdate, commit bool) error {
-	es.mu.Lock()
-	defer es.mu.Unlock()
-
-	es.uncommitted = append(es.uncommitted, cau)
-	if !commit {
-		return nil
+func (et *ephemeralWalletUpdateTxn) UpdateStateElements(elements []types.StateElement) error {
+	for _, se := range elements {
+		utxo := et.store.utxos[types.SiacoinOutputID(se.ID)]
+		utxo.StateElement = se
+		et.store.utxos[types.SiacoinOutputID(se.ID)] = utxo
 	}
+	return nil
+}
 
-	walletAddress := types.StandardUnlockHash(es.privateKey.PublicKey())
+func (et *ephemeralWalletUpdateTxn) AddTransactions(transactions []wallet.Transaction) error {
+	// transactions are added in reverse order to make the most recent transaction the first
+	slices.Reverse(transactions)
+	et.store.transactions = append(transactions, et.store.transactions...)
+	return nil
+}
 
-	for _, update := range es.uncommitted {
-		// cache the source of new immature outputs to show payout transactions
-		siacoinOutputSources := map[types.SiacoinOutputID]wallet.TransactionSource{
-			update.Block.ID().FoundationOutputID(): wallet.TxnSourceFoundationPayout,
-		}
-		// add the miner payouts
-		for i := range update.Block.MinerPayouts {
-			siacoinOutputSources[update.Block.ID().MinerOutputID(i)] = wallet.TxnSourceMinerPayout
+func (et *ephemeralWalletUpdateTxn) AddSiacoinElements(elements []wallet.SiacoinElement) error {
+	for _, se := range elements {
+		if _, ok := et.store.utxos[types.SiacoinOutputID(se.ID)]; ok {
+			return fmt.Errorf("siacoin element %q already exists", se.ID)
 		}
+		et.store.utxos[types.SiacoinOutputID(se.ID)] = se
+	}
+	return nil
+}
 
-		// add the file contract outputs
-		update.ForEachFileContractElement(func(fce types.FileContractElement, rev *types.FileContractElement, resolved bool, valid bool) {
-			if !resolved {
-				return
-			}
-
-			if valid {
-				for i := range fce.FileContract.ValidProofOutputs {
-					siacoinOutputSources[types.FileContractID(fce.ID).ValidOutputID(i)] = wallet.TxnSourceContract
-				}
-			} else {
-				for i := range fce.FileContract.MissedProofOutputs {
-					siacoinOutputSources[types.FileContractID(fce.ID).MissedOutputID(i)] = wallet.TxnSourceContract
-				}
-			}
-		})
-
-		// add matured transactions first since
-		for _, txn := range es.immaturePayoutTransactions[update.State.Index.Height] {
-			txn.Index = update.State.Index
-			txn.Timestamp = update.Block.Timestamp
-			// prepend the transaction to the wallet
-			es.transactions = append([]wallet.Transaction{txn}, es.transactions...)
+func (et *ephemeralWalletUpdateTxn) RemoveSiacoinElements(ids []types.SiacoinOutputID) error {
+	for _, id := range ids {
+		if _, ok := et.store.utxos[id]; !ok {
+			return fmt.Errorf("siacoin element %q does not exist", id)
 		}
+		delete(et.store.utxos, id)
+	}
+	return nil
+}
 
-		// add the block transactions
-		for _, txn := range update.Block.Transactions {
-			if !wallet.IsRelevantTransaction(txn, walletAddress) {
-				continue
-			}
-
-			wt := wallet.Transaction{
-				ID:          txn.ID(),
-				Index:       update.State.Index,
-				Transaction: txn,
-				Timestamp:   update.Block.Timestamp,
-				Source:      wallet.TxnSourceTransaction,
-			}
-
-			for _, sci := range txn.SiacoinInputs {
-				if sci.UnlockConditions.UnlockHash() != walletAddress {
-					continue
-				}
-
-				sce, ok := es.utxos[types.Hash256(sci.ParentID)]
-				if !ok {
-					return fmt.Errorf("missing relevant siacoin element %q", sci.ParentID)
-				}
-				wt.Outflow = wt.Outflow.Add(sce.SiacoinOutput.Value)
-			}
-
-			for _, sco := range txn.SiacoinOutputs {
-				if sco.Address != walletAddress {
-					continue
-				}
-
-				wt.Inflow = wt.Inflow.Add(sco.Value)
-			}
-
-			// prepend the transaction to the wallet
-			es.transactions = append([]wallet.Transaction{wt}, es.transactions...)
-
-			// add the siafund claim output IDs
-			for i := range txn.SiafundInputs {
-				siacoinOutputSources[txn.SiafundInputs[i].ParentID.ClaimOutputID()] = wallet.TxnSourceSiafundClaim
-			}
+func (et *ephemeralWalletUpdateTxn) RevertIndex(index types.ChainIndex) error {
+	// remove any transactions that were added in the reverted block
+	filtered := et.store.transactions[:0]
+	for i := range et.store.transactions {
+		if et.store.transactions[i].Index == index {
+			continue
 		}
+		filtered = append(filtered, et.store.transactions[i])
+	}
+	et.store.transactions = filtered
 
-		// update the utxo set
-		update.ForEachSiacoinElement(func(se types.SiacoinElement, spent bool) {
-			if se.SiacoinOutput.Address != walletAddress {
-				return
-			}
-
-			// update the utxo set
-			if spent {
-				delete(es.utxos, se.ID)
-			} else {
-				es.utxos[se.ID] = se
-			}
-
-			// create an immature payout transaction for any immature siacoin outputs
-			if se.MaturityHeight == 0 || spent {
-				return
-			}
-
-			source, ok := siacoinOutputSources[types.SiacoinOutputID(se.ID)]
-			if !ok {
-				panic("missing siacoin source")
-			}
-
-			es.immaturePayoutTransactions[se.MaturityHeight] = append(es.immaturePayoutTransactions[se.MaturityHeight], wallet.Transaction{
-				ID:          types.TransactionID(se.ID),
-				Transaction: types.Transaction{SiacoinOutputs: []types.SiacoinOutput{se.SiacoinOutput}},
-				Inflow:      se.SiacoinOutput.Value,
-				Source:      source,
-				// Index and Timestamp will be filled in later
-			})
-		})
-
-		// update the element proofs
-		for id, se := range es.utxos {
-			cau.UpdateElementProof(&se.StateElement)
-			es.utxos[id] = se
+	// remove any siacoin elements that were added in the reverted block
+	for id, se := range et.store.utxos {
+		if se.Index == index {
+			delete(et.store.utxos, id)
 		}
 	}
-
-	es.uncommitted = es.uncommitted[:0]
-	es.tip = cau.State.Index
 	return nil
 }
 
@@ -224,7 +116,7 @@ func (es *EphemeralWalletStore) TransactionCount() (uint64, error) {
 }
 
 // UnspentSiacoinElements returns the wallet's unspent siacoin outputs.
-func (es *EphemeralWalletStore) UnspentSiacoinElements() (utxos []types.SiacoinElement, _ error) {
+func (es *EphemeralWalletStore) UnspentSiacoinElements() (utxos []wallet.SiacoinElement, _ error) {
 	es.mu.Lock()
 	defer es.mu.Unlock()
 
@@ -241,12 +133,47 @@ func (es *EphemeralWalletStore) Tip() (types.ChainIndex, error) {
 	return es.tip, nil
 }
 
+// ProcessChainApplyUpdate implements chain.Subscriber.
+func (es *EphemeralWalletStore) ProcessChainApplyUpdate(cau *chain.ApplyUpdate, mayCommit bool) error {
+	es.mu.Lock()
+	defer es.mu.Unlock()
+
+	es.uncommitted = append(es.uncommitted, cau)
+	if !mayCommit {
+		return nil
+	}
+
+	address := types.StandardUnlockHash(es.privateKey.PublicKey())
+	ephemeralWalletUpdateTxn := &ephemeralWalletUpdateTxn{store: es}
+
+	if err := wallet.ApplyChainUpdates(ephemeralWalletUpdateTxn, address, es.uncommitted); err != nil {
+		return err
+	}
+	es.tip = cau.State.Index
+	es.uncommitted = nil
+	return nil
+}
+
+// ProcessChainRevertUpdate implements chain.Subscriber.
+func (es *EphemeralWalletStore) ProcessChainRevertUpdate(cru *chain.RevertUpdate) error {
+	es.mu.Lock()
+	defer es.mu.Unlock()
+
+	if len(es.uncommitted) > 0 && es.uncommitted[len(es.uncommitted)-1].State.Index == cru.State.Index {
+		es.uncommitted = es.uncommitted[:len(es.uncommitted)-1]
+		return nil
+	}
+
+	address := types.StandardUnlockHash(es.privateKey.PublicKey())
+	return wallet.RevertChainUpdate(&ephemeralWalletUpdateTxn{store: es}, address, cru)
+}
+
 // NewEphemeralWalletStore returns a new EphemeralWalletStore.
 func NewEphemeralWalletStore(pk types.PrivateKey) *EphemeralWalletStore {
 	return &EphemeralWalletStore{
 		privateKey: pk,
 
-		utxos:                      make(map[types.Hash256]types.SiacoinElement),
-		immaturePayoutTransactions: make(map[uint64][]wallet.Transaction),
+		utxos:        make(map[types.SiacoinOutputID]wallet.SiacoinElement),
+		transactions: make([]wallet.Transaction, 0),
 	}
 }