Merge pull request #756 from nspcc-dev/feature/unclaimed

core/state: do not unmarshal Unclaimed balances in account

pkg/core/blockchain.go

@@ -507,13 +507,16 @@ func (bc *Blockchain) storeBlock(block *block.Block) error {
 				}
 
 				if prevTXOutput.AssetID.Equals(GoverningTokenID()) {
-					account.Unclaimed = append(account.Unclaimed, state.UnclaimedBalance{
+					err = account.Unclaimed.Put(&state.UnclaimedBalance{
 						Tx:    input.PrevHash,
 						Index: input.PrevIndex,
 						Start: unspent.Height,
 						End:   block.Index,
 						Value: prevTXOutput.Amount,
 					})
+					if err != nil {
+						return err
+					}
 					if err = processTXWithValidatorsSubtract(prevTXOutput, account, cache); err != nil {
 						return err
 					}
@@ -614,19 +617,7 @@ func (bc *Blockchain) storeBlock(block *block.Block) error {
 					return err
 				}
 
-				var changed bool
+				changed := acc.Unclaimed.Remove(input.PrevHash, input.PrevIndex)
-				for i := range acc.Unclaimed {
-					if acc.Unclaimed[i].Tx == input.PrevHash && acc.Unclaimed[i].Index == input.PrevIndex {
-						last := len(acc.Unclaimed) - 1
-						if last > i {
-							acc.Unclaimed[i] = acc.Unclaimed[last]
-						}
-						acc.Unclaimed = acc.Unclaimed[:last]
-						changed = true
-						break
-					}
-				}
-
 				if !changed {
 					bc.log.Warn("no spent coin in the account",
 						zap.String("tx", tx.Hash().StringLE()),

pkg/core/state/account.go

@@ -34,7 +34,7 @@ type Account struct {
 	IsFrozen   bool
 	Votes      []*keys.PublicKey
 	Balances   map[util.Uint256][]UnspentBalance
-	Unclaimed  []UnclaimedBalance
+	Unclaimed  UnclaimedBalances
 }
 
 // NewAccount returns a new Account object.
@@ -45,7 +45,7 @@ func NewAccount(scriptHash util.Uint160) *Account {
 		IsFrozen:   false,
 		Votes:      []*keys.PublicKey{},
 		Balances:   make(map[util.Uint256][]UnspentBalance),
-		Unclaimed:  []UnclaimedBalance{},
+		Unclaimed:  UnclaimedBalances{Raw: []byte{}},
 	}
 }
 
@@ -69,7 +69,9 @@ func (s *Account) DecodeBinary(br *io.BinReader) {
 		s.Balances[key] = ubs
 	}
 
-	br.ReadArray(&s.Unclaimed)
+	lenBalances = br.ReadVarUint()
+	s.Unclaimed.Raw = make([]byte, lenBalances*UnclaimedBalanceSize)
+	br.ReadBytes(s.Unclaimed.Raw)
 }
 
 // EncodeBinary encodes Account to the given BinWriter.
@@ -88,7 +90,8 @@ func (s *Account) EncodeBinary(bw *io.BinWriter) {
 		}
 	}
 
-	bw.WriteArray(s.Unclaimed)
+	bw.WriteVarUint(uint64(s.Unclaimed.Size()))
+	bw.WriteBytes(s.Unclaimed.Raw)
 }
 
 // DecodeBinary implements io.Serializable interface.

pkg/core/state/unclaimed.go

@@ -0,0 +1,69 @@
+package state
+
+import (
+	"bytes"
+	"encoding/binary"
+
+	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+)
+
+// UnclaimedBalanceSize is a size of the UnclaimedBalance struct in bytes.
+const UnclaimedBalanceSize = util.Uint256Size + 2 + 4 + 4 + 8
+
+// UnclaimedBalances is a slice of UnclaimedBalance.
+type UnclaimedBalances struct {
+	Raw []byte
+}
+
+// Size returns an amount of store unclaimed balances.
+func (bs *UnclaimedBalances) Size() int {
+	return len(bs.Raw) / UnclaimedBalanceSize
+}
+
+// ForEach iterates over all unclaimed balances.
+func (bs *UnclaimedBalances) ForEach(f func(*UnclaimedBalance) error) error {
+	b := new(UnclaimedBalance)
+	for i := 0; i < len(bs.Raw); i += UnclaimedBalanceSize {
+		r := io.NewBinReaderFromBuf(bs.Raw[i : i+UnclaimedBalanceSize])
+		b.DecodeBinary(r)
+		if r.Err != nil {
+			return r.Err
+		} else if err := f(b); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Remove removes specified unclaim from the list and returns
+// false if it wasn't found.
+func (bs *UnclaimedBalances) Remove(tx util.Uint256, index uint16) bool {
+	const keySize = util.Uint256Size + 2
+	key := make([]byte, keySize)
+	copy(key, tx[:])
+	binary.LittleEndian.PutUint16(key[util.Uint256Size:], index)
+
+	for i := 0; i < len(bs.Raw); i += UnclaimedBalanceSize {
+		if bytes.Equal(bs.Raw[i:i+keySize], key) {
+			lastIndex := len(bs.Raw) - UnclaimedBalanceSize
+			if i != lastIndex {
+				copy(bs.Raw[i:i+UnclaimedBalanceSize], bs.Raw[lastIndex:])
+			}
+			bs.Raw = bs.Raw[:lastIndex]
+			return true
+		}
+	}
+	return false
+}
+
+// Put puts new unclaim in a list.
+func (bs *UnclaimedBalances) Put(b *UnclaimedBalance) error {
+	w := io.NewBufBinWriter()
+	b.EncodeBinary(w.BinWriter)
+	if w.Err != nil {
+		return w.Err
+	}
+	bs.Raw = append(bs.Raw, w.Bytes()...)
+	return nil
+}

pkg/core/state/unclaimed_test.go

@@ -0,0 +1,80 @@
+package state
+
+import (
+	"encoding/binary"
+	gio "io"
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUnclaimedBalance_Structure(t *testing.T) {
+	b := randomUnclaimed(t)
+	w := io.NewBufBinWriter()
+	b.EncodeBinary(w.BinWriter)
+	require.NoError(t, w.Err)
+
+	buf := w.Bytes()
+	require.Equal(t, UnclaimedBalanceSize, len(buf))
+	require.Equal(t, b.Tx.BytesBE(), buf[:util.Uint256Size])
+	require.Equal(t, b.Index, binary.LittleEndian.Uint16(buf[util.Uint256Size:]))
+}
+
+func TestUnclaimedBalances_Put(t *testing.T) {
+	bs := new(UnclaimedBalances)
+	b1 := randomUnclaimed(t)
+	b2 := randomUnclaimed(t)
+	b3 := randomUnclaimed(t)
+
+	require.NoError(t, bs.Put(b1))
+	require.Equal(t, 1, bs.Size())
+	require.NoError(t, bs.Put(b2))
+	require.Equal(t, 2, bs.Size())
+	require.NoError(t, bs.Put(b3))
+	require.Equal(t, 3, bs.Size())
+	require.True(t, bs.Remove(b2.Tx, b2.Index))
+	require.Equal(t, 2, bs.Size())
+	require.False(t, bs.Remove(b2.Tx, b2.Index))
+	require.Equal(t, 2, bs.Size())
+	require.True(t, bs.Remove(b1.Tx, b1.Index))
+	require.Equal(t, 1, bs.Size())
+	require.True(t, bs.Remove(b3.Tx, b3.Index))
+	require.Equal(t, 0, bs.Size())
+}
+
+func TestUnclaimedBalances_ForEach(t *testing.T) {
+	bs := new(UnclaimedBalances)
+	b1 := randomUnclaimed(t)
+	b2 := randomUnclaimed(t)
+	b3 := randomUnclaimed(t)
+
+	require.NoError(t, bs.Put(b1))
+	require.NoError(t, bs.Put(b2))
+	require.NoError(t, bs.Put(b3))
+
+	var indices []uint16
+	err := bs.ForEach(func(b *UnclaimedBalance) error {
+		indices = append(indices, b.Index)
+		return nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, []uint16{b1.Index, b2.Index, b3.Index}, indices)
+}
+
+func randomUnclaimed(t *testing.T) *UnclaimedBalance {
+	b := new(UnclaimedBalance)
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	_, err := gio.ReadFull(r, b.Tx[:])
+	require.NoError(t, err)
+
+	b.Index = uint16(rand.Uint32())
+	b.Start = rand.Uint32()
+	b.End = rand.Uint32()
+	b.Value = util.Fixed8(rand.Int63())
+
+	return b
+}

pkg/rpc/response/result/unclaimed.go

@@ -20,12 +20,16 @@ func NewUnclaimed(a *state.Account, chain core.Blockchainer) (*Unclaimed, error)
 		unavailable util.Fixed8
 	)
 
-	for _, ucb := range a.Unclaimed {
+	err := a.Unclaimed.ForEach(func(ucb *state.UnclaimedBalance) error {
 		gen, sys, err := chain.CalculateClaimable(ucb.Value, ucb.Start, ucb.End)
 		if err != nil {
-			return nil, err
+			return err
 		}
 		available += gen + sys
+		return nil
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	blockHeight := chain.BlockHeight()

pkg/rpc/server/server.go

@@ -379,7 +379,13 @@ func (s *Server) getClaimable(ps request.Params) (interface{}, error) {
 
 	var unclaimed []state.UnclaimedBalance
 	if acc := s.chain.GetAccountState(u); acc != nil {
-		unclaimed = acc.Unclaimed
+		err := acc.Unclaimed.ForEach(func(b *state.UnclaimedBalance) error {
+			unclaimed = append(unclaimed, *b)
+			return nil
+		})
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	var sum util.Fixed8