snapshot with blockhashes

Author: Darioush Jalali

core/state/snapshot/snapshot.go

@@ -174,6 +174,10 @@ type Tree struct {
 
 	// Test hooks
 	onFlatten func() // Hook invoked when the bottom most diff layers are flattened
+
+	// XXX
+	children        map[common.Hash][]common.Hash
+	initiallyLoaded map[common.Hash]struct{}
 }
 
 // New attempts to load an already existing snapshot from a persistent key-value
@@ -195,10 +199,12 @@ type Tree struct {
 func New(config Config, diskdb ethdb.KeyValueStore, triedb *triedb.Database, root common.Hash) (*Tree, error) {
 	// Create a new, empty snapshot tree
 	snap := &Tree{
-		config: config,
-		diskdb: diskdb,
-		triedb: triedb,
-		layers: make(map[common.Hash]snapshot),
+		config:          config,
+		diskdb:          diskdb,
+		triedb:          triedb,
+		layers:          make(map[common.Hash]snapshot),
+		children:        make(map[common.Hash][]common.Hash),
+		initiallyLoaded: map[common.Hash]struct{}{root: {}},
 	}
 	// Attempt to load a previously persisted snapshot and rebuild one if failed
 	head, disabled, err := loadSnapshot(diskdb, triedb, root, config.CacheSize, config.Recovery, config.NoBuild)
@@ -221,6 +227,7 @@ func New(config Config, diskdb ethdb.KeyValueStore, triedb *triedb.Database, roo
 	// Existing snapshot loaded, seed all the layers
 	for head != nil {
 		snap.layers[head.Root()] = head
+		snap.initiallyLoaded[head.Root()] = struct{}{}
 		head = head.Parent()
 	}
 	return snap, nil
@@ -310,7 +317,17 @@ func (t *Tree) Snapshot(blockRoot common.Hash) Snapshot {
 	t.lock.RLock()
 	defer t.lock.RUnlock()
 
-	return t.layers[blockRoot]
+	return t.byRoot(blockRoot)
+}
+
+func (t *Tree) byRoot(blockRoot common.Hash) snapshot {
+	for _, snap := range t.layers {
+		if snap.Root() == blockRoot {
+			return snap
+		}
+	}
+
+	return nil
 }
 
 // Snapshots returns all visited layers from the topmost layer with specific
@@ -323,7 +340,7 @@ func (t *Tree) Snapshots(root common.Hash, limits int, nodisk bool) []Snapshot {
 	if limits == 0 {
 		return nil
 	}
-	layer := t.layers[root]
+	layer := t.byRoot(root)
 	if layer == nil {
 		return nil
 	}
@@ -348,28 +365,47 @@ func (t *Tree) Snapshots(root common.Hash, limits int, nodisk bool) []Snapshot {
 
 // Update adds a new snapshot into the tree, if that can be linked to an existing
 // old parent. It is disallowed to insert a disk layer (the origin of all).
-func (t *Tree) Update(blockRoot common.Hash, parentRoot common.Hash, destructs map[common.Hash]struct{}, accounts map[common.Hash][]byte, storage map[common.Hash]map[common.Hash][]byte) error {
+func (t *Tree) Update(blockRoot common.Hash, parentRoot common.Hash, destructs map[common.Hash]struct{}, accounts map[common.Hash][]byte, storage map[common.Hash]map[common.Hash][]byte,
+	opts ...LibEVMOption) error {
 	// Reject noop updates to avoid self-loops in the snapshot tree. This is a
 	// special case that can only happen for Clique networks where empty blocks
 	// don't modify the state (0 block subsidy).
 	//
 	// Although we could silently ignore this internally, it should be the caller's
 	// responsibility to avoid even attempting to insert such a snapshot.
-	if blockRoot == parentRoot {
+	opt := asLibEVMConfig(opts)
+	layer := blockRoot
+	if opt.hash != (common.Hash{}) {
+		layer = opt.hash
+	} else {
+		panic("block hash is required")
+	}
+
+	parentLayer := opt.parentHash
+	_, parentHashKnown := t.layers[opt.parentHash]
+	_, isInitiallyLoaded := t.initiallyLoaded[parentRoot]
+	if !parentHashKnown && isInitiallyLoaded {
+		// Allow use of parentRoot as the key in the tree if:
+		// - opt.parentHash is not known,
+		// - parentRoot is known to be initially loaded.
+		parentLayer = parentRoot
+	}
+	if layer == parentLayer {
 		return errSnapshotCycle
 	}
 	// Generate a new snapshot on top of the parent
-	parent := t.Snapshot(parentRoot)
+	parent := t.layers[parentLayer]
 	if parent == nil {
 		return fmt.Errorf("parent [%#x] snapshot missing", parentRoot)
 	}
 	snap := parent.(snapshot).Update(blockRoot, destructs, accounts, storage)
+	t.children[parentLayer] = append(t.children[parentLayer], layer)
 
 	// Save the new snapshot for later
 	t.lock.Lock()
 	defer t.lock.Unlock()
 
-	t.layers[snap.root] = snap
+	t.layers[layer] = snap
 	return nil
 }
 

core/state/snapshot/snapshot.libevm.go

@@ -0,0 +1,177 @@
+// The libevm additions to go-ethereum are free software: you can redistribute
+// them and/or modify them under the terms of the GNU Lesser General Public License
+// as published by the Free Software Foundation, either version 3 of the License,
+// or (at your option) any later version.
+//
+// The libevm additions are distributed in the hope that they will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser
+// General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see
+// <http://www.gnu.org/licenses/>.
+
+package snapshot
+
+import (
+	"fmt"
+
+	"github.com/ava-labs/libevm/common"
+)
+
+// A LibEVMOption configures default behaviour of this package.
+type LibEVMOption interface {
+	apply(*libevmConfig)
+}
+
+type libevmConfig struct {
+	preserveDescendantsOnCapZero bool
+	hash, parentHash             common.Hash
+}
+
+func asLibEVMConfig(opts []LibEVMOption) *libevmConfig {
+	c := new(libevmConfig)
+	for _, o := range opts {
+		o.apply(c)
+	}
+	return c
+}
+
+type libevmFuncOpt func(*libevmConfig)
+
+func (f libevmFuncOpt) apply(c *libevmConfig) { f(c) }
+
+// PreserveDescendantsOnCapZero signals to [Tree.Cap], if capping to zero layers
+// (i.e. flattening), that descendants of the flattened layer must be kept.
+// Without this option, the entire tree is replaced with the new base (disk)
+// layer.
+func PreserveDescendantsOnCapZero() LibEVMOption {
+	return libevmFuncOpt(func(c *libevmConfig) {
+		c.preserveDescendantsOnCapZero = true
+	})
+}
+
+func WithBlockHashes(hash, parentHash common.Hash) LibEVMOption {
+	return libevmFuncOpt(func(c *libevmConfig) {
+		c.hash = hash
+		c.parentHash = parentHash
+	})
+}
+
+func (t *Tree) updateLayersAfterCapZero(base *diskLayer, flattened *diffLayer, opts ...LibEVMOption) {
+	// Original geth behaviour
+	opt := asLibEVMConfig(opts)
+	if !opt.preserveDescendantsOnCapZero {
+		t.layers = map[common.Hash]snapshot{base.root: base}
+		t.children = make(map[common.Hash][]common.Hash)
+		return
+	}
+
+	children := t.children
+	baseHash := base.root
+	if opt.hash != (common.Hash{}) {
+		baseHash = opt.hash
+	}
+
+	newLayers := map[common.Hash]snapshot{baseHash: base}
+	newChildren := make(map[common.Hash][]common.Hash)
+	var keepChildren func(root common.Hash)
+	keepChildren = func(root common.Hash) {
+		for _, child := range children[root] {
+			childLayer := t.layers[child]
+			newLayers[child] = childLayer
+			newChildren[root] = append(newChildren[root], child)
+			keepChildren(child)
+		}
+	}
+	keepChildren(baseHash)
+
+	for _, child := range children[baseHash] {
+		d, ok := t.layers[child].(*diffLayer)
+		if !ok {
+			continue
+		}
+		d.lock.Lock()
+		d.parent = base
+		d.lock.Unlock()
+	}
+	t.layers = newLayers
+	t.children = newChildren
+}
+
+func (t *Tree) Flatten(hash common.Hash) error {
+	t.lock.Lock()
+	defer t.lock.Unlock()
+
+	// Retrieve the head snapshot to cap from
+	snap := t.layers[hash]
+	if snap == nil {
+		return fmt.Errorf("snapshot [%#x] missing", hash)
+	}
+	diff, ok := snap.(*diffLayer)
+	if !ok {
+		return fmt.Errorf("snapshot [%#x] is disk layer", hash)
+	}
+
+	diff.lock.RLock()
+	base := diffToDisk(diff.flatten().(*diffLayer))
+	diff.lock.RUnlock()
+
+	t.updateLayersAfterCapZero(
+		base,
+		diff,
+		PreserveDescendantsOnCapZero(),
+		WithBlockHashes(hash, common.Hash{}),
+	)
+	return nil
+}
+
+func (t *Tree) Discard(hash common.Hash) error {
+	return nil
+}
+
+func (t *Tree) AbortGeneration() {
+	t.lock.Lock()
+	defer t.lock.Unlock()
+
+	dl := t.disklayer()
+
+	dl.lock.Lock()
+	if dl.genAbort == nil {
+		dl.lock.Unlock()
+		return
+	}
+	if dl.genAbort != nil {
+		abort := make(chan *generatorStats)
+		dl.genAbort <- abort
+		dl.genAbort = nil
+		dl.lock.Unlock()
+		<-abort
+	}
+}
+
+func (t *Tree) NumStateLayers() int {
+	t.lock.RLock()
+	defer t.lock.RUnlock()
+
+	return len(t.layers)
+}
+
+func (t *Tree) NumBlockLayers() int {
+	return t.NumStateLayers()
+}
+func (t *Tree) DiskAccountIterator(seek common.Hash) AccountIterator {
+	t.lock.Lock()
+	defer t.lock.Unlock()
+
+	return t.disklayer().AccountIterator(seek)
+}
+
+func (t *Tree) DiskStorageIterator(account common.Hash, seek common.Hash) StorageIterator {
+	t.lock.Lock()
+	defer t.lock.Unlock()
+
+	it, _ := t.disklayer().StorageIterator(account, seek)
+	return it
+}