From 1ab4f81fc3b9ba7a99e98cd1b3eadbcdbea12b47 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 6 Jul 2020 17:17:49 +0300
Subject: [PATCH] dao: migrate nep5 balances with the contract

Fixes #1144. It's quite simple approach, we just update balance info right
upon contract migration. It will slow down migration transactions, but it
takes about 1-2 seconds to Seek through balances at mainnet's 3.8M, so the
approach should still work good enough. The other idea was to make lazy
updates (maintaining contract migration map), but it's more complicated to
implement (and implies that a balance get might also do a write).

There also is a concern about memory usage, it can give a spike of some tens
of megabytes, but that also is considered to be acceptable.
---
 pkg/core/dao/cacheddao.go | 68 ++++++++++++++++++++++++++++++++++++++-
 pkg/core/interop_neo.go   |  1 +
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/pkg/core/dao/cacheddao.go b/pkg/core/dao/cacheddao.go
index fc905f919..3a64b2b58 100644
--- a/pkg/core/dao/cacheddao.go
+++ b/pkg/core/dao/cacheddao.go
@@ -6,6 +6,7 @@ import (
 	"sort"
 
 	"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"
 )
@@ -21,6 +22,8 @@ type Cached struct {
 	balances  map[util.Uint160]*state.NEP5Balances
 	transfers map[util.Uint160]map[uint32]*state.NEP5TransferLog
 	storage   *itemCache
+
+	dropNEP5Cache bool
 }
 
 // NewCached returns new Cached wrapping around given backing store.
@@ -39,7 +42,7 @@ func NewCached(d DAO) *Cached {
 			}
 		}
 	}
-	return &Cached{dao, accs, ctrs, unspents, balances, transfers, st}
+	return &Cached{dao, accs, ctrs, unspents, balances, transfers, st, false}
 }
 
 // GetAccountStateOrNew retrieves Account from cache or underlying store
@@ -149,6 +152,65 @@ 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) {
@@ -162,6 +224,9 @@ func (cd *Cached) Persist() (int, error) {
 	// usage scenario it should be good enough if cd doesn't modify object
 	// caches (accounts/contracts/etc) in any way.
 	if ok {
+		if cd.dropNEP5Cache {
+			lowerCache.balances = make(map[util.Uint160]*state.NEP5Balances)
+		}
 		var simpleCache *Simple
 		for simpleCache == nil {
 			if err := lowerCache.FlushStorage(); err != nil {
@@ -220,6 +285,7 @@ func (cd *Cached) GetWrapped() DAO {
 		cd.balances,
 		cd.transfers,
 		cd.storage,
+		false,
 	}
 }
 
diff --git a/pkg/core/interop_neo.go b/pkg/core/interop_neo.go
index 418f943a6..ef69276ac 100644
--- a/pkg/core/interop_neo.go
+++ b/pkg/core/interop_neo.go
@@ -597,6 +597,7 @@ func (ic *interopContext) contractMigrate(v *vm.VM) error {
 					return err
 				}
 			}
+			ic.dao.MigrateNEP5Balances(hash, contract.ScriptHash())
 		}
 	}
 	v.Estack().PushVal(vm.NewInteropItem(contract))