Modified extension gadget (#1685)

### Description

This PR will replace old modified extension node PRs (there were
multiple because the witness generator was in a separate repository at
that time).

Support for modified extension nodes.

This happens when an extension node is replaced by another (shorter or
longer) extension node. For example, we have an extension node at `n1 n2
n3 n4 n5 n6` with branch with two leaves at positions `n` and `m`. If we
add a leaf at `n1 n2 n3 n4 n7` where `n5 != n7`, a new extension node is
inserted at `n1 n2 n3 n4` with a new branch with an extension node at
position `n5` (with only one nibble `n6`) and a leaf at position `n7`.
In this case, the S proof contains the extension node at `n1 n2 n3 n4 n5
n6` and no underlying branch and leaf (the modification happens at `n1
n2 n3 n4 n7` and only an extension node is find there), thus we need to
add a placeholder branch and a placeholder leaf.

### Type of change

- [x] New feature (non-breaking change which adds functionality)

---------

Co-authored-by: Chih Cheng Liang <[email protected]>

Author: miha-stopar

mpt-witness-generator/oracle/prefetch.go

@@ -6,11 +6,11 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"log"
 	"math/big"
 	"net/http"
 	"os"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
@@ -88,7 +88,7 @@ func toFilename(key string) string {
 }
 
 func cacheRead(key string) []byte {
-	dat, err := ioutil.ReadFile(toFilename(key))
+	dat, err := os.ReadFile(toFilename(key))
 	if err == nil {
 		return dat
 	}
@@ -101,19 +101,27 @@ func cacheExists(key string) bool {
 }
 
 func cacheWrite(key string, value []byte) {
-	ioutil.WriteFile(toFilename(key), value, 0644)
+	os.WriteFile(toFilename(key), value, 0644)
 }
 
 func getAPI(jsonData []byte) io.Reader {
 	key := hexutil.Encode(crypto.Keccak256(jsonData))
-	/* Note: switching between two testnets (to prepare tests with account in the first level)
-	if cacheExists(key) {
-		return bytes.NewReader(cacheRead(key))
+	var (
+		err     error
+		resp    *http.Response
+		retries int = 3
+	)
+	for retries > 0 {
+		resp, err = http.Post(NodeUrl, "application/json", bytes.NewBuffer(jsonData))
+		if err != nil {
+			retries -= 1
+			time.Sleep(1000)
+		} else {
+			break
+		}
 	}
-	*/
-	resp, _ := http.Post(NodeUrl, "application/json", bytes.NewBuffer(jsonData))
 	defer resp.Body.Close()
-	ret, _ := ioutil.ReadAll(resp.Body)
+	ret, _ := io.ReadAll(resp.Body)
 	cacheWrite(key, ret)
 	return bytes.NewReader(ret)
 }
@@ -239,13 +247,8 @@ func PrefetchBlock(blockNumber *big.Int, startBlock bool, hasher types.TrieHashe
 	r.Params[1] = true
 	jsonData, _ := json.Marshal(r)
 
-	/*dat, _ := ioutil.ReadAll(getAPI(jsonData))
-	fmt.Println(string(dat))*/
-
 	jr := jsonrespt{}
 	check(json.NewDecoder(getAPI(jsonData)).Decode(&jr))
-	//fmt.Println(jr.Result)
-	// blockHeader := types.Header(jr.Result)
 	blockHeader := jr.Result.ToHeader()
 
 	// put in the start block header
@@ -277,7 +280,7 @@ func PrefetchBlock(blockNumber *big.Int, startBlock bool, hasher types.TrieHashe
 		saveinput = append(saveinput, inputs[i].Bytes()[:]...)
 	}
 	key := fmt.Sprintf("/tmp/eth/%d", blockNumber.Uint64()-1)
-	ioutil.WriteFile(key, saveinput, 0644)
+	os.WriteFile(key, saveinput, 0644)
 
 	// save the txs
 	txs := make([]*types.Transaction, len(jr.Result.Transactions))
@@ -308,7 +311,11 @@ func getProofAccount(blockNumber *big.Int, addr common.Address, skey common.Hash
 	json.NewDecoder(getAPI(jsonData)).Decode(&jr)
 
 	if storage {
-		return jr.Result.StorageProof[0].Proof
+		if len(jr.Result.StorageProof) != 0 {
+			return jr.Result.StorageProof[0].Proof
+		} else {
+			return []string{}
+		}
 	} else {
 		return jr.Result.AccountProof
 	}

mpt-witness-generator/witness/branch.go

@@ -25,6 +25,7 @@ func isBranch(proofEl []byte) bool {
 // S occupies the first 34 columns, C occupies the next 34 columns.
 // The branch children are positioned each in its own row.
 func prepareBranchWitness(rows [][]byte, branch []byte, branchStart int, branchRLPOffset int) {
+	branchNodeRLPLen := 2 // we have two positions for RLP meta data
 	rowInd := 1
 	colInd := branchNodeRLPLen - 1
 
@@ -65,7 +66,7 @@ func prepareBranchWitness(rows [][]byte, branch []byte, branchStart int, branchR
 }
 
 func prepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes []byte, extValues [][]byte, key, driftedInd,
-	branchC16, branchC1 byte, isBranchSPlaceholder, isBranchCPlaceholder, isExtension, isSModExtension, isCModExtension bool) Node {
+	branchC16, branchC1 byte, isBranchSPlaceholder, isBranchCPlaceholder, isExtension bool) Node {
 	extensionNode := ExtensionNode{
 		ListRlpBytes: extListRlpBytes,
 	}
@@ -112,11 +113,10 @@ func prepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes []b
 	}
 
 	extensionBranch := ExtensionBranchNode{
-		IsExtension:    isExtension,
-		IsModExtension: [2]bool{isSModExtension, isCModExtension},
-		IsPlaceholder:  [2]bool{isBranchSPlaceholder, isBranchCPlaceholder},
-		Extension:      extensionNode,
-		Branch:         branchNode,
+		IsExtension:   isExtension,
+		IsPlaceholder: [2]bool{isBranchSPlaceholder, isBranchCPlaceholder},
+		Extension:     extensionNode,
+		Branch:        branchNode,
 	}
 
 	values := make([][]byte, 17)
@@ -282,15 +282,6 @@ func addBranchAndPlaceholder(proof1, proof2,
 
 	// Note that isModifiedExtNode happens also when we have a branch instead of shortExtNode
 	isModifiedExtNode := !isBranch(longExtNode) && !isShorterProofLastLeaf
-	isSModifiedExtNode := false
-	isCModifiedExtNode := false
-	if isModifiedExtNode {
-		if len1 < len2 {
-			isSModifiedExtNode = true
-		} else {
-			isCModifiedExtNode = true
-		}
-	}
 
 	if len1 > len2 {
 		// We now get the first nibble of the leaf that was turned into branch.
@@ -300,7 +291,7 @@ func addBranchAndPlaceholder(proof1, proof2,
 
 		node = prepareBranchNode(proof1[len1-2], proof1[len1-2], extNode, extNode, extListRlpBytes, extValues,
 			key[keyIndex+numberOfNibbles], driftedInd,
-			branchC16, branchC1, false, true, isExtension, isSModifiedExtNode, isCModifiedExtNode)
+			branchC16, branchC1, false, true, isExtension)
 
 		// We now get the first nibble of the leaf that was turned into branch.
 		// This first nibble presents the position of the leaf once it moved
@@ -313,7 +304,7 @@ func addBranchAndPlaceholder(proof1, proof2,
 
 		node = prepareBranchNode(proof2[len2-2], proof2[len2-2], extNode, extNode, extListRlpBytes, extValues,
 			key[keyIndex+numberOfNibbles], driftedInd,
-			branchC16, branchC1, true, false, isExtension, isSModifiedExtNode, isCModifiedExtNode)
+			branchC16, branchC1, true, false, isExtension)
 	}
 
 	return isModifiedExtNode, isExtension, numberOfNibbles, branchC16, node

mpt-witness-generator/witness/extension_node.go

@@ -1,97 +1,5 @@
 package witness
 
-// setExtNodeSelectors sets in the branch init row the information about the extension node.
-func setExtNodeSelectors(row, proofEl []byte, numberOfNibbles int, branchC16 byte) {
-	row[isExtensionPos] = 1
-	if len(proofEl) > 56 { // 56 because there is 1 byte for length
-		// isCExtLongerThan55 doesn't need to be set here
-		row[isSExtLongerThan55Pos] = 1
-	}
-
-	if len(proofEl) < 32 {
-		// isExtNodeSNonHashed doesn't need to be set here
-		row[isExtNodeSNonHashedPos] = 1
-	}
-
-	if numberOfNibbles == 1 {
-		if branchC16 == 1 {
-			row[isExtShortC16Pos] = 1
-		} else {
-			row[isExtShortC1Pos] = 1
-		}
-	} else {
-		if numberOfNibbles%2 == 0 {
-			if branchC16 == 1 {
-				row[isExtLongEvenC16Pos] = 1
-			} else {
-				row[isExtLongEvenC1Pos] = 1
-			}
-		} else {
-			if branchC16 == 1 {
-				row[isExtLongOddC16Pos] = 1
-			} else {
-				row[isExtLongOddC1Pos] = 1
-			}
-		}
-	}
-}
-
-func prepareEmptyExtensionRows(beforeModification, afterModification bool) [][]byte {
-	ext_row1 := make([]byte, rowLen)
-	ext_row2 := make([]byte, rowLen)
-	if !beforeModification && !afterModification {
-		ext_row1 = append(ext_row1, 16)
-		ext_row2 = append(ext_row2, 17)
-	} else if beforeModification {
-		ext_row1 = append(ext_row1, 20)
-		ext_row2 = append(ext_row2, 21)
-	} else if afterModification {
-		ext_row1 = append(ext_row1, 22)
-		ext_row2 = append(ext_row2, 23)
-	}
-
-	return [][]byte{ext_row1, ext_row2}
-}
-
-// TODO: remove when Nodes are fully implemented
-func prepareExtensionRows(extNibbles [][]byte, extensionNodeInd int, proofEl1, proofEl2 []byte, beforeModification, afterModification bool) (byte, []byte, []byte) {
-	var extensionRowS []byte
-	var extensionRowC []byte
-
-	extRows := prepareEmptyExtensionRows(beforeModification, afterModification)
-	extensionRowS = extRows[0]
-	extensionRowC = extRows[1]
-	prepareExtensionRow(extensionRowS, proofEl1, true)
-	prepareExtensionRow(extensionRowC, proofEl2, false)
-
-	evenNumberOfNibbles := proofEl1[2] == 0
-	keyLen := getExtensionNodeKeyLen(proofEl1)
-	numberOfNibbles := getExtensionNumberOfNibbles(proofEl1)
-
-	// We need nibbles as witness to compute key RLC, so we set them
-	// into extensionRowC s_advices (we can do this because both extension
-	// nodes have the same key, so we can have this info only in one).
-	// There can be more up to 64 nibbles, but there is only 32 bytes
-	// in extensionRowC s_advices. So we store every second nibble (having
-	// the whole byte and one nibble is enough to compute the other nibble).
-	startNibblePos := 2 // we don't need any nibbles for case keyLen = 1
-	if keyLen > 1 {
-		if evenNumberOfNibbles {
-			startNibblePos = 1
-		} else {
-			startNibblePos = 2
-		}
-	}
-	ind := 0
-	for j := startNibblePos; j < len(extNibbles[extensionNodeInd]); j += 2 {
-		extensionRowC[branchNodeRLPLen+ind] =
-			extNibbles[extensionNodeInd][j]
-		ind++
-	}
-
-	return numberOfNibbles, extensionRowS, extensionRowC
-}
-
 func prepareExtensions(extNibbles [][]byte, extensionNodeInd int, proofEl1, proofEl2 []byte) (byte, []byte, [][]byte) {
 	var values [][]byte
 	v1 := make([]byte, valueLen)
@@ -197,77 +105,6 @@ func getExtensionNodeNibbles(proofEl []byte) []byte {
 	return nibbles
 }
 
-// TODO: remove when Nodes are fully implemented
-func prepareExtensionRow(witnessRow, proofEl []byte, setKey bool) {
-	// storageProof[i]:
-	// [228,130,0,149,160,114,253,150,133,18,192,156,19,241,162,51,210,24,1,151,16,48,7,177,42,60,49,34,230,254,242,79,132,165,90,75,249]
-	// Note that the first element (228 in this case) can go much higher - for example, if there
-	// are 40 nibbles, this would take 20 bytes which would make the first element 248.
-
-	// If only one nibble in key:
-	// [226,16,160,172,105,12...
-	// Could also be non-hashed branch:
-	// [223,16,221,198,132,32,0,0,0,1,198,132,32,0,0,0,1,128,128,128,128,128,128,128,128,128,128,128,128,128,128,128]
-
-	// Extension node with non-hashed branch:
-	// List contains up to 55 bytes (192 + 55)
-	// [247,160,16,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,213,128,194,32,1,128,194,32,1,128,128,128,128,128,128,128,128,128,128,128,128,128]
-
-	// List contains more than 55 bytes
-	// [248,58,159,16,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,217,128,196,130,32,0,1,128,196,130,32,0,1,128,128,128,128,128,128,128,128,128,128,128,128,128]
-
-	// Note that the extension node can be much shorter than the one above - in case when
-	// there are less nibbles, so we cannot say that 226 appears as the first byte only
-	// when there are hashed nodes in the branch and there is only one nibble.
-	// Branch with two non-hashed nodes (that's the shortest possible branch):
-	// [217,128,196,130,32,0,1,128,196,130,32,0,1,128,128,128,128,128,128,128,128,128,128,128,128,128]
-	// Note: branch contains at least 26 bytes. 192 + 26 = 218
-
-	/*
-		If proofEl[0] <= 247 (length at most 55, so proofEl[1] doesn't specify the length of the whole
-			remaining stream, only of the next substream)
-		  If proofEl[1] <= 128:
-		    There is only 1 byte for nibbles (keyLen = 1) and this is proofEl[1].
-		  Else:
-		    Nibbles are stored in more than 1 byte, proofEl[1] specifies the length of bytes.
-		Else:
-		  proofEl[1] contains the length of the remaining stream.
-		  proofEl[2] specifies the length of the bytes (for storing nibbles).
-		  Note that we can't have only one nibble in this case.
-	*/
-
-	if setKey {
-		witnessRow[0] = proofEl[0]
-		witnessRow[1] = proofEl[1]
-	}
-
-	lenKey, startKey := getExtensionLenStartKey(proofEl)
-	if startKey == 3 {
-		witnessRow[2] = proofEl[2]
-	}
-
-	if setKey {
-		for j := 0; j < lenKey; j++ {
-			witnessRow[startKey+j] = proofEl[startKey+j]
-		}
-	}
-
-	encodedNodeLen := proofEl[startKey+lenKey]
-	nodeLen := byte(0)
-	start := branch2start
-	if encodedNodeLen > 192 {
-		// we have a list, that means a non-hashed node
-		nodeLen = encodedNodeLen - 192
-	} else if encodedNodeLen == 160 {
-		// hashed-node
-		nodeLen = encodedNodeLen - 128
-	}
-	witnessRow[start] = encodedNodeLen
-	for j := 0; j < int(nodeLen); j++ {
-		witnessRow[start+1+j] = proofEl[startKey+lenKey+1+j]
-	}
-}
-
 func prepareExtension(v1, v2, proofEl []byte, setKey bool) []byte {
 	// storageProof[i]:
 	// [228,130,0,149,160,114,253,150,133,18,192,156,19,241,162,51,210,24,1,151,16,48,7,177,42,60,49,34,230,254,242,79,132,165,90,75,249]