core: rework balance accounting

Store all unspents instead of balance values. Bump store version as this is an
incompatible change.

pkg/core/account_state.go

@@ -69,13 +69,21 @@ func (a Accounts) commit(store storage.Store) error {
 	return nil
 }
 
+// UnspentBalance contains input/output transactons that sum up into the
+// account balance for the given asset.
+type UnspentBalance struct {
+	Tx    util.Uint256
+	Index uint16
+	Value util.Fixed8
+}
+
 // AccountState represents the state of a NEO account.
 type AccountState struct {
 	Version    uint8
 	ScriptHash util.Uint160
 	IsFrozen   bool
 	Votes      []*keys.PublicKey
-	Balances   map[util.Uint256]util.Fixed8
+	Balances   map[util.Uint256][]UnspentBalance
 }
 
 // NewAccountState returns a new AccountState object.
@@ -85,7 +93,7 @@ func NewAccountState(scriptHash util.Uint160) *AccountState {
 		ScriptHash: scriptHash,
 		IsFrozen:   false,
 		Votes:      []*keys.PublicKey{},
-		Balances:   make(map[util.Uint256]util.Fixed8),
+		Balances:   make(map[util.Uint256][]UnspentBalance),
 	}
 }
 
@@ -96,14 +104,14 @@ func (s *AccountState) DecodeBinary(br *io.BinReader) {
 	br.ReadLE(&s.IsFrozen)
 	br.ReadArray(&s.Votes)
 
-	s.Balances = make(map[util.Uint256]util.Fixed8)
+	s.Balances = make(map[util.Uint256][]UnspentBalance)
 	lenBalances := br.ReadVarUint()
 	for i := 0; i < int(lenBalances); i++ {
 		key := util.Uint256{}
 		br.ReadLE(&key)
-		var val util.Fixed8
+		ubs := make([]UnspentBalance, 0)
-		br.ReadLE(&val)
+		br.ReadArray(&ubs)
-		s.Balances[key] = val
+		s.Balances[key] = ubs
 	}
 }
 
@@ -114,21 +122,37 @@ func (s *AccountState) EncodeBinary(bw *io.BinWriter) {
 	bw.WriteLE(s.IsFrozen)
 	bw.WriteArray(s.Votes)
 
-	balances := s.nonZeroBalances()
+	bw.WriteVarUint(uint64(len(s.Balances)))
-	bw.WriteVarUint(uint64(len(balances)))
+	for k, v := range s.Balances {
-	for k, v := range balances {
 		bw.WriteLE(k)
-		bw.WriteLE(v)
+		bw.WriteArray(v)
 	}
 }
 
-// nonZeroBalances returns only the non-zero balances for the account.
+// DecodeBinary implements io.Serializable interface.
-func (s *AccountState) nonZeroBalances() map[util.Uint256]util.Fixed8 {
+func (u *UnspentBalance) DecodeBinary(r *io.BinReader) {
-	b := make(map[util.Uint256]util.Fixed8)
+	u.Tx.DecodeBinary(r)
-	for k, v := range s.Balances {
+	r.ReadLE(&u.Index)
-		if v > 0 {
+	r.ReadLE(&u.Value)
-			b[k] = v
+}
-		}
+
-	}
+// EncodeBinary implements io.Serializable interface.
-	return b
+func (u UnspentBalance) EncodeBinary(w *io.BinWriter) {
+	u.Tx.EncodeBinary(w)
+	w.WriteLE(u.Index)
+	w.WriteLE(u.Value)
+}
+
+// GetBalanceValues sums all unspent outputs and returns a map of asset IDs to
+// overall balances.
+func (s *AccountState) GetBalanceValues() map[util.Uint256]util.Fixed8 {
+	res := make(map[util.Uint256]util.Fixed8)
+	for k, v := range s.Balances {
+		balance := util.Fixed8(0)
+		for _, b := range v {
+			balance += b.Value
+		}
+		res[k] = balance
+	}
+	return res
 }

pkg/core/account_state_test.go

@@ -12,11 +12,18 @@ import (
 func TestDecodeEncodeAccountState(t *testing.T) {
 	var (
 		n        = 10
-		balances = make(map[util.Uint256]util.Fixed8)
+		balances = make(map[util.Uint256][]UnspentBalance)
 		votes    = make([]*keys.PublicKey, n)
 	)
 	for i := 0; i < n; i++ {
-		balances[randomUint256()] = util.Fixed8(int64(randomInt(1, 10000)))
+		asset := randomUint256()
+		for j := 0; j < i+1; j++ {
+			balances[asset] = append(balances[asset], UnspentBalance{
+				Tx:    randomUint256(),
+				Index: uint16(randomInt(0, 65535)),
+				Value: util.Fixed8(int64(randomInt(1, 10000))),
+			})
+		}
 		k, err := keys.NewPrivateKey()
 		assert.Nil(t, err)
 		votes[i] = k.PublicKey()
@@ -48,3 +55,18 @@ func TestDecodeEncodeAccountState(t *testing.T) {
 	}
 	assert.Equal(t, a.Balances, aDecode.Balances)
 }
+
+func TestAccountStateBalanceValues(t *testing.T) {
+	asset1 := randomUint256()
+	asset2 := randomUint256()
+	as := AccountState{Balances: make(map[util.Uint256][]UnspentBalance)}
+	ref := 0
+	for i := 0; i < 10; i++ {
+		ref += i
+		as.Balances[asset1] = append(as.Balances[asset1], UnspentBalance{Value: util.Fixed8(i)})
+		as.Balances[asset2] = append(as.Balances[asset2], UnspentBalance{Value: util.Fixed8(i * 10)})
+	}
+	bVals := as.GetBalanceValues()
+	assert.Equal(t, util.Fixed8(ref), bVals[asset1])
+	assert.Equal(t, util.Fixed8(ref*10), bVals[asset2])
+}

pkg/core/blockchain.go

@@ -23,7 +23,7 @@ import (
 // Tuning parameters.
 const (
 	headerBatchCount = 2000
-	version          = "0.0.1"
+	version          = "0.0.2"
 
 	// This one comes from C# code and it's different from the constant used
 	// when creating an asset with Neo.Asset.Create interop call. It looks
@@ -361,16 +361,16 @@ func (bc *Blockchain) storeBlock(block *Block) error {
 		unspentCoins[tx.Hash()] = NewUnspentCoinState(len(tx.Outputs))
 
 		// Process TX outputs.
-		for _, output := range tx.Outputs {
+		for index, output := range tx.Outputs {
 			account, err := accounts.getAndUpdate(bc.store, output.ScriptHash)
 			if err != nil {
 				return err
 			}
-			if _, ok := account.Balances[output.AssetID]; ok {
+			account.Balances[output.AssetID] = append(account.Balances[output.AssetID], UnspentBalance{
-				account.Balances[output.AssetID] += output.Amount
+				Tx:    tx.Hash(),
-			} else {
+				Index: uint16(index),
-				account.Balances[output.AssetID] = output.Amount
+				Value: output.Amount,
-			}
+			})
 		}
 
 		// Process TX inputs that are grouped by previous hash.
@@ -398,7 +398,21 @@ func (bc *Blockchain) storeBlock(block *Block) error {
 					spentCoins[input.PrevHash] = spentCoin
 				}
 
-				account.Balances[prevTXOutput.AssetID] -= prevTXOutput.Amount
+				balancesLen := len(account.Balances[prevTXOutput.AssetID])
+				if balancesLen <= 1 {
+					delete(account.Balances, prevTXOutput.AssetID)
+				} else {
+					var gotTx bool
+					for index, balance := range account.Balances[prevTXOutput.AssetID] {
+						if !gotTx && balance.Tx.Equals(input.PrevHash) && balance.Index == input.PrevIndex {
+							gotTx = true
+						}
+						if gotTx && index+1 < balancesLen {
+							account.Balances[prevTXOutput.AssetID][index] = account.Balances[prevTXOutput.AssetID][index+1]
+						}
+					}
+					account.Balances[prevTXOutput.AssetID] = account.Balances[prevTXOutput.AssetID][:balancesLen-1]
+				}
 			}
 		}
 

pkg/core/interop_neo.go

@@ -338,7 +338,7 @@ func (ic *interopContext) accountGetBalance(v *vm.VM) error {
 	if err != nil {
 		return err
 	}
-	balance, ok := acc.Balances[ashash]
+	balance, ok := acc.GetBalanceValues()[ashash]
 	if !ok {
 		balance = util.Fixed8(0)
 	}

pkg/rpc/wrappers/account_state.go

@@ -35,7 +35,7 @@ type Balance struct {
 // NewAccountState creates a new AccountState wrapper.
 func NewAccountState(a *core.AccountState) AccountState {
 	balances := make(Balances, 0, len(a.Balances))
-	for k, v := range a.Balances {
+	for k, v := range a.GetBalanceValues() {
 		balances = append(balances, Balance{
 			Asset: k,
 			Value: v,