feat(witness): generate the storage witness from the world-state scope

src/Nethermind/Nethermind.Benchmark/Store/PatriciaTrieWitnessGeneratorBenchmarks.cs

@@ -0,0 +1,142 @@
+// SPDX-FileCopyrightText: 2026 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using BenchmarkDotNet.Attributes;
+using Nethermind.Core;
+using Nethermind.Core.Collections;
+using Nethermind.Core.Crypto;
+using Nethermind.Db;
+using Nethermind.Logging;
+using Nethermind.Trie;
+using Nethermind.Trie.Pruning;
+
+namespace Nethermind.Benchmarks.Store
+{
+    /// <summary>
+    /// Compares the new mutation-free <see cref="PatriciaTrieWitnessGenerator"/> (sequential and parallel) against
+    /// the old "capture trie reads during the actual mutation" technique it replaces.
+    /// </summary>
+    [MemoryDiagnoser]
+    [MinIterationTime(1000)]
+    public class PatriciaTrieWitnessGeneratorBenchmarks
+    {
+        [Params(100_000)]
+        public int TrieSize { get; set; }
+
+        [Params(1_000, 5_000)]
+        public int TouchedCount { get; set; }
+
+        [Params(0.0, 0.5)]
+        public double DeleteFraction { get; set; }
+
+        private MemDb _db;
+        private Hash256 _root;
+        private PatriciaTrieWitnessGenerator.PathEntry[] _entries;
+        private Hash256[] _reads;
+        private Hash256[] _deletes;
+
+        [GlobalSetup]
+        public void Setup()
+        {
+            Random rng = new(0);
+
+            MemDb db = new();
+            RawScopedTrieStore store = new(db);
+            PatriciaTree tree = new(store, LimboLogs.Instance);
+
+            Hash256[] keys = new Hash256[TrieSize];
+            using ArrayPoolListRef<PatriciaTree.BulkSetEntry> bulk = new(TrieSize);
+            for (int i = 0; i < TrieSize; i++)
+            {
+                byte[] keyBuf = new byte[32];
+                rng.NextBytes(keyBuf);
+                byte[] valueBuf = new byte[32];
+                rng.NextBytes(valueBuf);
+                keys[i] = new Hash256(keyBuf);
+                bulk.Add(new PatriciaTree.BulkSetEntry(keys[i], valueBuf));
+            }
+            tree.BulkSet(bulk);
+            tree.Commit();
+
+            _db = db;
+            _root = tree.RootHash;
+
+            int deleteCount = (int)(TouchedCount * DeleteFraction);
+            _entries = new PatriciaTrieWitnessGenerator.PathEntry[TouchedCount];
+            List<Hash256> reads = [];
+            List<Hash256> deletes = [];
+            for (int i = 0; i < TouchedCount; i++)
+            {
+                Hash256 key = keys[rng.Next(TrieSize)];
+                bool isDeleted = i < deleteCount;
+                _entries[i] = new PatriciaTrieWitnessGenerator.PathEntry(
+                    key,
+                    isDeleted ? PatriciaTrieWitnessGenerator.AccessType.Delete : PatriciaTrieWitnessGenerator.AccessType.Read);
+                (isDeleted ? deletes : reads).Add(key);
+            }
+
+            _reads = [.. reads];
+            _deletes = [.. deletes];
+        }
+
+        [Benchmark(Baseline = true)]
+        public int Old_CaptureDuringMutation()
+        {
+            CapturingScopedTrieStore store = new(new RawScopedTrieStore(_db));
+            PatriciaTree tree = new(store, LimboLogs.Instance) { RootHash = _root };
+            foreach (Hash256 key in _reads) tree.Get(key.Bytes);
+            foreach (Hash256 key in _deletes) tree.Set(key.Bytes, (byte[])null);
+            tree.UpdateRootHash();
+            return store.Captured.Count;
+        }
+
+        [Benchmark]
+        public int New_Sequential()
+        {
+            CountingSink sink = new();
+            PatriciaTrieWitnessGenerator.Generate(new RawScopedTrieStore(_db), _root, _entries, sink, parallelize: false);
+            return sink.Count;
+        }
+
+        [Benchmark]
+        public int New_Parallel()
+        {
+            CountingSink sink = new();
+            PatriciaTrieWitnessGenerator.Generate(new RawScopedTrieStore(_db), _root, _entries, sink, parallelize: true);
+            return sink.Count;
+        }
+
+        private sealed class CountingSink : PatriciaTrieWitnessGenerator.ISink
+        {
+            private int _count;
+            public int Count => _count;
+            public void Add(in TreePath path, TrieNode node) => Interlocked.Increment(ref _count);
+        }
+
+        private sealed class CapturingScopedTrieStore(IScopedTrieStore baseStore) : IScopedTrieStore
+        {
+            public Dictionary<Hash256AsKey, byte[]> Captured { get; } = [];
+
+            public TrieNode FindCachedOrUnknown(in TreePath path, Hash256 hash) => baseStore.FindCachedOrUnknown(in path, hash);
+
+            public byte[] LoadRlp(in TreePath path, Hash256 hash, ReadFlags flags = ReadFlags.None) => Capture(hash, baseStore.LoadRlp(in path, hash, flags));
+
+            public byte[] TryLoadRlp(in TreePath path, Hash256 hash, ReadFlags flags = ReadFlags.None) => Capture(hash, baseStore.TryLoadRlp(in path, hash, flags));
+
+            private byte[] Capture(Hash256 hash, byte[] rlp)
+            {
+                if (rlp is not null) Captured[hash] = rlp;
+                return rlp;
+            }
+
+            public ITrieNodeResolver GetStorageTrieNodeResolver(Hash256 address) => baseStore.GetStorageTrieNodeResolver(address);
+
+            public INodeStorage.KeyScheme Scheme => baseStore.Scheme;
+
+            public ICommitter BeginCommit(TrieNode root, WriteFlags writeFlags = WriteFlags.None) => baseStore.BeginCommit(root, writeFlags);
+        }
+    }
+}

src/Nethermind/Nethermind.Consensus/Stateless/SingleCallWitnessCollector.cs

@@ -2,12 +2,15 @@
 // SPDX-License-Identifier: LGPL-3.0-only
 
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using Nethermind.Blockchain.Tracing;
 using Nethermind.Core;
 using Nethermind.Evm;
+using Nethermind.Evm.State;
 using Nethermind.Evm.Tracing;
 using Nethermind.Evm.TransactionProcessing;
+using Nethermind.State;
 
 namespace Nethermind.Consensus.Stateless;
 
@@ -41,15 +44,17 @@ public interface ISingleCallWitnessCollector
 }
 
 public class SingleCallWitnessCollector(
-    WitnessGeneratingWorldState worldState,
+    IWorldState worldState,
+    AccessWitnessScopeProvider accessWitness,
+    WitnessGeneratingHeaderFinder headerFinder,
     ITransactionProcessor transactionProcessor) : ISingleCallWitnessCollector
 {
     public SingleCallWitnessResult ExecuteCallAndCollectWitness(BlockHeader blockHeader, Transaction transaction, CancellationToken cancellationToken = default)
     {
         // Uses blockHeader (not parentHeader) intentionally: for a single call we want the
         // post-state of the target block. Block-level witness uses parentHeader because it
         // needs the pre-state to re-execute the block's transactions.
-        using IDisposable? scope = worldState.BeginScope(blockHeader);
+        using IDisposable? scope = worldState.BeginScope(blockHeader, trackWitness: true);
 
         // Mirror BlockchainBridge.CallAndRestore: ignore the caller-supplied nonce and resolve it
         // from the scoped state. Without this, a proof_call request that includes `from` but omits
@@ -63,11 +68,12 @@ public SingleCallWitnessResult ExecuteCallAndCollectWitness(BlockHeader blockHea
         CallOutputTracer tracer = new();
         TransactionResult txResult = transactionProcessor.CallAndRestore(transaction, blockHeader, tracer.WithCancellation(cancellationToken));
 
+        IReadOnlyList<byte[]> stateNodes = worldState.Witness ?? throw new InvalidOperationException("Witness tracking was not enabled for this scope.");
         return new SingleCallWitnessResult(
             Output: tracer.ReturnValue,
             Error: txResult.GetErrorMessage(tracer.Error),
             ExecutionReverted: txResult.EvmExceptionType == EvmExceptionType.Revert,
             InputError: !txResult.TransactionExecuted,
-            Witness: worldState.GetWitness(blockHeader));
+            Witness: WitnessAssembler.Build(stateNodes, accessWitness.TouchedKeys, accessWitness.Codes, headerFinder, blockHeader));
     }
 }

src/Nethermind/Nethermind.Consensus/Stateless/WitnessAssembler.cs

@@ -0,0 +1,71 @@
+// SPDX-FileCopyrightText: 2026 Demerzel Solutions Limited
+// SPDX-License-Identifier: LGPL-3.0-only
+
+using System.Collections.Generic;
+using Nethermind.Core;
+using Nethermind.Core.Collections;
+using Nethermind.Core.Extensions;
+using Nethermind.Int256;
+
+namespace Nethermind.Consensus.Stateless;
+
+/// <summary>
+/// Assembles a full execution <see cref="Witness"/> from its independently produced parts: the storage
+/// witness (state-trie node RLPs), the touched keys and contract code, and the ancestor block headers.
+/// </summary>
+internal static class WitnessAssembler
+{
+    public static Witness Build(
+        IReadOnlyList<byte[]> stateNodes,
+        IReadOnlyDictionary<AddressAsKey, HashSet<UInt256>> touchedKeys,
+        IReadOnlyCollection<byte[]> codes,
+        WitnessGeneratingHeaderFinder headerFinder,
+        BlockHeader parentHeader)
+    {
+        // New pool-rented buffers added here must also be disposed in the catch below.
+        ArrayPoolList<byte[]>? codeList = null;
+        ArrayPoolList<byte[]>? state = null;
+        ArrayPoolList<byte[]>? keys = null;
+        try
+        {
+            codeList = new ArrayPoolList<byte[]>(codes.Count);
+            foreach (byte[] code in codes)
+                codeList.Add(code);
+
+            state = new ArrayPoolList<byte[]>(stateNodes.Count);
+            foreach (byte[] node in stateNodes)
+                state.Add(node);
+
+            int totalKeysCount = touchedKeys.Count;
+            foreach (KeyValuePair<AddressAsKey, HashSet<UInt256>> kvp in touchedKeys)
+                totalKeysCount += kvp.Value.Count;
+
+            keys = new ArrayPoolList<byte[]>(totalKeysCount);
+            // Keys ordered like: <addr1><addr2><slot1-of-addr2><slot2-of-addr2><addr3><slot1-of-addr3>
+            foreach (KeyValuePair<AddressAsKey, HashSet<UInt256>> kvp in touchedKeys)
+            {
+                keys.Add(kvp.Key.Value.Bytes.ToArray());
+                foreach (UInt256 slot in kvp.Value)
+                    keys.Add(slot.ToBigEndian());
+            }
+
+            return new Witness
+            {
+                Codes = codeList,
+                State = state,
+                Keys = keys,
+                Headers = headerFinder.GetWitnessHeaders(parentHeader.Hash!)
+            };
+        }
+        catch
+        {
+            // Any failure mid-build returns the rented buffers before propagating, else they leak:
+            // an OOM while filling a list, or GetWitnessHeaders throwing because a walked ancestor
+            // header vanished (reorg/prune between the call and the witness build).
+            codeList?.Dispose();
+            state?.Dispose();
+            keys?.Dispose();
+            throw;
+        }
+    }
+}