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,6 +979,7 @@ 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 _, ok := s.snap.(*pluginSnapshot); !ok && s.snaps != nil {
 				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)
 				}
@@ -983,6 +991,7 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 					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
 	}
 	return root, err