core: track NEP5 balances

pkg/core/blockchain.go

@@ -29,7 +29,7 @@ import (
 // Tuning parameters.
 const (
 	headerBatchCount = 2000
-	version          = "0.0.5"
+	version          = "0.0.6"
 
 	// 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
@@ -714,8 +714,7 @@ func (bc *Blockchain) storeBlock(block *block.Block) error {
 						}
 						amount = emit.BytesToInt(bs)
 					}
-					// TODO: #498
+					bc.processNEP5Transfer(cache, tx, block, note.ScriptHash, from, to, amount.Int64())
-					_, _, _, _ = op, from, to, amount
 				}
 			} else {
 				bc.log.Warn("contract invocation failed",
@@ -755,6 +754,50 @@ func (bc *Blockchain) storeBlock(block *block.Block) error {
 	return nil
 }
 
+func parseUint160(addr []byte) *util.Uint160 {
+	if u, err := util.Uint160DecodeBytesBE(addr); err == nil {
+		return &u
+	}
+	return nil
+}
+
+func (bc *Blockchain) processNEP5Transfer(cache *cachedDao, tx *transaction.Transaction, b *block.Block, sc util.Uint160, from, to []byte, amount int64) {
+	toAddr := parseUint160(to)
+	fromAddr := parseUint160(from)
+	if fromAddr != nil {
+		acc, err := cache.GetAccountStateOrNew(*fromAddr)
+		if err != nil {
+			return
+		}
+		bs := acc.NEP5Balances[sc]
+		if bs == nil {
+			bs = new(state.NEP5Tracker)
+			acc.NEP5Balances[sc] = bs
+		}
+		bs.Balance -= amount
+		bs.LastUpdatedBlock = b.Index
+		if err := cache.PutAccountState(acc); err != nil {
+			return
+		}
+	}
+	if toAddr != nil {
+		acc, err := cache.GetAccountStateOrNew(*toAddr)
+		if err != nil {
+			return
+		}
+		bs := acc.NEP5Balances[sc]
+		if bs == nil {
+			bs = new(state.NEP5Tracker)
+			acc.NEP5Balances[sc] = bs
+		}
+		bs.Balance += amount
+		bs.LastUpdatedBlock = b.Index
+		if err := cache.PutAccountState(acc); err != nil {
+			return
+		}
+	}
+}
+
 // LastBatch returns last persisted storage batch.
 func (bc *Blockchain) LastBatch() *storage.MemBatch {
 	return bc.lastBatch

pkg/core/state/account.go

@@ -35,6 +35,9 @@ type Account struct {
 	Votes      []*keys.PublicKey
 	Balances   map[util.Uint256][]UnspentBalance
 	Unclaimed  []UnclaimedBalance
+	// NEP5Balances is a map of the NEP5 contract hashes
+	// to the corresponding structures.
+	NEP5Balances map[util.Uint160]*NEP5Tracker
 }
 
 // NewAccount returns a new Account object.
@@ -46,6 +49,8 @@ func NewAccount(scriptHash util.Uint160) *Account {
 		Votes:      []*keys.PublicKey{},
 		Balances:   make(map[util.Uint256][]UnspentBalance),
 		Unclaimed:  []UnclaimedBalance{},
+
+		NEP5Balances: make(map[util.Uint160]*NEP5Tracker),
 	}
 }
 
@@ -70,6 +75,16 @@ func (s *Account) DecodeBinary(br *io.BinReader) {
 	}
 
 	br.ReadArray(&s.Unclaimed)
+
+	lenBalances = br.ReadVarUint()
+	s.NEP5Balances = make(map[util.Uint160]*NEP5Tracker, lenBalances)
+	for i := 0; i < int(lenBalances); i++ {
+		var key util.Uint160
+		var tr NEP5Tracker
+		br.ReadBytes(key[:])
+		tr.DecodeBinary(br)
+		s.NEP5Balances[key] = &tr
+	}
 }
 
 // EncodeBinary encodes Account to the given BinWriter.
@@ -89,6 +104,12 @@ func (s *Account) EncodeBinary(bw *io.BinWriter) {
 	}
 
 	bw.WriteArray(s.Unclaimed)
+
+	bw.WriteVarUint(uint64(len(s.NEP5Balances)))
+	for k, v := range s.NEP5Balances {
+		bw.WriteBytes(k[:])
+		v.EncodeBinary(bw)
+	}
 }
 
 // DecodeBinary implements io.Serializable interface.

pkg/core/state/nep5.go

@@ -0,0 +1,26 @@
+package state
+
+import (
+	"github.com/nspcc-dev/neo-go/pkg/io"
+)
+
+// NEP5Tracker contains info about a single account in a NEP5 contract.
+type NEP5Tracker struct {
+	// Balance is the current balance of the account.
+	Balance int64
+	// LastUpdatedBlock is a number of block when last `transfer` to or from the
+	// account occured.
+	LastUpdatedBlock uint32
+}
+
+// EncodeBinary implements io.Serializable interface.
+func (t *NEP5Tracker) EncodeBinary(w *io.BinWriter) {
+	w.WriteU64LE(uint64(t.Balance))
+	w.WriteU32LE(t.LastUpdatedBlock)
+}
+
+// DecodeBinary implements io.Serializable interface.
+func (t *NEP5Tracker) DecodeBinary(r *io.BinReader) {
+	t.Balance = int64(r.ReadU64LE())
+	t.LastUpdatedBlock = r.ReadU32LE()
+}

pkg/core/state/nep5_test.go

@@ -0,0 +1,30 @@
+package state
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/pkg/io"
+	"github.com/nspcc-dev/neo-go/pkg/util"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNEP5Tracker_EncodeBinary(t *testing.T) {
+	expected := &NEP5Tracker{
+		Balance:          int64(rand.Uint64()),
+		LastUpdatedBlock: rand.Uint32(),
+	}
+
+	testEncodeDecode(t, expected, new(NEP5Tracker))
+}
+
+func testEncodeDecode(t *testing.T, expected, actual io.Serializable) {
+	w := io.NewBufBinWriter()
+	expected.EncodeBinary(w.BinWriter)
+	require.NoError(t, w.Err)
+
+	r := io.NewBinReaderFromBuf(w.Bytes())
+	actual.DecodeBinary(r)
+	require.NoError(t, r.Err)
+	require.Equal(t, expected, actual)
+}