Initial implementation for copying subtries

planetarium · Jan 19, 2024 · 9ebf85b · 9ebf85b
1 parent fb17f71
commit 9ebf85b
Showing 1 changed file with 28 additions and 1 deletion.
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;
@@ -64,6 +63,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();