From 484cb7a3c60b83c97af101c34d26ecd58c1fa6d2 Mon Sep 17 00:00:00 2001
From: Miha Stopar <stopar.miha@gmail.com>
Date: Mon, 5 Feb 2024 11:19:34 +0100
Subject: [PATCH] getDriftedPosition fix; test added

---
 mpt-witness-generator/witness/branch.go       | 34 ++++++-----
 .../gen_witness_from_local_blockchain_test.go | 56 +++++++++++++++++++
 2 files changed, 76 insertions(+), 14 deletions(-)

diff --git a/mpt-witness-generator/witness/branch.go b/mpt-witness-generator/witness/branch.go
index 9cbe21ea5d1..082d56f8e94 100644
--- a/mpt-witness-generator/witness/branch.go
+++ b/mpt-witness-generator/witness/branch.go
@@ -154,21 +154,27 @@ func prepareBranchNode(branch1, branch2, extNode1, extNode2, extListRlpBytes []b
 func getDriftedPosition(leafKeyRow []byte, numberOfNibbles int) byte {
 	var nibbles []byte
 	if leafKeyRow[0] != 248 {
-		keyLen := int(leafKeyRow[1] - 128)
-		if (leafKeyRow[2] != 32) && (leafKeyRow[2] != 0) { // second term is for extension node
-			if leafKeyRow[2] < 32 { // extension node
-				nibbles = append(nibbles, leafKeyRow[2]-16)
-			} else { // leaf
-				nibbles = append(nibbles, leafKeyRow[2]-48)
+		var keyLen int
+		if leafKeyRow[1] > 128 {
+			keyLen = int(leafKeyRow[1] - 128)
+			if (leafKeyRow[2] != 32) && (leafKeyRow[2] != 0) { // second term is for extension node
+				if leafKeyRow[2] < 32 { // extension node
+					nibbles = append(nibbles, leafKeyRow[2]-16)
+				} else { // leaf
+					nibbles = append(nibbles, leafKeyRow[2]-48)
+				}
 			}
-		}
-		for i := 0; i < keyLen-1; i++ { // -1 because the first byte doesn't have any nibbles
-			b := leafKeyRow[3+i]
-			n1 := b / 16
-			n2 := b - n1*16
-			nibbles = append(nibbles, n1)
-			nibbles = append(nibbles, n2)
-		}
+			for i := 0; i < keyLen-1; i++ { // -1 because the first byte doesn't have any nibbles
+				b := leafKeyRow[3+i]
+				n1 := b / 16
+				n2 := b - n1*16
+				nibbles = append(nibbles, n1)
+				nibbles = append(nibbles, n2)
+			}
+		} else {
+			keyLen = 1
+			nibbles = append(nibbles, leafKeyRow[1]-16)
+		}	
 	} else {
 		keyLen := int(leafKeyRow[2] - 128)
 		if (leafKeyRow[3] != 32) && (leafKeyRow[3] != 0) { // second term is for extension node
diff --git a/mpt-witness-generator/witness/gen_witness_from_local_blockchain_test.go b/mpt-witness-generator/witness/gen_witness_from_local_blockchain_test.go
index 4508e8e22e6..958a5a0a5de 100644
--- a/mpt-witness-generator/witness/gen_witness_from_local_blockchain_test.go
+++ b/mpt-witness-generator/witness/gen_witness_from_local_blockchain_test.go
@@ -920,3 +920,59 @@ func TestExtNodeDeletedExtShortIsBranchFirstLevel(t *testing.T) {
 
 	ExtNodeDeleted(key1, key2, key3, "ExtNodeDeletedExtShortIsBranchFirstLevel")
 }
+
+func TestExtensionIntoBranch(t *testing.T) {
+	SkipIfNoGeth(t)
+	oracle.NodeUrl = oracle.LocalUrl
+
+	blockNum := 0
+	blockNumberParent := big.NewInt(int64(blockNum))
+	blockHeaderParent := oracle.PrefetchBlock(blockNumberParent, true, nil)
+	database := state.NewDatabase(blockHeaderParent)
+	statedb, _ := state.New(blockHeaderParent.Root, database, nil)
+	addr := common.HexToAddress("0x50efbf12580138bc623c95757286df4e24eb81c9")
+	
+	statedb.DisableLoadingRemoteAccounts()
+
+	statedb.CreateAccount(addr)
+
+	oracle.PreventHashingInSecureTrie = true // to store the unchanged key
+
+	val0 := common.BigToHash(big.NewInt(int64(1)))
+	key0 := common.HexToHash("0x1200000000000000000000000000000000000000000000000000000000000000")
+	statedb.SetState(addr, key0, val0)
+
+	key00 := common.HexToHash("0x0000000000000000000000000000000000000000000000000000000000000000")
+	statedb.SetState(addr, key00, val0)
+	// After insertion of key0 and key00, we have a branch B with two leaves - at position 1 and 0.
+
+	key1 := common.HexToHash("0x1234561000000000000000000000000000000000000000000000000000000000")
+
+	// make the value long to have a hashed branch
+	v1 := common.FromHex("0xbbefaa12580138bc263c95757826df4e24eb81c9aaaaaaaaaaaaaaaaaaaaaaaa")
+	val1 := common.BytesToHash(v1)
+	statedb.SetState(addr, key1, val1)
+	// After insertion of key1, we have a branch B with a leaf at position 0
+	// and an extension node E1 at position 1 (with one nibble: 2)
+	// The underlying branch B1 of E1 has two leaves - at position 0 and 3.
+
+	key2 := common.HexToHash("0x1434563000000000000000000000000000000000000000000000000000000000")
+	// After inserting key2, we have a branch B with two nodes - a leaf at position 0 and branch B1
+	// at position 1. At position 1 we have a branch B2 at position 2 (used to be E1's nibble)
+	// and a leaf at position 4 (newly added leaf).
+	// Branch B2 has two leaves - at position 0 and 3.
+
+	v1 = common.FromHex("0xbb")
+	val := common.BytesToHash(v1)
+	trieMod := TrieModification{
+		Type:    StorageChanged,
+		Key:     key2,
+		Value:   val,
+		Address: addr,
+	}
+	trieModifications := []TrieModification{trieMod}
+
+	prepareWitness("ExtensionIntoBranch", trieModifications, statedb)
+
+	oracle.PreventHashingInSecureTrie = false
+}