Add RedistributeV2

wallet/wallet.go

@@ -555,14 +555,11 @@ func (sw *SingleAddressWallet) UnconfirmedTransactions() (annotated []Event, err
 	return annotated, nil
 }
 
-// Redistribute returns a transaction that redistributes money in the wallet by
-// selecting a minimal set of inputs to cover the creation of the requested
-// outputs. It also returns a list of output IDs that need to be signed.
-func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte types.Currency) (txns []types.Transaction, toSign []types.Hash256, err error) {
+func (sw *SingleAddressWallet) selectRedistributeUTXOs(bh uint64, outputs int, amount types.Currency) ([]types.SiacoinElement, int, error) {
 	// fetch outputs from the store
 	elements, err := sw.store.UnspentSiacoinElements()
 	if err != nil {
-		return nil, nil, err
+		return nil, 0, err
 	}
 
 	// fetch outputs currently in the pool
@@ -573,13 +570,6 @@ func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte type
 		}
 	}
 
-	// grab current height
-	state := sw.cm.TipState()
-	bh := state.Index.Height
-
-	sw.mu.Lock()
-	defer sw.mu.Unlock()
-
 	// adjust the number of desired outputs for any output we encounter that is
 	// unused, matured and has the same value
 	utxos := make([]types.SiacoinElement, 0, len(elements))
@@ -598,6 +588,25 @@ func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte type
 			utxos = append(utxos, sce)
 		}
 	}
+	// desc sort
+	sort.Slice(utxos, func(i, j int) bool {
+		return utxos[i].SiacoinOutput.Value.Cmp(utxos[j].SiacoinOutput.Value) > 0
+	})
+	return utxos, outputs, nil
+}
+
+// Redistribute returns a transaction that redistributes money in the wallet by
+// selecting a minimal set of inputs to cover the creation of the requested
+// outputs. It also returns a list of output IDs that need to be signed.
+func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte types.Currency) (txns []types.Transaction, toSign []types.Hash256, err error) {
+	state := sw.cm.TipState()
+	sw.mu.Lock()
+	defer sw.mu.Unlock()
+
+	utxos, outputs, err := sw.selectRedistributeUTXOs(state.Index.Height, outputs, amount)
+	if err != nil {
+		return nil, nil, err
+	}
 
 	// return early if we don't have to defrag at all
 	if outputs <= 0 {
@@ -679,6 +688,96 @@ func (sw *SingleAddressWallet) Redistribute(outputs int, amount, feePerByte type
 	return
 }
 
+// RedistributeV2 returns a transaction that redistributes money in the wallet
+// by selecting a minimal set of inputs to cover the creation of the requested
+// outputs. It also returns a list of output IDs that need to be signed.
+func (sw *SingleAddressWallet) RedistributeV2(outputs int, amount, feePerByte types.Currency) (txns []types.V2Transaction, toSign [][]int, err error) {
+	state := sw.cm.TipState()
+	sw.mu.Lock()
+	defer sw.mu.Unlock()
+
+	utxos, outputs, err := sw.selectRedistributeUTXOs(state.Index.Height, outputs, amount)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// return early if we don't have to defrag at all
+	if outputs <= 0 {
+		return nil, nil, nil
+	}
+
+	// in case of an error we need to free all inputs
+	defer func() {
+		if err != nil {
+			for txnIdx, toSignTxn := range toSign {
+				for i := range toSignTxn {
+					delete(sw.locked, txns[txnIdx].SiacoinInputs[i].Parent.ID)
+				}
+			}
+		}
+	}()
+
+	// prepare defrag transactions
+	for outputs > 0 {
+		var txn types.V2Transaction
+		for i := 0; i < outputs && i < redistributeBatchSize; i++ {
+			txn.SiacoinOutputs = append(txn.SiacoinOutputs, types.SiacoinOutput{
+				Value:   amount,
+				Address: sw.addr,
+			})
+		}
+		outputs -= len(txn.SiacoinOutputs)
+
+		// estimate the fees
+		outputFees := feePerByte.Mul64(state.V2TransactionWeight(txn))
+		feePerInput := feePerByte.Mul64(bytesPerInput)
+
+		// collect outputs that cover the total amount
+		var inputs []types.SiacoinElement
+		want := amount.Mul64(uint64(len(txn.SiacoinOutputs)))
+		for _, sce := range utxos {
+			inputs = append(inputs, sce)
+			fee := feePerInput.Mul64(uint64(len(inputs))).Add(outputFees)
+			if SumOutputs(inputs).Cmp(want.Add(fee)) > 0 {
+				break
+			}
+		}
+
+		// not enough outputs found
+		fee := feePerInput.Mul64(uint64(len(inputs))).Add(outputFees)
+		if sumOut := SumOutputs(inputs); sumOut.Cmp(want.Add(fee)) < 0 {
+			return nil, nil, fmt.Errorf("%w: inputs %v < needed %v + txnFee %v", ErrNotEnoughFunds, sumOut.String(), want.String(), fee.String())
+		}
+
+		// set the miner fee
+		if !fee.IsZero() {
+			txn.MinerFee = fee
+		}
+
+		// add the change output
+		change := SumOutputs(inputs).Sub(want.Add(fee))
+		if !change.IsZero() {
+			txn.SiacoinOutputs = append(txn.SiacoinOutputs, types.SiacoinOutput{
+				Value:   change,
+				Address: sw.addr,
+			})
+		}
+
+		// add the inputs
+		toSignTxn := make([]int, 0, len(inputs))
+		for _, sce := range inputs {
+			toSignTxn = append(toSignTxn, len(txn.SiacoinInputs))
+			txn.SiacoinInputs = append(txn.SiacoinInputs, types.V2SiacoinInput{
+				Parent: sce,
+			})
+			sw.locked[sce.ID] = time.Now().Add(sw.cfg.ReservationDuration)
+		}
+		txns = append(txns, txn)
+		toSign = append(toSign, toSignTxn)
+	}
+	return
+}
+
 // ReleaseInputs is a helper function that releases the inputs of txn for use in
 // other transactions. It should only be called on transactions that are invalid
 // or will never be broadcast.

wallet/wallet_test.go

@@ -456,6 +456,101 @@ func TestWalletRedistribute(t *testing.T) {
 	}
 }
 
+func TestWalletRedistributeV2(t *testing.T) {
+	// create wallet store
+	pk := types.GeneratePrivateKey()
+	ws := testutil.NewEphemeralWalletStore(pk)
+
+	// create chain store
+	network, genesis := testutil.Network()
+	network.HardforkV2.AllowHeight = 1 // allow V2 transactions from the start
+	cs, tipState, err := chain.NewDBStore(chain.NewMemDB(), network, genesis)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// create chain manager and subscribe the wallet
+	cm := chain.NewManager(cs, tipState)
+	// create wallet
+	l := zaptest.NewLogger(t)
+	w, err := wallet.NewSingleAddressWallet(pk, cm, ws, wallet.WithLogger(l.Named("wallet")))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer w.Close()
+
+	// fund the wallet
+	mineAndSync(t, cm, ws, w, w.Address(), 1)
+	mineAndSync(t, cm, ws, w, types.VoidAddress, cm.TipState().MaturityHeight()-1)
+
+	redistribute := func(amount types.Currency, n int) error {
+		txns, toSign, err := w.RedistributeV2(n, amount, types.ZeroCurrency)
+		if err != nil {
+			return fmt.Errorf("redistribute failed: %w", err)
+		} else if len(txns) == 0 {
+			return nil
+		}
+
+		for i := 0; i < len(txns); i++ {
+			w.SignV2Inputs(cm.TipState(), &txns[i], toSign[i])
+		}
+		if _, err := cm.AddV2PoolTransactions(cm.Tip(), txns); err != nil {
+			return fmt.Errorf("failed to add transactions to pool: %w", err)
+		}
+		mineAndSync(t, cm, ws, w, types.VoidAddress, 1)
+		return nil
+	}
+
+	assertOutputs := func(amount types.Currency, n int) error {
+		utxos, err := w.SpendableOutputs()
+		if err != nil {
+			return fmt.Errorf("failed to get unspent outputs: %w", err)
+		}
+		var count int
+		for _, utxo := range utxos {
+			if utxo.SiacoinOutput.Value.Equals(amount) {
+				count++
+			}
+		}
+		if count != n {
+			return fmt.Errorf("expected %v outputs of %v, got %v", n, amount, count)
+		}
+		return nil
+	}
+
+	// assert we have one output
+	assertOutputs(tipState.BlockReward(), 1)
+
+	// redistribute the wallet into 4 outputs of 75KS
+	amount := types.Siacoins(75e3)
+	if err := redistribute(amount, 4); err != nil {
+		t.Fatal(err)
+	}
+	assertOutputs(amount, 4)
+
+	// redistribute the wallet into 4 outputs of 50KS
+	amount = types.Siacoins(50e3)
+	if err := redistribute(amount, 4); err != nil {
+		t.Fatal(err)
+	}
+	assertOutputs(amount, 4)
+
+	// redistribute the wallet into 3 outputs of 101KS - expect ErrNotEnoughFunds
+	if err := redistribute(types.Siacoins(101e3), 3); !errors.Is(err, wallet.ErrNotEnoughFunds) {
+		t.Fatalf("expected ErrNotEnoughFunds, got %v", err)
+	}
+
+	// redistribute the wallet into 3 outputs of 50KS - assert this is a no-op
+	txns, toSign, err := w.RedistributeV2(3, amount, types.ZeroCurrency)
+	if err != nil {
+		t.Fatal(err)
+	} else if len(txns) != 0 {
+		t.Fatalf("expected no transactions, got %v", len(txns))
+	} else if len(toSign) != 0 {
+		t.Fatalf("expected no ids, got %v", len(toSign))
+	}
+}
+
 func TestReorg(t *testing.T) {
 	// create wallet store
 	pk := types.GeneratePrivateKey()