From a33e422b794438190ab4610080a8426b1597ead7 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 7 Jul 2021 14:04:29 +0300
Subject: [PATCH] mpt: do not create an Extension for the last child of a
 Branch

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>
---
 pkg/core/mpt/batch.go      |  5 ++++-
 pkg/core/mpt/batch_test.go | 16 ++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/pkg/core/mpt/batch.go b/pkg/core/mpt/batch.go
index 1b782498b..a03a6cbc8 100644
--- a/pkg/core/mpt/batch.go
+++ b/pkg/core/mpt/batch.go
@@ -192,7 +192,10 @@ func (t *Trie) stripBranch(b *BranchNode) (Node, error) {
 	case n == 0:
 		return new(HashNode), nil
 	case n == 1:
-		return t.mergeExtension([]byte{lastIndex}, b.Children[lastIndex])
+		if lastIndex != lastChild {
+			return t.mergeExtension([]byte{lastIndex}, b.Children[lastIndex])
+		}
+		return b.Children[lastIndex], nil
 	default:
 		t.addRef(b.Hash(), b.bytes)
 		return b, nil
diff --git a/pkg/core/mpt/batch_test.go b/pkg/core/mpt/batch_test.go
index ee673825a..6585d5204 100644
--- a/pkg/core/mpt/batch_test.go
+++ b/pkg/core/mpt/batch_test.go
@@ -174,6 +174,22 @@ func TestTrie_PutBatchBranch(t *testing.T) {
 			testPut(t, ps, tr1, tr2)
 			require.IsType(t, (*ExtensionNode)(nil), tr1.root)
 		})
+		t.Run("non-empty child is last node", func(t *testing.T) {
+			tr1 := NewTrie(new(HashNode), false, newTestStore())
+			tr2 := NewTrie(new(HashNode), false, newTestStore())
+			require.NoError(t, tr1.Put([]byte{0x00, 2}, []byte("value1")))
+			require.NoError(t, tr2.Put([]byte{0x00, 2}, []byte("value1")))
+			require.NoError(t, tr1.Put([]byte{0x00}, []byte("value2")))
+			require.NoError(t, tr2.Put([]byte{0x00}, []byte("value2")))
+
+			tr1.Flush()
+			tr1.Collapse(1)
+			tr2.Flush()
+			tr2.Collapse(1)
+
+			var ps = pairs{{[]byte{0x00, 2}, nil}}
+			testPut(t, ps, tr1, tr2)
+		})
 	})
 	t.Run("incomplete put, transform to extension", func(t *testing.T) {
 		tr1, tr2 := prepareBranch(t)