Fixed end root problem

mpt-witness-generator/README.md

@@ -71,14 +71,12 @@ What is changed compared to geth:
 
 ## Generate witnesses
 
-To generate witnesses for MPT circuit, go into witness folder and execute
+To generate witnesses for MPT circuit, execute
 
 ```
-go test gen_witness_from_infura_blockchain_test.go prepare_witness.go leaf.go extension_node.go modified_extension_node.go nodes.go test_tools.go branch.go util.go
+go test -v ./...
 ```
 
-to generate the tests that use Infura blockchain.
-
 To generate the tests that use a local blockchain you need a local `geth`. You would
 need to run something like:
 ```
@@ -93,14 +91,16 @@ database if you already have some accounts:
 geth removedb
 ```
 
-And to generate the tests:
+The witness files will appear in generated_witnesses folder.
+
+## Format the code
+
+To format the code use:
 
 ```
-go test gen_witness_from_local_blockchain_test.go prepare_witness.go leaf.go extension_node.go modified_extension_node.go nodes.go test_tools.go branch.go util.go
+gofmt -w ./*
 ```
 
-The witness files will appear in generated_witnesses folder.
-
 ## Calling from Rust
 
 Build:

mpt-witness-generator/state/statedb.go

@@ -410,6 +410,15 @@ func (s *StateDB) SetCode(addr common.Address, code []byte) {
 	}
 }
 
+func (s *StateDB) SetCodeHash(addr common.Address, codeHash []byte) {
+	s.SetStateObjectIfExists(addr)
+	stateObject := s.GetOrNewStateObject(addr)
+	if stateObject != nil {
+		// Only codeHash is the correct one, but we don't need the actual code for the MPT witness generator.
+		stateObject.SetCode(common.BytesToHash(codeHash), codeHash)
+	}
+}
+
 func (s *StateDB) SetState(addr common.Address, key, value common.Hash) {
 	s.SetStateObjectIfExists(addr)
 	stateObject := s.GetOrNewStateObject(addr)

mpt-witness-generator/witness/prepare_witness.go

@@ -1,7 +1,6 @@
 package witness
 
 import (
-	"encoding/hex"
 	"fmt"
 	"math/big"
 
@@ -116,7 +115,10 @@ func obtainAccountProofAndConvertToWitness(i int, tMod TrieModification, tModsLe
 	} else if tMod.Type == BalanceChanged {
 		statedb.SetBalance(addr, tMod.Balance)
 	} else if tMod.Type == CodeHashChanged {
-		statedb.SetCode(addr, tMod.CodeHash)
+		statedb.SetCodeHash(addr, tMod.CodeHash)
+		// For cases when the wrong account is obtained by PrefetchAccount:
+		statedb.SetBalance(addr, big.NewInt(0))
+		statedb.SetNonce(addr, 0)
 	} else if tMod.Type == AccountCreate {
 		statedb.CreateAccount(tMod.Address)
 	} else if tMod.Type == AccountDestructed {
@@ -201,7 +203,7 @@ func obtainTwoProofsAndConvertToWitness(trieModifications []TrieModification, st
 			addrh := crypto.Keccak256(addr.Bytes())
 			accountAddr := trie.KeybytesToHex(addrh)
 
-			ap := oracle.PrefetchAccount(statedb.Db.BlockNumber, tMod.Address, nil)
+			oracle.PrefetchAccount(statedb.Db.BlockNumber, tMod.Address, nil)
 			oracle.PrefetchStorage(statedb.Db.BlockNumber, addr, tMod.Key, nil)
 
 			if specialTest == 1 {
@@ -211,9 +213,13 @@ func obtainTwoProofsAndConvertToWitness(trieModifications []TrieModification, st
 			accountProof, aNeighbourNode1, aExtNibbles1, aIsLastLeaf1, aIsNeighbourNodeHashed1, err := statedb.GetProof(addr)
 			check(err)
 
-			if !statedb.Exist(addr) && len(ap) > 0 {
-				ret, _ := hex.DecodeString(ap[len(ap)-1][2:])
-				statedb.SetStateObjectFromEncoding(addr, ret)
+			// When the account has not been created yet and PrefetchAccount gets the wrong
+			// account - because the first part of the address is the same and
+			// the queried address doesn't have the account yet.
+			if !statedb.Exist(addr) {
+				// Note: the storage modification should not be the first modification for the account that does
+				// not exist yet.
+				panic("The account should exist at this point - created by SetNonce, SetBalance, or SetCodehash")
 			}
 
 			storageProof, neighbourNode1, extNibbles1, isLastLeaf1, isNeighbourNodeHashed1, err := statedb.GetStorageProof(addr, tMod.Key)