[neox-2.x] apply mpt to storages and persist roots

neo.UnitTests/TestBlockchain.cs

@@ -25,6 +25,7 @@ public static NeoSystem InitializeMockNeoSystem()
                 mockSnapshot.SetupGet(p => p.Assets).Returns(new TestDataCache<UInt256, AssetState>());
                 mockSnapshot.SetupGet(p => p.Contracts).Returns(new TestDataCache<UInt160, ContractState>());
                 mockSnapshot.SetupGet(p => p.Storages).Returns(new TestDataCache<StorageKey, StorageItem>());
+                mockSnapshot.SetupGet(p => p.StateRoots).Returns(new TestDataCache<UInt32Wrapper, StateRootState>());
                 mockSnapshot.SetupGet(p => p.HeaderHashList)
                     .Returns(new TestDataCache<UInt32Wrapper, HeaderHashList>());
                 mockSnapshot.SetupGet(p => p.ValidatorsCount).Returns(new TestMetaDataCache<ValidatorsCountState>());

neo/IO/Caching/DataCache.cs

@@ -60,7 +60,7 @@ public void Add(TKey key, TValue value)
 
         protected abstract void AddInternal(TKey key, TValue value);
 
-        public void Commit()
+        public virtual void Commit()
         {
             foreach (Trackable trackable in GetChangeSet())
                 switch (trackable.State)

neo/Ledger/StateRootState.cs

@@ -0,0 +1,59 @@
+using Neo.IO;
+using Neo.IO.Json;
+using Neo.Network.P2P.Payloads;
+using System.IO;
+
+namespace Neo.Ledger
+{
+    public enum StateRootVerifyFlag : byte
+    {
+        Unverified = 0x00,
+        Verified = 0x01,
+        Invalid = 0x03,
+    }
+
+    public class StateRootState : StateBase, ICloneable<StateRootState>
+    {
+        public StateRootVerifyFlag Flag;
+        public StateRoot StateRoot;
+
+        public override int Size => base.Size + sizeof(StateRootVerifyFlag) + StateRoot.Size;
+
+        StateRootState ICloneable<StateRootState>.Clone()
+        {
+            return new StateRootState
+            {
+                Flag = Flag,
+                StateRoot = StateRoot,
+            };
+        }
+
+        public override void Deserialize(BinaryReader reader)
+        {
+            base.Deserialize(reader);
+            Flag = (StateRootVerifyFlag)reader.ReadByte();
+            StateRoot = reader.ReadSerializable<StateRoot>();
+        }
+
+        void ICloneable<StateRootState>.FromReplica(StateRootState replica)
+        {
+            Flag = replica.Flag;
+            StateRoot = replica.StateRoot;
+        }
+
+        public override void Serialize(BinaryWriter writer)
+        {
+            base.Serialize(writer);
+            writer.Write((byte)Flag);
+            writer.Write(StateRoot);
+        }
+
+        public override JObject ToJson()
+        {
+            JObject json = new JObject();
+            json["flag"] = Flag;
+            json["stateroot"] = StateRoot.ToJson();
+            return json;
+        }
+    }
+}

neo/Persistence/CloneSnapshot.cs

@@ -16,6 +16,7 @@ internal class CloneSnapshot : Snapshot
         public override DataCache<UInt256, AssetState> Assets { get; }
         public override DataCache<UInt160, ContractState> Contracts { get; }
         public override DataCache<StorageKey, StorageItem> Storages { get; }
+        public override DataCache<UInt32Wrapper, StateRootState> StateRoots { get; }
         public override DataCache<UInt32Wrapper, HeaderHashList> HeaderHashList { get; }
         public override MetaDataCache<ValidatorsCountState> ValidatorsCount { get; }
         public override MetaDataCache<HashIndexState> BlockHashIndex { get; }
@@ -33,6 +34,7 @@ public CloneSnapshot(Snapshot snapshot)
             this.Assets = snapshot.Assets.CreateSnapshot();
             this.Contracts = snapshot.Contracts.CreateSnapshot();
             this.Storages = snapshot.Storages.CreateSnapshot();
+            this.StateRoots = snapshot.StateRoots.CreateSnapshot();
             this.HeaderHashList = snapshot.HeaderHashList.CreateSnapshot();
             this.ValidatorsCount = snapshot.ValidatorsCount.CreateSnapshot();
             this.BlockHashIndex = snapshot.BlockHashIndex.CreateSnapshot();

neo/Persistence/IPersistence.cs

@@ -16,6 +16,7 @@ public interface IPersistence
         DataCache<UInt256, AssetState> Assets { get; }
         DataCache<UInt160, ContractState> Contracts { get; }
         DataCache<StorageKey, StorageItem> Storages { get; }
+        DataCache<UInt32Wrapper, StateRootState> StateRoots { get; }
         DataCache<UInt32Wrapper, HeaderHashList> HeaderHashList { get; }
         MetaDataCache<ValidatorsCountState> ValidatorsCount { get; }
         MetaDataCache<HashIndexState> BlockHashIndex { get; }

neo/Persistence/LevelDB/DbCacheWithTrie.cs

@@ -0,0 +1,70 @@
+using Neo.IO;
+using Neo.IO.Caching;
+using Neo.IO.Data.LevelDB;
+using Neo.Trie.MPT;
+using System;
+using System.Collections.Generic;
+
+namespace Neo.Persistence.LevelDB
+{
+    public class DbCacheWithTrie<TKey, TValue> : DataCache<TKey, TValue>
+        where TKey : IEquatable<TKey>, ISerializable, new()
+        where TValue : class, ICloneable<TValue>, ISerializable, new()
+    {
+        private readonly DB db;
+        private readonly ReadOptions options;
+        private readonly WriteBatch batch;
+        private readonly byte prefix;
+        private MPTTrie mptTrie;
+        private DbTrieStore trieDb;
+
+        public DbCacheWithTrie(DB db, ReadOptions options, WriteBatch batch, byte prefix)
+        {
+            this.db = db;
+            this.options = options ?? ReadOptions.Default;
+            this.batch = batch;
+            this.prefix = prefix;
+            this.trieDb = new DbTrieStore(db, options, batch, Prefixes.DATA_MPT);
+            this.mptTrie = new MPTTrie(trieDb.GetRoot(), trieDb);
+        }
+
+        protected override void AddInternal(TKey key, TValue value)
+        {
+            batch?.Put(prefix, key, value);
+            mptTrie.Put(key.ToArray(), value.ToArray());
+        }
+
+        public override void DeleteInternal(TKey key)
+        {
+            batch?.Delete(prefix, key);
+            mptTrie.TryDelete(key.ToArray());
+        }
+
+        protected override IEnumerable<KeyValuePair<TKey, TValue>> FindInternal(byte[] key_prefix)
+        {
+            return db.Find(options, SliceBuilder.Begin(prefix).Add(key_prefix), (k, v) => new KeyValuePair<TKey, TValue>(k.ToArray().AsSerializable<TKey>(1), v.ToArray().AsSerializable<TValue>()));
+        }
+
+        protected override TValue GetInternal(TKey key)
+        {
+            return db.Get<TValue>(options, prefix, key);
+        }
+
+        protected override TValue TryGetInternal(TKey key)
+        {
+            return db.TryGet<TValue>(options, prefix, key);
+        }
+
+        protected override void UpdateInternal(TKey key, TValue value)
+        {
+            batch?.Put(prefix, key, value);
+            mptTrie.Put(key.ToArray(), value.ToArray());
+        }
+
+        public override void Commit()
+        {
+            base.Commit();
+            trieDb.PutRoot(mptTrie.GetRoot());
+        }
+    }
+}

neo/Persistence/LevelDB/DbSnapshot.cs

@@ -12,7 +12,6 @@ internal class DbSnapshot : Snapshot
         private readonly DB db;
         private readonly LSnapshot snapshot;
         private readonly WriteBatch batch;
-
         public override DataCache<UInt256, BlockState> Blocks { get; }
         public override DataCache<UInt256, TransactionState> Transactions { get; }
         public override DataCache<UInt160, AccountState> Accounts { get; }
@@ -22,6 +21,7 @@ internal class DbSnapshot : Snapshot
         public override DataCache<UInt256, AssetState> Assets { get; }
         public override DataCache<UInt160, ContractState> Contracts { get; }
         public override DataCache<StorageKey, StorageItem> Storages { get; }
+        public override DataCache<UInt32Wrapper, StateRootState> StateRoots { get; }
         public override DataCache<UInt32Wrapper, HeaderHashList> HeaderHashList { get; }
         public override MetaDataCache<ValidatorsCountState> ValidatorsCount { get; }
         public override MetaDataCache<HashIndexState> BlockHashIndex { get; }
@@ -41,7 +41,8 @@ public DbSnapshot(DB db)
             Validators = new DbCache<ECPoint, ValidatorState>(db, options, batch, Prefixes.ST_Validator);
             Assets = new DbCache<UInt256, AssetState>(db, options, batch, Prefixes.ST_Asset);
             Contracts = new DbCache<UInt160, ContractState>(db, options, batch, Prefixes.ST_Contract);
-            Storages = new DbCache<StorageKey, StorageItem>(db, options, batch, Prefixes.ST_Storage);
+            Storages = new DbCacheWithTrie<StorageKey, StorageItem>(db, options, batch, Prefixes.ST_Storage);
+            StateRoots = new DbCache<UInt32Wrapper, StateRootState>(db, options, batch, Prefixes.ST_StateRoot);
             HeaderHashList = new DbCache<UInt32Wrapper, HeaderHashList>(db, options, batch, Prefixes.IX_HeaderHashList);
             ValidatorsCount = new DbMetaDataCache<ValidatorsCountState>(db, options, batch, Prefixes.IX_ValidatorsCount);
             BlockHashIndex = new DbMetaDataCache<HashIndexState>(db, options, batch, Prefixes.IX_CurrentBlock);

neo/Persistence/LevelDB/DbTrieStore.cs

@@ -0,0 +1,54 @@
+using Neo.IO.Data.LevelDB;
+using Neo.Trie;
+using Neo.Trie.MPT;
+using System;
+using System.Text;
+
+namespace Neo.Persistence.LevelDB
+{
+    public class DbTrieStore : IKVStore
+    {
+        private readonly DB db;
+        private readonly ReadOptions options;
+        private readonly WriteBatch batch;
+        private readonly byte prefix;
+
+        private readonly byte[] ROOT_KEY = Encoding.ASCII.GetBytes("CURRENT_ROOT");
+
+        public DbTrieStore(DB db, ReadOptions options, WriteBatch batch, byte prefix)
+        {
+            this.db = db;
+            this.options = options;
+            this.batch = batch;
+            this.prefix = prefix;
+        }
+
+        private byte[] StoreKey(byte[] key)
+        {
+            return new byte[] { prefix }.Concat(key);
+        }
+
+        public byte[] Get(byte[] key)
+        {
+            var result = db.TryGet(options, StoreKey(key), out Slice value);
+            return result ? value.ToArray() : null;
+        }
+
+        public void Put(byte[] key, byte[] value)
+        {
+            batch.Put(StoreKey(key), value);
+        }
+
+        public UInt256 GetRoot()
+        {
+            var result = db.TryGet(options, StoreKey(ROOT_KEY), out Slice value);
+            return result ? new UInt256(value.ToArray()) : null;
+        }
+
+        public void PutRoot(UInt256 root)
+        {
+            if (root is null) return;
+            batch.Put(StoreKey(ROOT_KEY), root.ToArray());
+        }
+    }
+}

neo/Persistence/LevelDB/LevelDBStore.cs

@@ -77,6 +77,11 @@ public override DataCache<StorageKey, StorageItem> GetStorages()
             return new DbCache<StorageKey, StorageItem>(db, null, null, Prefixes.ST_Storage);
         }
 
+        public override DataCache<UInt32Wrapper, StateRootState> GetStateRoots()
+        {
+            return new DbCache<UInt32Wrapper, StateRootState>(db, null, null, Prefixes.ST_StateRoot);
+        }
+
         public override DataCache<UInt256, TransactionState> GetTransactions()
         {
             return new DbCache<UInt256, TransactionState>(db, null, null, Prefixes.DATA_Transaction);

neo/Persistence/LevelDB/Prefixes.cs

@@ -4,6 +4,7 @@ internal static class Prefixes
     {
         public const byte DATA_Block = 0x01;
         public const byte DATA_Transaction = 0x02;
+        public const byte DATA_MPT = 0x03;
 
         public const byte ST_Account = 0x40;
         public const byte ST_Coin = 0x44;
@@ -12,6 +13,7 @@ internal static class Prefixes
         public const byte ST_Asset = 0x4c;
         public const byte ST_Contract = 0x50;
         public const byte ST_Storage = 0x70;
+        public const byte ST_StateRoot = 0x71;
 
         public const byte IX_HeaderHashList = 0x80;
         public const byte IX_ValidatorsCount = 0x90;

neo/Persistence/Snapshot.cs

@@ -22,6 +22,7 @@ public abstract class Snapshot : IDisposable, IPersistence, IScriptTable
         public abstract DataCache<UInt256, AssetState> Assets { get; }
         public abstract DataCache<UInt160, ContractState> Contracts { get; }
         public abstract DataCache<StorageKey, StorageItem> Storages { get; }
+        public abstract DataCache<UInt32Wrapper, StateRootState> StateRoots { get; }
         public abstract DataCache<UInt32Wrapper, HeaderHashList> HeaderHashList { get; }
         public abstract MetaDataCache<ValidatorsCountState> ValidatorsCount { get; }
         public abstract MetaDataCache<HashIndexState> BlockHashIndex { get; }
@@ -134,6 +135,7 @@ public virtual void Commit()
             Assets.Commit();
             Contracts.Commit();
             Storages.Commit();
+            StateRoots.Commit();
             HeaderHashList.Commit();
             ValidatorsCount.Commit();
             BlockHashIndex.Commit();

neo/Persistence/Store.cs

@@ -16,6 +16,7 @@ public abstract class Store : IPersistence
         DataCache<UInt256, AssetState> IPersistence.Assets => GetAssets();
         DataCache<UInt160, ContractState> IPersistence.Contracts => GetContracts();
         DataCache<StorageKey, StorageItem> IPersistence.Storages => GetStorages();
+        DataCache<UInt32Wrapper, StateRootState> IPersistence.StateRoots => GetStateRoots();
         DataCache<UInt32Wrapper, HeaderHashList> IPersistence.HeaderHashList => GetHeaderHashList();
         MetaDataCache<ValidatorsCountState> IPersistence.ValidatorsCount => GetValidatorsCount();
         MetaDataCache<HashIndexState> IPersistence.BlockHashIndex => GetBlockHashIndex();
@@ -31,6 +32,7 @@ public abstract class Store : IPersistence
         public abstract DataCache<UInt256, AssetState> GetAssets();
         public abstract DataCache<UInt160, ContractState> GetContracts();
         public abstract DataCache<StorageKey, StorageItem> GetStorages();
+        public abstract DataCache<UInt32Wrapper, StateRootState> GetStateRoots();
         public abstract DataCache<UInt32Wrapper, HeaderHashList> GetHeaderHashList();
         public abstract MetaDataCache<ValidatorsCountState> GetValidatorsCount();
         public abstract MetaDataCache<HashIndexState> GetBlockHashIndex();