Implement snapshot if snapshot trie is not available

We want to be able to capture StateUpdates even if the Geth snapshot
trie is in a weird state and can't offer the snapshot we're looking
for. This adds our own implementation of the Snapshot() interface
so that we can continue collecting the necessary information to make
it available to the StateUpdates hook.

core/state/plugin_hooks.go

@@ -1,12 +1,34 @@
 package state
 
 import (
+	"fmt"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/plugins"
+	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	"github.com/openrelayxyz/plugeth-utils/core"
 )
 
+type pluginSnapshot struct {
+	root common.Hash
+}
+
+func (s *pluginSnapshot) Root() common.Hash {
+	return s.root
+}
+
+func (s *pluginSnapshot) Account(hash common.Hash) (*snapshot.Account, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (s *pluginSnapshot) AccountRLP(hash common.Hash) ([]byte, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
+func (s *pluginSnapshot) Storage(accountHash, storageHash common.Hash) ([]byte, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
 func PluginStateUpdate(pl *plugins.PluginLoader, blockRoot, parentRoot common.Hash, destructs map[common.Hash]struct{}, accounts map[common.Hash][]byte, storage map[common.Hash]map[common.Hash][]byte, codeUpdates map[common.Hash][]byte) {
 	fnList := pl.Lookup("StateUpdate", func(item interface{}) bool {
 		_, ok := item.(func(core.Hash, core.Hash, map[core.Hash]struct{}, map[core.Hash][]byte, map[core.Hash]map[core.Hash][]byte, map[core.Hash][]byte))

core/state/statedb.go

@@ -151,6 +151,13 @@ func New(root common.Hash, db Database, snaps *snapshot.Tree) (*StateDB, error)
 			sdb.snapStorage = make(map[common.Hash]map[common.Hash][]byte)
 		}
 	}
+	if sdb.snap == nil {
+		log.Debug("Snapshots not availble. Using plugin snapshot.")
+		sdb.snap = &pluginSnapshot{root}
+		sdb.snapDestructs = make(map[common.Hash]struct{})
+		sdb.snapAccounts = make(map[common.Hash][]byte)
+		sdb.snapStorage = make(map[common.Hash]map[common.Hash][]byte)
+	}
 	return sdb, nil
 }
 
@@ -972,15 +979,17 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 		// Only update if there's a state transition (skip empty Clique blocks)
 		if parent := s.snap.Root(); parent != root {
 			pluginStateUpdate(root, parent, s.snapDestructs, s.snapAccounts, s.snapStorage, codeUpdates)
-			if err := s.snaps.Update(root, parent, s.snapDestructs, s.snapAccounts, s.snapStorage); err != nil {
+			if _, ok := s.snap.(*pluginSnapshot); !ok && s.snaps != nil {
-				log.Warn("Failed to update snapshot tree", "from", parent, "to", root, "err", err)
+				if err := s.snaps.Update(root, parent, s.snapDestructs, s.snapAccounts, s.snapStorage); err != nil {
-			}
+					log.Warn("Failed to update snapshot tree", "from", parent, "to", root, "err", err)
-			// Keep 128 diff layers in the memory, persistent layer is 129th.
+				}
-			// - head layer is paired with HEAD state
+				// Keep 128 diff layers in the memory, persistent layer is 129th.
-			// - head-1 layer is paired with HEAD-1 state
+				// - head layer is paired with HEAD state
-			// - head-127 layer(bottom-most diff layer) is paired with HEAD-127 state
+				// - head-1 layer is paired with HEAD-1 state
-			if err := s.snaps.Cap(root, 128); err != nil {
+				// - head-127 layer(bottom-most diff layer) is paired with HEAD-127 state
-				log.Warn("Failed to cap snapshot tree", "root", root, "layers", 128, "err", err)
+				if err := s.snaps.Cap(root, 128); err != nil {
+					log.Warn("Failed to cap snapshot tree", "root", root, "layers", 128, "err", err)
+				}
 			}
 		}
 		s.snap, s.snapDestructs, s.snapAccounts, s.snapStorage = nil, nil, nil, nil