From 242b24af9f21cd4b35e2e609b12371f41528da3d Mon Sep 17 00:00:00 2001
From: Martin HS <martin@swende.se>
Date: Tue, 30 Apr 2024 13:51:04 +0200
Subject: [PATCH] trie/trienode: minor speedup in nodeset merging (#29683)

---
 trie/trienode/node.go      |  7 ++++-
 trie/trienode/node_test.go | 61 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 trie/trienode/node_test.go

diff --git a/trie/trienode/node.go b/trie/trienode/node.go
index 055db8822e..bc93e3ca88 100644
--- a/trie/trienode/node.go
+++ b/trie/trienode/node.go
@@ -114,7 +114,12 @@ func (set *NodeSet) Merge(owner common.Hash, nodes map[string]*Node) error {
 				set.updates -= 1
 			}
 		}
-		set.AddNode([]byte(path), node)
+		if node.IsDeleted() {
+			set.deletes += 1
+		} else {
+			set.updates += 1
+		}
+		set.Nodes[path] = node
 	}
 	return nil
 }
diff --git a/trie/trienode/node_test.go b/trie/trienode/node_test.go
new file mode 100644
index 0000000000..bcb3a2202b
--- /dev/null
+++ b/trie/trienode/node_test.go
@@ -0,0 +1,61 @@
+// Copyright 2023 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>
+
+package trienode
+
+import (
+	"crypto/rand"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
+)
+
+func BenchmarkMerge(b *testing.B) {
+	b.Run("1K", func(b *testing.B) {
+		benchmarkMerge(b, 1000)
+	})
+	b.Run("10K", func(b *testing.B) {
+		benchmarkMerge(b, 10_000)
+	})
+}
+
+func benchmarkMerge(b *testing.B, count int) {
+	x := NewNodeSet(common.Hash{})
+	y := NewNodeSet(common.Hash{})
+	addNode := func(s *NodeSet) {
+		path := make([]byte, 4)
+		rand.Read(path)
+		blob := make([]byte, 32)
+		rand.Read(blob)
+		hash := crypto.Keccak256Hash(blob)
+		s.AddNode(path, New(hash, blob))
+	}
+	for i := 0; i < count; i++ {
+		// Random path of 4 nibbles
+		addNode(x)
+		addNode(y)
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// Store set x into a backup
+		z := NewNodeSet(common.Hash{})
+		z.Merge(common.Hash{}, x.Nodes)
+		// Merge y into x
+		x.Merge(common.Hash{}, y.Nodes)
+		x = z
+	}
+}