wip: implement block-level witness dumping

core/blockchain.go

@@ -1484,10 +1484,14 @@ func (bc *BlockChain) WriteBlockAndSetHead(block *types.Block, receipts []*types
 
 // writeBlockAndSetHead is the internal implementation of WriteBlockAndSetHead.
 // This function expects the chain mutex to be held.
-func (bc *BlockChain) writeBlockAndSetHead(block *types.Block, receipts []*types.Receipt, logs []*types.Log, state *state.StateDB, emitHeadEvent bool) (status WriteStatus, err error) {
-	if err := bc.writeBlockWithState(block, receipts, state); err != nil {
+func (bc *BlockChain) writeBlockAndSetHead(block *types.Block, receipts []*types.Receipt, logs []*types.Log, st *state.StateDB, emitHeadEvent bool) (status WriteStatus, err error) {
+	if err := bc.writeBlockWithState(block, receipts, st); err != nil {
 		return NonStatTy, err
 	}
+
+	st.Witness.Block = block
+	state.DumpWitnessToFile(st.Witness)
+
 	currentBlock := bc.CurrentBlock()
 	reorg, err := bc.forker.ReorgNeeded(currentBlock, block.Header())
 	if err != nil {

core/blockchain_test.go

@@ -17,6 +17,7 @@
 package core
 
 import (
+	"crypto/sha256"
 	"errors"
 	"fmt"
 	"math/big"
@@ -4717,7 +4718,29 @@ func TestEIP3651(t *testing.T) {
 	}
 }
 
-func TestSingleWitness(t *testing.T) {
+func hashAccount(addr common.Address) (output common.Hash) {
+	hasher := sha256.New()
+	hasher.Write(addr[:])
+	res := hasher.Sum(nil)
+	copy(output[:], res[:])
+	return output
+}
+
+/*
+func debugPrintTrie(db ethdb.Database, account common.Address) {
+	var bc BlockChain
+	latestBlock := bc.GetBlockByNumber(1)
+	storageRoot := bc.StateAt(latestBlock.Root()).GetStorageRoot(account)
+
+	tr, err := trie.New(trie.StorageTrieID(latestBlock.Root(), hashAccount(account), storageRoot), s.Database())
+	if err != nil {
+		panic(err)
+	}
+	trIterator
+}
+*/
+
+func TestWitnessStorageClear(t *testing.T) {
 	var (
 		engine = ethash.NewFaker()
 
@@ -4728,15 +4751,14 @@ func TestSingleWitness(t *testing.T) {
 		bb        = common.HexToAddress("0x000000000000000000000000000000000000bbbb")
 		aaStorage = make(map[common.Hash]common.Hash) // Initial storage in AA
 	)
-	// Populate two slots
+	// Populate one slots
 	aaStorage[common.HexToHash("01")] = common.HexToHash("01")
+	aaStorage[common.HexToHash("02")] = common.HexToHash("01")
 
 	code := []byte{
-		byte(vm.PUSH1), 0x1, //
-		byte(vm.NUMBER),     // value = number + 1
-		byte(vm.ADD),        //
-		byte(vm.PUSH1), 0x3, // location
-		byte(vm.SSTORE), // Set slot[3] = number + 1
+		byte(vm.PUSH1), 0x0, // value
+		byte(vm.PUSH1), 0x2, // key
+		byte(vm.SSTORE),
 	}
 	gspec := &Genesis{
 		Config: params.TestChainConfig,
@@ -4746,6 +4768,7 @@ func TestSingleWitness(t *testing.T) {
 			bb: {
 				Code:    code,
 				Balance: big.NewInt(1),
+				Storage: aaStorage,
 			},
 		},
 	}
@@ -4761,8 +4784,9 @@ func TestSingleWitness(t *testing.T) {
 	// Import the canonical chain
 	cache := DefaultCacheConfigWithScheme(rawdb.PathScheme)
 	cache.SnapshotLimit = 0 // disable snapshot
-	cache.SnapshotLimit = 500 // enable snapshot
-	chain, err := NewBlockChain(rawdb.NewMemoryDatabase(), cache, gspec, nil, engine, vm.Config{
+	//cache.SnapshotLimit = 500 // enable snapshot
+	db := rawdb.NewMemoryDatabase()
+	chain, err := NewBlockChain(db, cache, gspec, nil, engine, vm.Config{
 		Tracer: logger.NewJSONLogger(nil, os.Stdout),
 	}, nil, nil)
 	if err != nil {

core/state/state_witness.go

@@ -1,36 +1,90 @@
 package state
 
 import (
+	"bytes"
 	"fmt"
-
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/rlp"
+	"io"
+	"os"
 )
 
-type witness struct {
+type Witness struct {
+	Block *types.Block
+	Root  common.Hash
+	Lists map[common.Hash]map[string][]byte
+}
+
+type EncodeWitness struct {
+	block  types.Block
 	root   common.Hash
-	lists  []map[string][]byte
 	owners []common.Hash
+	paths  [][]string
+	nodes  [][][]byte
 }
 
-func newWitness(originalRoot common.Hash) *witness {
-	return &witness{root: originalRoot}
+func (w *Witness) EncodeRLP() []byte {
+	var e EncodeWitness
+	for owner, nodeMap := range w.Lists {
+		e.owners = append(e.owners, owner)
+		var paths []string
+		var nodes [][]byte
+
+		for path, node := range nodeMap {
+			paths = append(paths, path)
+			nodes = append(nodes, node)
+		}
+		e.paths = append(e.paths, paths)
+		e.nodes = append(e.nodes, nodes)
+	}
+	res := new(bytes.Buffer)
+	if err := e.encode(res); err != nil {
+		panic(err)
+	}
+	return res.Bytes()
 }
 
-func (w *witness) addAccessList(owner common.Hash, list map[string][]byte) {
-	//fmt.Printf("Adding owner %x len %d\n", owner, len(list))
+// generated by rlpgen
+func (obj *EncodeWitness) encode(_w io.Writer) error {
+	w := rlp.NewEncoderBuffer(_w)
+	_tmp0 := w.List()
+	w.ListEnd(_tmp0)
+	return w.Flush()
+}
+
+func newWitness(originalRoot common.Hash) *Witness {
+	return &Witness{Root: originalRoot}
+}
+
+func (w *Witness) addAccessList(owner common.Hash, list map[string][]byte) {
 	if len(list) > 0 {
-		w.lists = append(w.lists, list)
-		w.owners = append(w.owners, owner)
+		w.Lists[owner] = list
 	}
 }
 
-func (w *witness) Dump() {
-	fmt.Printf("[witness] Root %x\n", w.root)
-	for i, list := range w.lists {
-		owner := w.owners[i]
+func (w *Witness) Dump() {
+	/*
+		fmt.Printf("[witness] Root %x\n", w.Root)
+		for i, list := range w.Lists {
+			owner := w.Owners[i]
 			fmt.Printf("[witness] Owner %#x, %d entries: \n", owner, len(list))
 			for path, v := range list {
 				fmt.Printf("[witness] - '%#x': %#x\n", path, v)
 			}
 		}
+	*/
+}
+
+func DumpWitnessToFile(w *Witness) {
+	enc := w.EncodeRLP()
+	path, err := os.Getwd()
+	if err != nil {
+		panic("shite")
+	}
+	outputFName := fmt.Sprintf("%d-%x.rlp\n", w.Block.NumberU64(), w.Block.Hash())
+	err = os.WriteFile(path+"/"+outputFName, enc, 0644)
+	if err != nil {
+		panic("shite 2")
+	}
 }

core/state/statedb.go

@@ -84,6 +84,9 @@ type StateDB struct {
 	stateObjectsDirty    map[common.Address]struct{}            // State objects modified in the current execution
 	stateObjectsDestruct map[common.Address]*types.StateAccount // State objects destructed in the block along with its previous value
 
+	usedBlockHashes map[common.Hash]struct{}
+	codes           map[common.Hash]Code
+
 	// DB error.
 	// State objects are used by the consensus core and VM which are
 	// unable to deal with database-level errors. Any error that occurs
@@ -138,6 +141,8 @@ type StateDB struct {
 
 	// Testing hooks
 	onCommit func(states *triestate.Set) // Hook invoked when commit is performed
+
+	Witness *Witness
 }
 
 // New creates a new state from a given trie.
@@ -172,6 +177,14 @@ func New(root common.Hash, db Database, snaps *snapshot.Tree) (*StateDB, error)
 	return sdb, nil
 }
 
+func (s *StateDB) MarkUsedBlockHash(hash common.Hash) {
+	s.usedBlockHashes[hash] = struct{}{}
+}
+
+func (s *StateDB) MarkWitnessCode(hash common.Hash, code Code) {
+	s.codes[hash] = code
+}
+
 // StartPrefetcher initializes a new trie prefetcher to pull in nodes from the
 // state trie concurrently while the state is mutated so that when we reach the
 // commit phase, most of the needed data is already hot.
@@ -316,8 +329,15 @@ func (s *StateDB) TxIndex() int {
 func (s *StateDB) GetCode(addr common.Address) []byte {
 	stateObject := s.getStateObject(addr)
 	if stateObject != nil {
-		return stateObject.Code()
+		code := stateObject.Code()
+		if code != nil {
+			var codeHash common.Hash
+			copy(codeHash[:], stateObject.CodeHash())
+			s.MarkWitnessCode(codeHash, stateObject.Code())
 		}
+		return code
+	}
+
 	return nil
 }
 
@@ -1156,6 +1176,19 @@ func (s *StateDB) handleDestruction(nodes *trienode.MergedNodeSet) (map[common.A
 	return incomplete, nil
 }
 
+type BlockProof struct {
+	blockHashes []common.Hash
+	codeHashes  []common.Hash
+	codes       []Code
+	witness     Witness
+}
+
+/*
+func (s *StateDB) GetBlockProof() {
+
+}
+*/
+
 // Commit writes the state to the underlying in-memory trie database.
 // Once the state is committed, tries cached in stateDB (including account
 // trie, storage tries) will no longer be functional. A new state instance
@@ -1306,7 +1339,8 @@ func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool) (common.Hash, er
 	s.stateObjectsDirty = make(map[common.Address]struct{})
 	s.stateObjectsDestruct = make(map[common.Address]*types.StateAccount)
 
-	w.Dump()
+	//w.Dump()
+	s.Witness = w
 
 	return root, nil
 }

core/vm/instructions.go

@@ -447,7 +447,11 @@ func opBlockhash(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) (
 		lower = upper - 256
 	}
 	if num64 >= lower && num64 < upper {
-		num.SetBytes(interpreter.evm.Context.GetHash(num64).Bytes())
+		res := interpreter.evm.Context.GetHash(num64).Bytes()
+		var bh common.Hash
+		copy(bh[:], res[:])
+		interpreter.evm.StateDB.MarkUsedBlockHash(bh)
+		num.SetBytes(res[:])
 	} else {
 		num.Clear()
 	}

core/vm/interface.go

@@ -78,6 +78,8 @@ type StateDB interface {
 
 	AddLog(*types.Log)
 	AddPreimage(common.Hash, []byte)
+
+	MarkUsedBlockHash(common.Hash)
 }
 
 // CallContext provides a basic interface for the EVM calling conventions. The EVM