From 11a2536058c64336808ce304f99fc57928b55dfa Mon Sep 17 00:00:00 2001 From: keithsue Date: Mon, 9 Sep 2024 17:13:59 +0800 Subject: [PATCH] optimize rune balances --- x/btcbridge/keeper/tss.go | 4 ++-- x/btcbridge/keeper/utxo.go | 12 ++++++++---- x/btcbridge/keeper/withdraw.go | 4 ++-- x/btcbridge/types/bitcoin.go | 16 +++++----------- x/btcbridge/types/runes.go | 29 ++++++++++++++++++++++++++--- 5 files changed, 43 insertions(+), 22 deletions(-) diff --git a/x/btcbridge/keeper/tss.go b/x/btcbridge/keeper/tss.go index 4a368d50..71a47808 100644 --- a/x/btcbridge/keeper/tss.go +++ b/x/btcbridge/keeper/tss.go @@ -290,7 +290,7 @@ func (k Keeper) handleTransferTx(ctx sdk.Context, p *psbt.Packet, sourceVault, d } } - runeBalances := make([]*types.RuneBalance, 0) + runeBalances := make(types.RuneBalances, 0) for i, ti := range p.UnsignedTx.TxIn { hash := ti.PreviousOutPoint.Hash.String() @@ -371,7 +371,7 @@ func (k Keeper) handleTransferTx(ctx sdk.Context, p *psbt.Packet, sourceVault, d Amount: uint64(out.Value), PubKeyScript: out.PkScript, IsLocked: false, - Runes: types.GetCompactRuneBalances(runeBalances), + Runes: runeBalances.Compact(), } k.saveUTXO(ctx, utxo) diff --git a/x/btcbridge/keeper/utxo.go b/x/btcbridge/keeper/utxo.go index 4fe48b23..8f864bb3 100644 --- a/x/btcbridge/keeper/utxo.go +++ b/x/btcbridge/keeper/utxo.go @@ -25,7 +25,7 @@ type UTXOViewKeeper interface { GetUnlockedUTXOsByAddr(ctx sdk.Context, addr string) []*types.UTXO GetOrderedUTXOsByAddr(ctx sdk.Context, addr string) []*types.UTXO - GetTargetRunesUTXOs(ctx sdk.Context, addr string, runeId string, targetAmount uint128.Uint128) ([]*types.UTXO, uint128.Uint128) + GetTargetRunesUTXOs(ctx sdk.Context, addr string, runeId string, targetAmount uint128.Uint128) ([]*types.UTXO, []*types.RuneBalance) IterateAllUTXOs(ctx sdk.Context, cb func(utxo *types.UTXO) (stop bool)) IterateUTXOsByAddr(ctx sdk.Context, addr string, cb func(addr string, utxo *types.UTXO) (stop bool)) @@ -150,10 +150,11 @@ func (bvk *BaseUTXOViewKeeper) GetOrderedUTXOsByAddr(ctx sdk.Context, addr strin } // GetTargetRunesUTXOs gets the unlocked runes utxos targeted by the given params -func (bvk *BaseUTXOViewKeeper) GetTargetRunesUTXOs(ctx sdk.Context, addr string, runeId string, targetAmount uint128.Uint128) ([]*types.UTXO, uint128.Uint128) { +func (bvk *BaseUTXOViewKeeper) GetTargetRunesUTXOs(ctx sdk.Context, addr string, runeId string, targetAmount uint128.Uint128) ([]*types.UTXO, []*types.RuneBalance) { utxos := make([]*types.UTXO, 0) totalAmount := uint128.Zero + totalRuneBalances := make(types.RuneBalances, 0) bvk.IterateRunesUTXOs(ctx, addr, runeId, func(addr string, id string, amount uint128.Uint128, utxo *types.UTXO) (stop bool) { if utxo.IsLocked { @@ -163,15 +164,18 @@ func (bvk *BaseUTXOViewKeeper) GetTargetRunesUTXOs(ctx sdk.Context, addr string, utxos = append(utxos, utxo) totalAmount = totalAmount.Add(amount) + totalRuneBalances = append(totalRuneBalances, utxo.Runes...) return totalAmount.Cmp(targetAmount) >= 0 }) if totalAmount.Cmp(targetAmount) < 0 { - return nil, uint128.Zero + return nil, nil } - return utxos, totalAmount.Sub(targetAmount) + runeBalancesDelta := totalRuneBalances.Compact().Update(runeId, totalAmount.Sub(targetAmount)) + + return utxos, runeBalancesDelta } func (bvk *BaseUTXOViewKeeper) IterateAllUTXOs(ctx sdk.Context, cb func(utxo *types.UTXO) (stop bool)) { diff --git a/x/btcbridge/keeper/withdraw.go b/x/btcbridge/keeper/withdraw.go index 2c7d342c..8dd723e4 100644 --- a/x/btcbridge/keeper/withdraw.go +++ b/x/btcbridge/keeper/withdraw.go @@ -93,14 +93,14 @@ func (k Keeper) NewRunesWithdrawRequest(ctx sdk.Context, sender string, amount s runeAmount := uint128.FromBig(amount.Amount.BigInt()) - runesUTXOs, amountDelta := k.GetTargetRunesUTXOs(ctx, vault, runeId.ToString(), runeAmount) + runesUTXOs, runeBalancesDelta := k.GetTargetRunesUTXOs(ctx, vault, runeId.ToString(), runeAmount) if len(runesUTXOs) == 0 { return nil, types.ErrInsufficientUTXOs } paymentUTXOIterator := k.GetUTXOIteratorByAddr(ctx, btcVault) - psbt, selectedUTXOs, changeUTXO, runesChangeUTXO, err := types.BuildRunesPsbt(runesUTXOs, paymentUTXOIterator, sender, runeId.ToString(), runeAmount, feeRate, amountDelta, vault, btcVault) + psbt, selectedUTXOs, changeUTXO, runesChangeUTXO, err := types.BuildRunesPsbt(runesUTXOs, paymentUTXOIterator, sender, runeId.ToString(), runeAmount, feeRate, runeBalancesDelta, vault, btcVault) if err != nil { return nil, err } diff --git a/x/btcbridge/types/bitcoin.go b/x/btcbridge/types/bitcoin.go index 86f7ac36..d394b6a5 100644 --- a/x/btcbridge/types/bitcoin.go +++ b/x/btcbridge/types/bitcoin.go @@ -71,7 +71,7 @@ func BuildPsbt(utxoIterator UTXOIterator, recipient string, amount int64, feeRat // BuildRunesPsbt builds a bitcoin psbt for runes edict from the given params. // Assume that the utxo script type is witness. -func BuildRunesPsbt(utxos []*UTXO, paymentUTXOIterator UTXOIterator, recipient string, runeId string, amount uint128.Uint128, feeRate int64, runesChangeAmount uint128.Uint128, runesChange string, change string) (*psbt.Packet, []*UTXO, *UTXO, *UTXO, error) { +func BuildRunesPsbt(utxos []*UTXO, paymentUTXOIterator UTXOIterator, recipient string, runeId string, amount uint128.Uint128, feeRate int64, runeBalancesDelta []*RuneBalance, runesChange string, change string) (*psbt.Packet, []*UTXO, *UTXO, *UTXO, error) { chaincfg := sdk.GetConfig().GetBtcChainCfg() recipientAddr, err := btcutil.DecodeAddress(recipient, chaincfg) @@ -107,9 +107,8 @@ func BuildRunesPsbt(utxos []*UTXO, paymentUTXOIterator UTXOIterator, recipient s var runesChangeUTXO *UTXO edictOutputIndex := uint32(1) - if runesChangeAmount.Cmp64(0) > 0 { - // we can guarantee that every runes UTXO only includes a single rune by the deposit policy - runesChangeUTXO = GetRunesChangeUTXO(runeId, runesChangeAmount, runesChange, runesChangePkScript, 1) + if len(runeBalancesDelta) > 0 { + runesChangeUTXO = GetRunesChangeUTXO(runeId, runeBalancesDelta, runesChange, runesChangePkScript, 1) // allocate the remaining runes to the first non-OP_RETURN output by default txOuts = append(txOuts, wire.NewTxOut(RunesOutValue, runesChangePkScript)) @@ -281,18 +280,13 @@ func GetChangeUTXO(tx *wire.MsgTx, change string) *UTXO { } // GetRunesChangeUTXO gets the runes change utxo. -func GetRunesChangeUTXO(runeId string, changeAmount uint128.Uint128, change string, changePkScript []byte, outIndex uint32) *UTXO { +func GetRunesChangeUTXO(runeId string, runeBalancesDelta []*RuneBalance, change string, changePkScript []byte, outIndex uint32) *UTXO { return &UTXO{ Vout: uint64(outIndex), Address: change, Amount: RunesOutValue, PubKeyScript: changePkScript, - Runes: []*RuneBalance{ - { - Id: runeId, - Amount: changeAmount.String(), - }, - }, + Runes: runeBalancesDelta, } } diff --git a/x/btcbridge/types/runes.go b/x/btcbridge/types/runes.go index 91060821..55a9356f 100644 --- a/x/btcbridge/types/runes.go +++ b/x/btcbridge/types/runes.go @@ -83,6 +83,7 @@ func ParseEdicts(tx *wire.MsgTx, payload []byte) ([]*Edict, error) { return nil, ErrInvalidRunes } + // actually we only support one edict for now, so delta is unnecessary edict := Edict{ Id: &RuneId{ Block: integers[i].Big().Uint64(), @@ -197,11 +198,14 @@ func (e *Edict) MustMarshalLEB128() []byte { return payload } -// GetCompactRuneBalances gets the compact rune balances -func GetCompactRuneBalances(runeBalances []*RuneBalance) []*RuneBalance { +// RuneBalances defines a set of rune balances +type RuneBalances []*RuneBalance + +// Compact gets the compact rune balances +func (rbs RuneBalances) Compact() RuneBalances { balanceMap := make(map[string]uint128.Uint128) - for _, balance := range runeBalances { + for _, balance := range rbs { balanceMap[balance.Id] = balanceMap[balance.Id].Add(RuneAmountFromString(balance.Amount)) } @@ -213,6 +217,25 @@ func GetCompactRuneBalances(runeBalances []*RuneBalance) []*RuneBalance { return compactBalances } +// Update updates the balance to the specified amount for the given rune id +// The rune balance will be removed if the given amount is zero +// Assume that the given RuneBalances is compact +func (rbs RuneBalances) Update(id string, amount uint128.Uint128) RuneBalances { + for i, balance := range rbs { + if balance.Id == id { + if !amount.IsZero() { + rbs[i].Amount = amount.String() + } else { + rbs = append(rbs[:i], rbs[i+1:]...) + } + + break + } + } + + return rbs +} + // RuneAmountFromString converts the given string to the rune amount // Panic if any error occurred func RuneAmountFromString(str string) uint128.Uint128 {