From 25428d53f194ee7486f394f42defdd6f638e56e0 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 7 Jul 2021 15:46:51 +0300
Subject: [PATCH] mpt: strip branch node in a simple implementation

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>
---
 pkg/core/mpt/trie.go      |  1 +
 pkg/core/mpt/trie_test.go | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/pkg/core/mpt/trie.go b/pkg/core/mpt/trie.go
index 6f4b452a4..fea6141e6 100644
--- a/pkg/core/mpt/trie.go
+++ b/pkg/core/mpt/trie.go
@@ -310,6 +310,7 @@ func (t *Trie) deleteFromExtension(n *ExtensionNode, path []byte) (Node, error)
 		if nxt.IsEmpty() {
 			return nxt, nil
 		}
+		n.next = nxt
 	default:
 		n.next = r
 	}
diff --git a/pkg/core/mpt/trie_test.go b/pkg/core/mpt/trie_test.go
index 18d0d43b6..25ddd561c 100644
--- a/pkg/core/mpt/trie_test.go
+++ b/pkg/core/mpt/trie_test.go
@@ -395,6 +395,24 @@ func testTrieDelete(t *testing.T, enableGC bool) {
 					tr.testHas(t, []byte{}, nil)
 					tr.testHas(t, []byte{0x56}, []byte{0x34})
 					require.IsType(t, (*ExtensionNode)(nil), tr.root)
+
+					t.Run("WithHash, branch node replaced", func(t *testing.T) {
+						ch := NewLeafNode([]byte{5, 6})
+						h := ch.Hash()
+
+						b := NewBranchNode()
+						b.Children[3] = NewExtensionNode([]byte{4}, NewLeafNode([]byte{1, 2, 3}))
+						b.Children[lastChild] = NewHashNode(h)
+
+						tr := NewTrie(NewExtensionNode([]byte{1, 2}, b), enableGC, newTestStore())
+						tr.putToStore(ch)
+
+						require.NoError(t, tr.Delete([]byte{0x12, 0x34}))
+						tr.testHas(t, []byte{0x12, 0x34}, nil)
+						tr.testHas(t, []byte{0x12}, []byte{5, 6})
+						require.IsType(t, (*ExtensionNode)(nil), tr.root)
+						require.Equal(t, h, tr.root.(*ExtensionNode).next.Hash())
+					})
 				})
 
 				t.Run("LeaveLeaf", func(t *testing.T) {