Merge pull request #3619 from greymistcube/refactor/copy-states

♻️ Refactor `TrieStateStore.CopyStates()`
planetarium · Jan 19, 2024 · 2b8491a · 2b8491a
2 parents fb17f71 + d4e83e7
commit 2b8491a
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 2 deletions.
diff --git a/Libplanet.Store/TrieStateStore.cs b/Libplanet.Store/TrieStateStore.cs
@@ -1,5 +1,4 @@
 using System;
-using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Security.Cryptography;
@@ -34,8 +33,17 @@ public TrieStateStore(IKeyValueStore stateKeyValueStore)
         public IKeyValueStore StateKeyValueStore { get; }
 
         /// <summary>
+        /// <para>
         /// Copies states under state root hashes of given <paramref name="stateRootHashes"/>
         /// to <paramref name="targetStateStore"/>.
+        /// </para>
+        /// <para>
+        /// Under the hood, this not only copies states directly associated
+        /// with <paramref name="stateRootHashes"/>, but also automatically copies states
+        /// that are not directly associated with <paramref name="stateRootHashes"/>
+        /// but associated with "subtries" with references stored in <see cref="ITrie"/>s
+        /// of <paramref name="stateRootHashes"/>.
+        /// </para>
         /// </summary>
         /// <param name="stateRootHashes">The state root hashes of states to copy.</param>
         /// <param name="targetStateStore">The target state store to copy state root hashes.</param>
@@ -64,6 +72,34 @@ public void CopyStates(
                     targetKeyValueStore.Set(key, value);
                     count++;
                 }
+
+                // FIXME: Probably not the right place to implement this.
+                // It'd be better to have it in Libplanet.Action.State.
+                if (stateTrie.Get(new KeyBytes(Array.Empty<byte>())) is { } metadata)
+                {
+                    foreach (var (path, hash) in stateTrie.IterateValues())
+                    {
+                        // Ignore metadata
+                        if (path.Length > 0)
+                        {
+                            HashDigest<SHA256> accountStateRootHash = new HashDigest<SHA256>(hash);
+                            MerkleTrie accountStateTrie =
+                                (MerkleTrie)GetStateRoot(accountStateRootHash);
+                            if (!accountStateTrie.Recorded)
+                            {
+                                throw new ArgumentException(
+                                    $"Failed to find a state root for given " +
+                                    $"state root hash {accountStateRootHash}.");
+                            }
+
+                            foreach (var (key, value) in accountStateTrie.IterateKeyValuePairs())
+                            {
+                                targetKeyValueStore.Set(key, value);
+                                count++;
+                            }
+                        }
+                    }
+                }
             }
 
             stopwatch.Stop();

diff --git a/Libplanet.Tests/Store/TrieStateStoreTest.cs b/Libplanet.Tests/Store/TrieStateStoreTest.cs
@@ -4,7 +4,9 @@
 using System.Linq;
 using System.Security.Cryptography;
 using Bencodex.Types;
+using Libplanet.Action.State;
 using Libplanet.Common;
+using Libplanet.Crypto;
 using Libplanet.Store;
 using Libplanet.Store.Trie;
 using Xunit;
@@ -76,7 +78,7 @@ public void CopyStates()
             List<(KeyBytes, IValue)> kvs = Enumerable.Range(0, 1_000)
                 .Select(_ =>
                 (
-                    new KeyBytes(GetRandomBytes(random.Next(20))),
+                    new KeyBytes(GetRandomBytes(random.Next(1, 20))),
                     (IValue)new Binary(GetRandomBytes(20))
                 ))
                 .ToList();
@@ -117,6 +119,80 @@ public void CopyStates()
                 targetStateStore.GetStateRoot(trie.Hash).IterateValues().Count());
         }
 
+        [Fact]
+        public void CopyWorldStates()
+        {
+            var stateStore = new TrieStateStore(_stateKeyValueStore);
+            IKeyValueStore targetStateKeyValueStore = new MemoryKeyValueStore();
+            var targetStateStore = new TrieStateStore(targetStateKeyValueStore);
+            Random random = new Random();
+            Dictionary<Address, List<(KeyBytes, IValue)>> data = Enumerable
+                .Range(0, 20)
+                .Select(_ => new Address(GetRandomBytes(Address.Size)))
+                .ToDictionary(
+                    address => address,
+                    _ => Enumerable
+                        .Range(0, 100)
+                        .Select(__ =>
+                        (
+                            new KeyBytes(GetRandomBytes(random.Next(20))),
+                            (IValue)new Binary(GetRandomBytes(20))
+                        ))
+                        .ToList());
+
+            ITrie worldTrie = stateStore.GetStateRoot(null);
+            worldTrie = worldTrie.SetMetadata(new TrieMetadata(5));
+
+            List<HashDigest<SHA256>> accountHashes = new List<HashDigest<SHA256>>();
+            foreach (var elem in data)
+            {
+                ITrie trie = stateStore.GetStateRoot(null);
+                foreach (var kv in elem.Value)
+                {
+                    trie = trie.Set(kv.Item1, kv.Item2);
+                }
+
+                trie = stateStore.Commit(trie);
+                worldTrie = worldTrie.Set(new KeyBytes(elem.Key.ByteArray), trie.Hash.Bencoded);
+                accountHashes.Add(trie.Hash);
+            }
+
+            worldTrie = stateStore.Commit(worldTrie);
+            int prevStatesCount = _stateKeyValueStore.ListKeys().Count();
+
+            // NOTE: Avoid possible collision of KeyBytes, just in case.
+            _stateKeyValueStore.Set(
+                new KeyBytes(GetRandomBytes(30)),
+                ByteUtil.ParseHex("00"));
+            _stateKeyValueStore.Set(
+                new KeyBytes(GetRandomBytes(40)),
+                ByteUtil.ParseHex("00"));
+
+            Assert.Equal(prevStatesCount + 2, _stateKeyValueStore.ListKeys().Count());
+            Assert.Empty(targetStateKeyValueStore.ListKeys());
+
+            stateStore.CopyStates(
+                ImmutableHashSet<HashDigest<SHA256>>.Empty.Add(worldTrie.Hash),
+                targetStateStore);
+
+            // It will stay at the same count of nodes.
+            // FIXME: Bencodex fingerprints also should be tracked.
+            //        https://github.com/planetarium/libplanet/issues/1653
+            Assert.Equal(prevStatesCount, targetStateKeyValueStore.ListKeys().Count());
+            Assert.Equal(
+                worldTrie.IterateNodes().Count(),
+                targetStateStore.GetStateRoot(worldTrie.Hash).IterateNodes().Count());
+            Assert.Equal(
+                worldTrie.IterateValues().Count(),
+                targetStateStore.GetStateRoot(worldTrie.Hash).IterateValues().Count());
+            Assert.Equal(
+                stateStore.GetStateRoot(accountHashes.First()).IterateNodes().Count(),
+                targetStateStore.GetStateRoot(accountHashes.First()).IterateNodes().Count());
+            Assert.Equal(
+                stateStore.GetStateRoot(accountHashes.First()).IterateValues().Count(),
+                targetStateStore.GetStateRoot(accountHashes.First()).IterateValues().Count());
+        }
+
         [Fact]
 #pragma warning disable S2699 // Tests should include assertions
         public void IdempotentDispose()