Fix batched trie visitor missed storage

src/Nethermind/Nethermind.Trie.Test/VisitingTests.cs

@@ -23,11 +23,11 @@ namespace Nethermind.Trie.Test;
 public class VisitingTests
 {
     [TestCaseSource(nameof(GetOptions))]
-    public void Visitors_state(VisitingOptions options)
+    public void Visitors_state(VisitingOptions options, INodeStorage.KeyScheme scheme)
     {
         MemDb memDb = new();
 
-        using TrieStore trieStore = TestTrieStoreFactory.Build(memDb, Prune.WhenCacheReaches(1.MB()), Persist.EveryBlock, LimboLogs.Instance);
+        using ITrieStore trieStore = TestTrieStoreFactory.Build(new NodeStorage(memDb, scheme), LimboLogs.Instance);
         PatriciaTree patriciaTree = new(trieStore, LimboLogs.Instance);
 
         Span<byte> raw = stackalloc byte[32];
@@ -62,11 +62,11 @@ public void Visitors_state(VisitingOptions options)
     }
 
     [TestCaseSource(nameof(GetOptions))]
-    public void Visitors_storage(VisitingOptions options)
+    public void Visitors_storage(VisitingOptions options, INodeStorage.KeyScheme scheme)
     {
         MemDb memDb = new();
 
-        using TrieStore trieStore = TestTrieStoreFactory.Build(memDb, Prune.WhenCacheReaches(1.MB()), Persist.EveryBlock, LimboLogs.Instance);
+        using ITrieStore trieStore = TestTrieStoreFactory.Build(new NodeStorage(memDb, scheme), LimboLogs.Instance);
 
         byte[] value = Enumerable.Range(1, 32).Select(static i => (byte)i).ToArray();
         Hash256 stateRootHash = Keccak.Zero;
@@ -109,8 +109,11 @@ public void Visitors_storage(VisitingOptions options)
 
         stateTree.Accept(visitor, stateTree.RootHash, options);
 
+        int totalPath = 0;
+
         foreach (var path in visitor.LeafPaths)
         {
+            totalPath++;
             if (path.Length == 64)
             {
                 AssertPath(path);
@@ -127,6 +130,8 @@ public void Visitors_storage(VisitingOptions options)
             }
         }
 
+        totalPath.Should().Be(4160);
+
         return;
 
         static void AssertPath(ReadOnlySpan<byte> path)
@@ -142,13 +147,23 @@ private static IEnumerable<TestCaseData> GetOptions()
     {
         yield return new TestCaseData(new VisitingOptions
         {
-        }).SetName("Default");
+        }, INodeStorage.KeyScheme.HalfPath).SetName("Default");
+
+        yield return new TestCaseData(new VisitingOptions
+        {
+        }, INodeStorage.KeyScheme.Hash).SetName("Default Hash");
+
+        yield return new TestCaseData(new VisitingOptions
+        {
+            MaxDegreeOfParallelism = Environment.ProcessorCount,
+            FullScanMemoryBudget = 1.MiB(),
+        }, INodeStorage.KeyScheme.HalfPath).SetName("Parallel");
 
         yield return new TestCaseData(new VisitingOptions
         {
             MaxDegreeOfParallelism = Environment.ProcessorCount,
             FullScanMemoryBudget = 1.MiB(),
-        }).SetName("Parallel");
+        }, INodeStorage.KeyScheme.Hash).SetName("Parallel Hash");
     }
 
     public class AppendingVisitor(bool expectAccount) : ITreeVisitor<AppendingVisitor.PathGatheringContext>

src/Nethermind/Nethermind.Trie/BatchedTrieVisitor.cs

@@ -367,10 +367,10 @@ private void BatchedThread()
 
         return;
 
-        static void ThrowUnableToResolve(in SmallTrieVisitContext ctx)
+        void ThrowUnableToResolve(in SmallTrieVisitContext ctx)
         {
             throw new TrieException(
-                $"Unable to resolve node without Keccak. ctx: {ctx.Level}, {ctx.ExpectAccounts}, {ctx.IsStorage}");
+                $"Unable to resolve node without Keccak. ctx: {ctx.Level}, {_visitor.ExpectAccounts}, {ctx.IsStorage}");
         }
     }
 
@@ -434,7 +434,7 @@ internal void AcceptResolvedNode(TrieNode node, in TNodeContext nodeContext, ITr
                 {
                     _visitor.VisitLeaf(nodeContext, node);
 
-                    if (!trieVisitContext.IsStorage && trieVisitContext.ExpectAccounts) // can combine these conditions
+                    if (!trieVisitContext.IsStorage && _visitor.ExpectAccounts) // can combine these conditions
                     {
                         TNodeContext childContext = nodeContext.Add(node.Key!);
 

src/Nethermind/Nethermind.Trie/VisitContext.cs

@@ -62,7 +62,6 @@ public SmallTrieVisitContext(TrieVisitContext trieVisitContext)
         private byte _flags = 0;
 
         private const byte StorageFlag = 1;
-        private const byte ExpectAccountsFlag = 2;
 
         public bool IsStorage
         {
@@ -79,21 +78,5 @@ internal set
                 }
             }
         }
-
-        public bool ExpectAccounts
-        {
-            readonly get => (_flags & ExpectAccountsFlag) == ExpectAccountsFlag;
-            internal set
-            {
-                if (value)
-                {
-                    _flags = (byte)(_flags | ExpectAccountsFlag);
-                }
-                else
-                {
-                    _flags = (byte)(_flags & ~ExpectAccountsFlag);
-                }
-            }
-        }
     }
 }