perf: encode a branch node in one pass

src/Nethermind/Nethermind.Trie.Test/TrieNodeTests.cs

@@ -719,6 +719,32 @@ public void Can_encode_branch_with_unresolved_children()
         restoredNode.RlpEncode(NullTrieNodeResolver.Instance, ref emptyPath);
     }
 
+    [Test]
+    public void Can_encode_branch_with_every_child_a_hash()
+    {
+        TrieNode node = new(NodeType.Branch);
+        for (int i = 0; i < TrieNode.BranchesCount; i++)
+        {
+            node.SetChild(i, new TrieNode(NodeType.Unknown, Keccak.Compute([(byte)i])));
+        }
+
+        TreePath emptyPath = TreePath.Empty;
+        CappedArray<byte> rlp = node.RlpEncode(NullTrieNodeResolver.Instance, ref emptyPath);
+
+        TrieNode restoredNode = new(NodeType.Unknown, rlp);
+        restoredNode.ResolveNode(NullTrieNodeResolver.Instance, TreePath.Empty);
+
+        using (Assert.EnterMultipleScope())
+        {
+            // The widest a branch encodes to: sixteen 33-byte hash items plus the value and the header.
+            Assert.That(rlp.Length, Is.EqualTo(532), "RLP length");
+            for (int i = 0; i < TrieNode.BranchesCount; i++)
+            {
+                Assert.That(restoredNode.GetChildHash(i), Is.EqualTo(Keccak.Compute([(byte)i])), $"child {i}");
+            }
+        }
+    }
+
     [Test]
     public void Size_of_a_heavy_leaf_is_correct()
     {

src/Nethermind/Nethermind.Trie/TrieNode.Decoder.cs

@@ -142,19 +142,61 @@ public static CappedArray<byte> EncodeLeaf(TrieNode node, ICappedArrayPool? pool
             [DoesNotReturn, StackTraceHidden]
             private static void ThrowNullKey(TrieNode node) => throw new TrieException($"Hex prefix of a leaf node is null at node {node.Keccak}");
 
+            /// <summary>Scratch for one branch's children, before the sequence header is known.</summary>
+            /// <remarks>Sixteen children at most, each either a 33-byte hash item or a node small enough
+            /// to be embedded, which the trie only does below 32 bytes. A struct local rather than a
+            /// stackalloc, for the reason given on <c>KeccakHash</c>'s state buffer: localloc would pin
+            /// the method at Tier0-FullOpts and add stack-probe overhead per call. Writes go through a
+            /// span, so an over-long branch throws rather than running off the buffer.</remarks>
+            [InlineArray(BranchesCount * Rlp.LengthOfKeccakRlp)]
+            private struct BranchScratch
+            {
+                private byte _element0;
+            }
+
+            [SkipLocalsInit]
             public static CappedArray<byte> RlpEncodeBranch(TrieNode item, ITrieNodeResolver tree, ref TreePath path, ICappedArrayPool? pool, bool canBeParallel)
             {
                 Metrics.IncrementTreeNodeRlpEncodings();
 
                 const int valueRlpLength = 1;
-                int contentLength = valueRlpLength + (UseParallel(canBeParallel, item) ? GetChildrenRlpLengthForBranchParallel(tree, ref path, item, pool, canBeParallel) : GetChildrenRlpLengthForBranch(tree, ref path, item, pool, canBeParallel));
-                int sequenceLength = Rlp.LengthOfSequence(contentLength);
-                CappedArray<byte> result = pool.SafeRent(sequenceLength);
-                Span<byte> resultSpan = result.AsSpan();
-                int position = Rlp.StartSequence(resultSpan, 0, contentLength);
-                WriteChildrenRlpBranch(tree, ref path, item, resultSpan.Slice(position, contentLength - valueRlpLength), pool, canBeParallel);
-                position = sequenceLength - valueRlpLength;
-                resultSpan[position] = 128;
+                int contentLength;
+                int sequenceLength;
+                CappedArray<byte> result;
+                Span<byte> resultSpan;
+                int position;
+
+                // The sequence header carries the children's length and CappedArray has no offset to write
+                // it backwards into, so the length has to be known before the children can go in. Writing
+                // them into a scratch buffer and copying them in behind the header trades the measuring
+                // walk for one bounded copy. The walk is kept where it does something else as well:
+                // spreading the children over cores, or collecting branch pairs for batched hashing.
+                bool useParallel = UseParallel(canBeParallel, item);
+                if (useParallel || (Avx512F.VL.IsSupported && HasBatchableChildPair(item)))
+                {
+                    contentLength = valueRlpLength + (useParallel
+                        ? GetChildrenRlpLengthForBranchParallel(tree, ref path, item, pool, canBeParallel)
+                        : GetChildrenRlpLengthForBranch(tree, ref path, item, pool, canBeParallel));
+                    sequenceLength = Rlp.LengthOfSequence(contentLength);
+                    result = pool.SafeRent(sequenceLength);
+                    resultSpan = result.AsSpan();
+                    position = Rlp.StartSequence(resultSpan, 0, contentLength);
+                    WriteChildrenRlpBranch(tree, ref path, item, resultSpan.Slice(position, contentLength - valueRlpLength), pool, canBeParallel);
+                    resultSpan[sequenceLength - valueRlpLength] = 128;
+
+                    return result;
+                }
+
+                Unsafe.SkipInit(out BranchScratch scratch);
+                Span<byte> children = scratch;
+                int childrenLength = WriteChildrenRlpBranch(tree, ref path, item, children, pool, canBeParallel);
+                contentLength = valueRlpLength + childrenLength;
+                sequenceLength = Rlp.LengthOfSequence(contentLength);
+                result = pool.SafeRent(sequenceLength);
+                resultSpan = result.AsSpan();
+                position = Rlp.StartSequence(resultSpan, 0, contentLength);
+                children[..childrenLength].CopyTo(resultSpan[position..]);
+                resultSpan[sequenceLength - valueRlpLength] = 128;
 
                 return result;
 
@@ -180,6 +222,27 @@ static bool UseParallel(bool canBeParallel, TrieNode item)
                 }
             }
 
+            /// <summary>Whether the measuring walk could pair up child hashes for <see cref="HashPreparedBranches" />.</summary>
+            /// <remarks>Only a dirty branch child is a candidate, and a lone candidate is hashed on its own, so
+            /// a branch without two of them gains nothing from the walk. The walk narrows the set further — a
+            /// candidate whose RLP is not a full branch drops out — so an upper bound is all this needs to be.</remarks>
+            private static bool HasBatchableChildPair(TrieNode item)
+            {
+                const int MinChildrenForBatchedHashing = 2;
+                int candidates = 0;
+                Debug.Assert(item._nodeData is BranchData, "Data is not BranchData");
+                BranchData branchData = Unsafe.As<BranchData>(item._nodeData!);
+                for (int i = 0; i < BranchesCount; i++)
+                {
+                    if (branchData[i] is TrieNode { IsBranch: true, Keccak: null } && ++candidates >= MinChildrenForBatchedHashing)
+                    {
+                        return true;
+                    }
+                }
+
+                return false;
+            }
+
             private static void HashPreparedBranches(TrieNode item, ushort candidateMask)
             {
                 int firstIndex = BitOperations.TrailingZeroCount(candidateMask);
@@ -444,20 +507,17 @@ private static int GetChildrenRlpLengthForBranchRlp(ITrieNodeResolver tree, ref
                 return totalLength;
             }
 
-            private static void WriteChildrenRlpBranch(ITrieNodeResolver tree, ref TreePath path, TrieNode item, Span<byte> destination, ICappedArrayPool? bufferPool, bool canBeParallel)
-            {
+            /// <summary>Writes a branch's sixteen children into <paramref name="destination" />, each as a hash
+            /// item or an embedded node.</summary>
+            /// <returns>The number of bytes written.</returns>
+            private static int WriteChildrenRlpBranch(ITrieNodeResolver tree, ref TreePath path, TrieNode item, Span<byte> destination, ICappedArrayPool? bufferPool, bool canBeParallel) =>
                 // Tail call optimized.
-                if (item.HasRlp)
-                {
-                    WriteChildrenRlpBranchRlp(tree, ref path, item, destination, bufferPool, canBeParallel);
-                }
-                else
-                {
-                    WriteChildrenRlpBranchNonRlp(tree, ref path, item, destination, bufferPool, canBeParallel);
-                }
-            }
+                item.HasRlp
+                    ? WriteChildrenRlpBranchRlp(tree, ref path, item, destination, bufferPool, canBeParallel)
+                    : WriteChildrenRlpBranchNonRlp(tree, ref path, item, destination, bufferPool, canBeParallel);
 
-            private static void WriteChildrenRlpBranchNonRlp(ITrieNodeResolver tree, ref TreePath path, TrieNode item, Span<byte> destination, ICappedArrayPool? bufferPool, bool canBeParallel)
+            /// <inheritdoc cref="WriteChildrenRlpBranch" />
+            private static int WriteChildrenRlpBranchNonRlp(ITrieNodeResolver tree, ref TreePath path, TrieNode item, Span<byte> destination, ICappedArrayPool? bufferPool, bool canBeParallel)
             {
                 int position = 0;
                 for (int i = 0; i < BranchesCount; i++)
@@ -492,9 +552,12 @@ private static void WriteChildrenRlpBranchNonRlp(ITrieNodeResolver tree, ref Tre
                         }
                     }
                 }
+
+                return position;
             }
 
-            private static void WriteChildrenRlpBranchRlp(ITrieNodeResolver tree, ref TreePath path, TrieNode item, Span<byte> destination, ICappedArrayPool? bufferPool, bool canBeParallel)
+            /// <inheritdoc cref="WriteChildrenRlpBranch" />
+            private static int WriteChildrenRlpBranchRlp(ITrieNodeResolver tree, ref TreePath path, TrieNode item, Span<byte> destination, ICappedArrayPool? bufferPool, bool canBeParallel)
             {
                 RlpReader rlpReader = item.RlpReader;
                 item.SeekChild(ref rlpReader, 0);
@@ -562,7 +625,10 @@ private static void WriteChildrenRlpBranchRlp(ITrieNodeResolver tree, ref TreePa
                 if (runStart >= 0)
                 {
                     rlpReader.Data.Slice(runStart, runLength).CopyTo(destination.Slice(position, runLength));
+                    position += runLength;
                 }
+
+                return position;
             }
         }
     }