tendermint/state/block_cache.go

package state

import (
	"sort"

	ac "github.com/tendermint/tendermint/account"
	"github.com/tendermint/tendermint/binary"
	. "github.com/tendermint/tendermint/common"
	dbm "github.com/tendermint/tendermint/db"
	"github.com/tendermint/tendermint/merkle"
)

func makeStorage(db dbm.DB, root []byte) merkle.Tree {
	storage := merkle.NewIAVLTree(
		binary.BasicCodec,
		binary.BasicCodec,
		1024,
		db,
	)
	storage.Load(root)
	return storage
}

type BlockCache struct {
	db       dbm.DB
	backend  *State
	accounts map[string]accountInfo
	storages map[Tuple256]Word256
}

func NewBlockCache(backend *State) *BlockCache {
	return &BlockCache{
		db:       backend.DB,
		backend:  backend,
		accounts: make(map[string]accountInfo),
		storages: make(map[Tuple256]Word256),
	}
}

func (cache *BlockCache) State() *State {
	return cache.backend
}

//-------------------------------------
// BlockCache.account

func (cache *BlockCache) GetAccount(addr []byte) *ac.Account {
	acc, storage, removed := unpack(cache.accounts[string(addr)])
	if removed {
		return nil
	} else if acc != nil {
		return acc
	} else {
		acc = cache.backend.GetAccount(addr)
		storage = makeStorage(cache.db, acc.StorageRoot)
		cache.accounts[string(addr)] = accountInfo{acc, storage, false}
		return acc
	}
}

func (cache *BlockCache) UpdateAccount(acc *ac.Account) {
	addr := acc.Address
	// SANITY CHECK
	_, storage, removed := unpack(cache.accounts[string(addr)])
	if removed {
		panic("UpdateAccount on a removed account")
	}
	// SANITY CHECK END
	cache.accounts[string(addr)] = accountInfo{acc, storage, false}
}

func (cache *BlockCache) RemoveAccount(addr []byte) {
	// SANITY CHECK
	_, _, removed := unpack(cache.accounts[string(addr)])
	if removed {
		panic("RemoveAccount on a removed account")
	}
	// SANITY CHECK END
	cache.accounts[string(addr)] = accountInfo{nil, nil, true}
}

// BlockCache.account
//-------------------------------------
// BlockCache.storage

func (cache *BlockCache) GetStorage(addr Word256, key Word256) (value Word256) {
	// Check cache
	value, ok := cache.storages[Tuple256{addr, key}]
	if ok {
		return value
	}

	// Get or load storage
	_, storage, removed := unpack(cache.accounts[string(addr.Prefix(20))])
	if removed {
		panic("GetStorage() on removed account")
	}

	// Load and set cache
	_, val_ := storage.Get(key.Bytes())
	value = Zero256
	if val_ != nil {
		value = RightPadWord256(val_.([]byte))
	}
	cache.storages[Tuple256{addr, key}] = value
	return value
}

// NOTE: Set value to zero to removed from the trie.
func (cache *BlockCache) SetStorage(addr Word256, key Word256, value Word256) {
	_, _, removed := unpack(cache.accounts[string(addr.Prefix(20))])
	if removed {
		panic("SetStorage() on a removed account")
	}
	cache.storages[Tuple256{addr, key}] = value
}

// BlockCache.storage
//-------------------------------------

// CONTRACT the updates are in deterministic order.
func (cache *BlockCache) Sync() {

	// Determine order for storage updates
	// The address comes first so it'll be grouped.
	storageKeys := make([]Tuple256, 0, len(cache.storages))
	for keyTuple := range cache.storages {
		storageKeys = append(storageKeys, keyTuple)
	}
	Tuple256Slice(storageKeys).Sort()

	// Update storage for all account/key.
	// Later we'll iterate over all the users and save storage + update storage root.
	var (
		curAddr       Word256
		curAcc        *ac.Account
		curAccRemoved bool
		curStorage    merkle.Tree
	)
	for _, storageKey := range storageKeys {
		addr, key := Tuple256Split(storageKey)
		if addr != curAddr || curAcc == nil {
			acc, storage, removed := unpack(cache.accounts[string(addr.Prefix(20))])
			curAddr = addr
			curAcc = acc
			curAccRemoved = removed
			curStorage = storage
		}
		if curAccRemoved {
			continue
		}
		value := cache.storages[storageKey]
		if value.IsZero() {
			curStorage.Remove(key.Bytes())
		} else {
			curStorage.Set(key.Bytes(), value.Bytes())
		}
	}

	// Determine order for accounts
	addrStrs := []string{}
	for addrStr := range cache.accounts {
		addrStrs = append(addrStrs, addrStr)
	}
	sort.Strings(addrStrs)

	// Update or delete accounts.
	for _, addrStr := range addrStrs {
		acc, storage, removed := unpack(cache.accounts[addrStr])
		if removed {
			removed := cache.backend.RemoveAccount(acc.Address)
			if !removed {
				panic(Fmt("Could not remove account to be removed: %X", acc.Address))
			}
		} else {
			if acc == nil {
				panic(Fmt("Account should not be nil for addr: %X", acc.Address))
			}
			acc.StorageRoot = storage.Save()
			cache.backend.UpdateAccount(acc)
		}
	}

}

//-----------------------------------------------------------------------------

type accountInfo struct {
	account *ac.Account
	storage merkle.Tree
	removed bool
}

func unpack(accInfo accountInfo) (*ac.Account, merkle.Tree, bool) {
	return accInfo.account, accInfo.storage, accInfo.removed
}