core: store nep5balances using contract id

Closes #1194

pkg/core/blockchain.go

@@ -715,15 +715,26 @@ func (bc *Blockchain) processNEP5Transfer(cache *dao.Cached, h util.Uint256, b *
 		Timestamp: b.Timestamp,
 		Tx:        h,
 	}
+	var id int32
+	nativeContract := bc.contracts.ByHash(sc)
+	if nativeContract != nil {
+		id = nativeContract.Metadata().ContractID
+	} else {
+		assetContract := bc.GetContractState(sc)
+		if assetContract == nil {
+			return
+		}
+		id = assetContract.ID
+	}
 	if !fromAddr.Equals(util.Uint160{}) {
 		balances, err := cache.GetNEP5Balances(fromAddr)
 		if err != nil {
 			return
 		}
-		bs := balances.Trackers[sc]
+		bs := balances.Trackers[id]
 		bs.Balance = *new(big.Int).Sub(&bs.Balance, amount)
 		bs.LastUpdatedBlock = b.Index
-		balances.Trackers[sc] = bs
+		balances.Trackers[id] = bs
 		transfer.Amount = *new(big.Int).Sub(&transfer.Amount, amount)
 		isBig, err := cache.AppendNEP5Transfer(fromAddr, balances.NextTransferBatch, transfer)
 		if err != nil {
@@ -741,10 +752,10 @@ func (bc *Blockchain) processNEP5Transfer(cache *dao.Cached, h util.Uint256, b *
 		if err != nil {
 			return
 		}
-		bs := balances.Trackers[sc]
+		bs := balances.Trackers[id]
 		bs.Balance = *new(big.Int).Add(&bs.Balance, amount)
 		bs.LastUpdatedBlock = b.Index
-		balances.Trackers[sc] = bs
+		balances.Trackers[id] = bs
 
 		transfer.Amount = *amount
 		isBig, err := cache.AppendNEP5Transfer(toAddr, balances.NextTransferBatch, transfer)
@@ -792,7 +803,7 @@ func (bc *Blockchain) GetUtilityTokenBalance(acc util.Uint160) *big.Int {
 	if err != nil {
 		return big.NewInt(0)
 	}
-	balance := bs.Trackers[bc.contracts.GAS.Hash].Balance
+	balance := bs.Trackers[bc.contracts.GAS.ContractID].Balance
 	return &balance
 }
 
@@ -803,7 +814,7 @@ func (bc *Blockchain) GetGoverningTokenBalance(acc util.Uint160) (*big.Int, uint
 	if err != nil {
 		return big.NewInt(0), 0
 	}
-	neo := bs.Trackers[bc.contracts.NEO.Hash]
+	neo := bs.Trackers[bc.contracts.NEO.ContractID]
 	return &neo.Balance, neo.LastUpdatedBlock
 }
 

pkg/core/dao/cacheddao.go

@@ -1,11 +1,9 @@
 package dao
 
 import (
-	"bytes"
 	"errors"
 
 	"github.com/nspcc-dev/neo-go/pkg/core/state"
-	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 )
@@ -125,65 +123,6 @@ func (cd *Cached) AppendNEP5Transfer(acc util.Uint160, index uint32, tr *state.N
 	return lg.Size() >= nep5TransferBatchSize, cd.PutNEP5TransferLog(acc, index, lg)
 }
 
-// MigrateNEP5Balances migrates NEP5 balances from old contract to the new one.
-func (cd *Cached) MigrateNEP5Balances(from, to util.Uint160) error {
-	var (
-		simpleDAO *Simple
-		cachedDAO = cd
-		ok        bool
-		w         = io.NewBufBinWriter()
-	)
-	for simpleDAO == nil {
-		simpleDAO, ok = cachedDAO.DAO.(*Simple)
-		if !ok {
-			cachedDAO, ok = cachedDAO.DAO.(*Cached)
-			if !ok {
-				panic("uknown DAO")
-			}
-		}
-	}
-	for acc, bs := range cd.balances {
-		err := simpleDAO.putNEP5Balances(acc, bs, w)
-		if err != nil {
-			return err
-		}
-		w.Reset()
-	}
-	cd.dropNEP5Cache = true
-	var store = simpleDAO.Store
-	// Create another layer of cache because we can't change original storage
-	// while seeking.
-	var upStore = storage.NewMemCachedStore(store)
-	store.Seek([]byte{byte(storage.STNEP5Balances)}, func(k, v []byte) {
-		if !bytes.Contains(v, from[:]) {
-			return
-		}
-		bs := state.NewNEP5Balances()
-		reader := io.NewBinReaderFromBuf(v)
-		bs.DecodeBinary(reader)
-		if reader.Err != nil {
-			panic("bad nep5 balances")
-		}
-		tr, ok := bs.Trackers[from]
-		if !ok {
-			return
-		}
-		delete(bs.Trackers, from)
-		bs.Trackers[to] = tr
-		w.Reset()
-		bs.EncodeBinary(w.BinWriter)
-		if w.Err != nil {
-			panic("error on nep5 balance encoding")
-		}
-		err := upStore.Put(k, w.Bytes())
-		if err != nil {
-			panic("can't put value in the DB")
-		}
-	})
-	_, err := upStore.Persist()
-	return err
-}
-
 // Persist flushes all the changes made into the (supposedly) persistent
 // underlying store.
 func (cd *Cached) Persist() (int, error) {

pkg/core/interop_neo.go

@@ -154,7 +154,6 @@ func contractUpdate(ic *interop.Context, v *vm.VM) error {
 		if err := ic.DAO.DeleteContractState(oldHash); err != nil {
 			return fmt.Errorf("failed to update script: %v", err)
 		}
-		ic.DAO.MigrateNEP5Balances(oldHash, newHash)
 	}
 	// if manifest was provided, update the old contract manifest and check associated
 	// storage items if needed

pkg/core/native/native_neo.go

@@ -210,7 +210,7 @@ func (n *NEO) unclaimedGas(ic *interop.Context, args []stackitem.Item) stackitem
 	if err != nil {
 		panic(err)
 	}
-	tr := bs.Trackers[n.Hash]
+	tr := bs.Trackers[n.ContractID]
 
 	gen := ic.Chain.CalculateClaimable(&tr.Balance, tr.LastUpdatedBlock, end)
 	return stackitem.NewBigInteger(gen)

pkg/core/native/native_nep5.go

@@ -210,7 +210,7 @@ func (c *nep5TokenNative) balanceOf(ic *interop.Context, args []stackitem.Item)
 	if err != nil {
 		panic(err)
 	}
-	balance := bs.Trackers[c.Hash].Balance
+	balance := bs.Trackers[c.ContractID].Balance
 	return stackitem.NewBigInteger(&balance)
 }
 

pkg/core/state/nep5.go

@@ -43,10 +43,10 @@ type NEP5Transfer struct {
 	Tx util.Uint256
 }
 
-// NEP5Balances is a map of the NEP5 contract hashes
+// NEP5Balances is a map of the NEP5 contract IDs
 // to the corresponding structures.
 type NEP5Balances struct {
-	Trackers map[util.Uint160]NEP5Tracker
+	Trackers map[int32]NEP5Tracker
 	// NextTransferBatch stores an index of the next transfer batch.
 	NextTransferBatch uint32
 }
@@ -54,7 +54,7 @@ type NEP5Balances struct {
 // NewNEP5Balances returns new NEP5Balances.
 func NewNEP5Balances() *NEP5Balances {
 	return &NEP5Balances{
-		Trackers: make(map[util.Uint160]NEP5Tracker),
+		Trackers: make(map[int32]NEP5Tracker),
 	}
 }
 
@@ -62,11 +62,10 @@ func NewNEP5Balances() *NEP5Balances {
 func (bs *NEP5Balances) DecodeBinary(r *io.BinReader) {
 	bs.NextTransferBatch = r.ReadU32LE()
 	lenBalances := r.ReadVarUint()
-	m := make(map[util.Uint160]NEP5Tracker, lenBalances)
+	m := make(map[int32]NEP5Tracker, lenBalances)
 	for i := 0; i < int(lenBalances); i++ {
-		var key util.Uint160
+		key := int32(r.ReadU32LE())
 		var tr NEP5Tracker
-		r.ReadBytes(key[:])
 		tr.DecodeBinary(r)
 		m[key] = tr
 	}
@@ -78,7 +77,7 @@ func (bs *NEP5Balances) EncodeBinary(w *io.BinWriter) {
 	w.WriteU32LE(bs.NextTransferBatch)
 	w.WriteVarUint(uint64(len(bs.Trackers)))
 	for k, v := range bs.Trackers {
-		w.WriteBytes(k[:])
+		w.WriteU32LE(uint32(k))
 		v.EncodeBinary(w)
 	}
 }

pkg/rpc/server/server.go

@@ -515,7 +515,11 @@ func (s *Server) getNEP5Balances(ps request.Params) (interface{}, *response.Erro
 	}
 	if as != nil {
 		cache := make(map[util.Uint160]int64)
-		for h, bal := range as.Trackers {
+		for id, bal := range as.Trackers {
+			h, err := s.chain.GetContractScriptHash(id)
+			if err != nil {
+				continue
+			}
 			dec, err := s.getDecimals(h, cache)
 			if err != nil {
 				continue