Merge pull request #3774 from greymistcube/feature/trie-sub-traversal

✨ Added methods to traverse subtries of a `MerkleTrie`

CHANGES.md

@@ -12,6 +12,8 @@ To be released.
 
  -  (Libplanet.Types) Updated `BlockMetadata.CurrentProtocolVersion`
     from 6 to 7.  [[#3769]]
+ -  (Libplanet.Store) Added `IterateSubTrieValues(KeyBytes)` and
+    `IterateSubTrieNodes(KeyBytes)` methods to `MerkleTrie`.  [[#3774]]
  -  (Libplanet.Types) Added `BlockMetadata.CurrencyAccountProtocolVersion`.
     [[#3775]]
 
@@ -30,6 +32,7 @@ To be released.
 ### CLI tools
 
 [#3769]: https://github.com/planetarium/libplanet/pull/3769
+[#3774]: https://github.com/planetarium/libplanet/pull/3774
 [#3775]: https://github.com/planetarium/libplanet/pull/3775
 
 

Libplanet.Store/Trie/MerkleTrie.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Linq;
 using System.Security.Cryptography;
 using Bencodex;
@@ -129,51 +130,63 @@ public ITrie Remove(in KeyBytes key)
         }
 
         /// <inheritdoc cref="ITrie.IterateNodes"/>
-        public IEnumerable<(Nibbles Path, INode Node)> IterateNodes()
+        public IEnumerable<(Nibbles Path, INode Node)> IterateNodes() =>
+            IterateSubTrieNodes(new KeyBytes(ImmutableArray<byte>.Empty));
+
+        /// <summary>
+        /// Iterates all values that can be reached from a sub-<see cref="ITrie"/> with
+        /// <see cref="INode"/> at <paramref name="rootPath"/> as its root.
+        /// </summary>
+        /// <param name="rootPath">The <see cref="KeyBytes"/> path of an <see cref="INode"/>
+        /// to use as the root of traversal.</param>
+        /// <returns>All <see cref="IValue"/>s together with its <em>full</em>
+        /// <see cref="KeyBytes"/> paths.</returns>
+        /// <remarks>
+        /// This requires an <see cref="INode"/> to exist at <paramref name="rootPath"/>.
+        /// As such, this does not necessarily return all <see cref="INode"/> with paths starting
+        /// with <paramref name="rootPath"/>.  In particular, if there doesn't exist
+        /// an <see cref="INode"/> at <paramref name="rootPath"/>, this returns nothing.
+        /// </remarks>
+        public IEnumerable<(KeyBytes Path, IValue Value)> IterateSubTrieValues(KeyBytes rootPath)
         {
-            if (Root is null)
+            foreach ((var path, var node) in IterateSubTrieNodes(rootPath))
             {
-                yield break;
+                if (node is ValueNode valueNode)
+                {
+                    yield return (path.ToKeyBytes(), valueNode.Value);
+                }
             }
+        }
 
-            var stack = new Stack<(Nibbles, INode)>();
-            stack.Push((Nibbles.Empty, Root));
-
-            while (stack.Count > 0)
+        /// <summary>
+        /// Iterates all sub-nodes from a sub-<see cref="ITrie"/> with <see cref="INode"/> at
+        /// <paramref name="rootPath"/> as its root.
+        /// </summary>
+        /// <param name="rootPath">The <see cref="KeyBytes"/> path of an <see cref="INode"/>
+        /// to use as the root of traversal.</param>
+        /// <returns>All <see cref="INode"/>s together with its <em>full</em>
+        /// <see cref="Nibbles"/> paths.</returns>
+        /// <remarks>
+        /// This requires an <see cref="INode"/> to exist at <paramref name="rootPath"/>.
+        /// As such, this does not necessarily return all <see cref="INode"/> with paths starting
+        /// with <paramref name="rootPath"/>.  In particular, if there doesn't exist
+        /// an <see cref="INode"/> at <paramref name="rootPath"/>, this returns nothing.
+        /// </remarks>
+        public IEnumerable<(Nibbles Path, INode Node)> IterateSubTrieNodes(KeyBytes rootPath)
+        {
+            Nibbles rootPrefix = Nibbles.FromKeyBytes(rootPath);
+            INode? root = GetNode(rootPrefix);
+            if (root is { } node)
             {
-                (Nibbles path, INode node) = stack.Pop();
-                yield return (path, node);
-                switch (node)
+                foreach (var pair in IterateNodes(rootPrefix, node))
                 {
-                    case FullNode fullNode:
-                        foreach (int index in Enumerable.Range(0, FullNode.ChildrenCount - 1))
-                        {
-                            if (fullNode.Children[index] is { } childNode)
-                            {
-                                stack.Push((path.Add((byte)index), childNode));
-                            }
-                        }
-
-                        if (fullNode.Value is { } fullNodeValue)
-                        {
-                            stack.Push((path, fullNodeValue));
-                        }
-
-                        break;
-
-                    case ShortNode shortNode:
-                        if (shortNode.Value is { } shortNodeValue)
-                        {
-                            stack.Push((path.AddRange(shortNode.Key), shortNodeValue));
-                        }
-
-                        break;
-
-                    case HashNode hashNode:
-                        stack.Push((path, UnhashNode(hashNode)));
-                        break;
+                    yield return pair;
                 }
             }
+            else
+            {
+                yield break;
+            }
         }
 
         /// <summary>
@@ -248,6 +261,51 @@ public ITrie Remove(in KeyBytes key)
             }
         }
 
+        private IEnumerable<(Nibbles Path, INode Node)> IterateNodes(
+            Nibbles rootPrefix,
+            INode root)
+        {
+            var stack = new Stack<(Nibbles, INode)>();
+            stack.Push((rootPrefix, root));
+
+            while (stack.Count > 0)
+            {
+                (Nibbles path, INode node) = stack.Pop();
+                yield return (path, node);
+                switch (node)
+                {
+                    case FullNode fullNode:
+                        foreach (int index in Enumerable.Range(0, FullNode.ChildrenCount - 1))
+                        {
+                            INode? child = fullNode.Children[index];
+                            if (!(child is null))
+                            {
+                                stack.Push((path.Add((byte)index), child));
+                            }
+                        }
+
+                        if (!(fullNode.Value is null))
+                        {
+                            stack.Push((path, fullNode.Value));
+                        }
+
+                        break;
+
+                    case ShortNode shortNode:
+                        if (!(shortNode.Value is null))
+                        {
+                            stack.Push((path.AddRange(shortNode.Key), shortNode.Value));
+                        }
+
+                        break;
+
+                    case HashNode hashNode:
+                        stack.Push((path, UnhashNode(hashNode)));
+                        break;
+                }
+            }
+        }
+
         /// <summary>
         /// Gets the concrete inner node corresponding to <paramref name="hashNode"/> from storage.
         /// </summary>

Libplanet.Tests/Store/Trie/MerkleTrieTest.cs

@@ -91,6 +91,51 @@ public void IterateNodes()
             Assert.Equal(4, merkleTrie.IterateNodes().Count());
         }
 
+        [Theory]
+        [InlineData(true, "_")]
+        [InlineData(false, "_")]
+        [InlineData(true, "_1ab3_639e")]
+        [InlineData(false, "_1ab3_639e")]
+        public void IterateSubTrie(bool commit, string extraKey)
+        {
+            IStateStore stateStore = new TrieStateStore(new MemoryKeyValueStore());
+            ITrie trie = stateStore.GetStateRoot(null);
+
+            string prefix = "_";
+            string[] keys =
+            {
+                "1b418c98",
+                "__3b8a",
+                "___",
+            };
+
+            foreach (var key in keys)
+            {
+                trie = trie.Set(new KeyBytes(Encoding.ASCII.GetBytes(key)), new Text(key));
+            }
+
+            trie = commit ? stateStore.Commit(trie) : trie;
+            Assert.Empty(
+                ((MerkleTrie)trie)
+                    .IterateSubTrieNodes(new KeyBytes(Encoding.ASCII.GetBytes(prefix))));
+            Assert.Empty(
+                ((MerkleTrie)trie)
+                    .IterateSubTrieValues(new KeyBytes(Encoding.ASCII.GetBytes(prefix))));
+
+            trie = trie.Set(new KeyBytes(Encoding.ASCII.GetBytes(extraKey)), new Text(extraKey));
+            trie = commit ? stateStore.Commit(trie) : trie;
+            Assert.Equal(
+                3,
+                ((MerkleTrie)trie)
+                    .IterateSubTrieNodes(new KeyBytes(Encoding.ASCII.GetBytes(prefix)))
+                    .Count(pair => pair.Node is ValueNode));
+            Assert.Equal(
+                3,
+                ((MerkleTrie)trie)
+                    .IterateSubTrieValues(new KeyBytes(Encoding.ASCII.GetBytes(prefix)))
+                    .Count());
+        }
+
         [Theory]
         [InlineData(true)]
         [InlineData(false)]