Add MPT (neo-project#1442)

src/neo/Cryptography/MPT/BranchNode.cs

@@ -0,0 +1,35 @@
+using System.IO;
+
+namespace Neo.Cryptography.MPT
+{
+    public class BranchNode : MPTNode
+    {
+        public const int ChildCount = 17;
+        public readonly MPTNode[] Children = new MPTNode[ChildCount];
+
+        protected override NodeType Type => NodeType.BranchNode;
+
+        public BranchNode()
+        {
+            for (int i = 0; i < ChildCount; i++)
+            {
+                Children[i] = HashNode.EmptyNode;
+            }
+        }
+
+        internal override void EncodeSpecific(BinaryWriter writer)
+        {
+            for (int i = 0; i < ChildCount; i++)
+                WriteHash(writer, Children[i].Hash);
+        }
+
+        internal override void DecodeSpecific(BinaryReader reader)
+        {
+            for (int i = 0; i < ChildCount; i++)
+            {
+                Children[i] = new HashNode();
+                Children[i].DecodeSpecific(reader);
+            }
+        }
+    }
+}

src/neo/Cryptography/MPT/ExtensionNode.cs

@@ -0,0 +1,30 @@
+using Neo.IO;
+using Neo.SmartContract;
+using System.IO;
+
+namespace Neo.Cryptography.MPT
+{
+    public class ExtensionNode : MPTNode
+    {
+        //max lenght when store StorageKey
+        public const int MaxKeyLength = (ApplicationEngine.MaxStorageValueSize + sizeof(int)) * 2;
+
+        public byte[] Key;
+        public MPTNode Next;
+
+        protected override NodeType Type => NodeType.ExtensionNode;
+
+        internal override void EncodeSpecific(BinaryWriter writer)
+        {
+            writer.WriteVarBytes(Key);
+            WriteHash(writer, Next.Hash);
+        }
+
+        internal override void DecodeSpecific(BinaryReader reader)
+        {
+            Key = reader.ReadVarBytes(MaxKeyLength);
+            Next = new HashNode();
+            Next.DecodeSpecific(reader);
+        }
+    }
+}

src/neo/Cryptography/MPT/HashNode.cs

@@ -0,0 +1,41 @@
+using Neo.IO;
+using System;
+using System.IO;
+
+namespace Neo.Cryptography.MPT
+{
+    public class HashNode : MPTNode
+    {
+        private UInt256 hash;
+
+        public override UInt256 Hash => hash;
+        protected override NodeType Type => NodeType.HashNode;
+        public bool IsEmpty => Hash is null;
+        public static HashNode EmptyNode { get; } = new HashNode();
+
+        public HashNode()
+        {
+        }
+
+        public HashNode(UInt256 hash)
+        {
+            this.hash = hash;
+        }
+
+        internal override void EncodeSpecific(BinaryWriter writer)
+        {
+            WriteHash(writer, hash);
+        }
+
+        internal override void DecodeSpecific(BinaryReader reader)
+        {
+            byte[] buffer = reader.ReadVarBytes(UInt256.Length);
+            hash = buffer.Length switch
+            {
+                0 => null,
+                UInt256.Length => new UInt256(buffer),
+                _ => throw new FormatException()
+            };
+        }
+    }
+}

src/neo/Cryptography/MPT/LeafNode.cs

@@ -0,0 +1,36 @@
+using Neo.IO;
+using Neo.SmartContract;
+using System;
+using System.IO;
+
+namespace Neo.Cryptography.MPT
+{
+    public class LeafNode : MPTNode
+    {
+        //the max size when store StorageItem
+        public const int MaxValueLength = 3 + ApplicationEngine.MaxStorageValueSize + sizeof(bool);
+
+        public byte[] Value;
+
+        protected override NodeType Type => NodeType.LeafNode;
+
+        public LeafNode()
+        {
+        }
+
+        public LeafNode(ReadOnlySpan<byte> value)
+        {
+            Value = value.ToArray();
+        }
+
+        internal override void EncodeSpecific(BinaryWriter writer)
+        {
+            writer.WriteVarBytes(Value);
+        }
+
+        internal override void DecodeSpecific(BinaryReader reader)
+        {
+            Value = reader.ReadVarBytes(MaxValueLength);
+        }
+    }
+}

src/neo/Cryptography/MPT/MPTNode.cs

@@ -0,0 +1,61 @@
+using Neo.IO;
+using Neo.IO.Caching;
+using System;
+using System.IO;
+
+namespace Neo.Cryptography.MPT
+{
+    public abstract class MPTNode
+    {
+        private UInt256 hash;
+
+        public virtual UInt256 Hash => hash ??= new UInt256(Crypto.Hash256(Encode()));
+        protected abstract NodeType Type { get; }
+
+        public void SetDirty()
+        {
+            hash = null;
+        }
+
+        public byte[] Encode()
+        {
+            using MemoryStream ms = new MemoryStream();
+            using BinaryWriter writer = new BinaryWriter(ms);
+
+            writer.Write((byte)Type);
+            EncodeSpecific(writer);
+            writer.Flush();
+
+            return ms.ToArray();
+        }
+
+        internal abstract void EncodeSpecific(BinaryWriter writer);
+
+        public static unsafe MPTNode Decode(ReadOnlySpan<byte> data)
+        {
+            if (data.IsEmpty) return null;
+
+            fixed (byte* pointer = data)
+            {
+                using UnmanagedMemoryStream stream = new UnmanagedMemoryStream(pointer, data.Length);
+                using BinaryReader reader = new BinaryReader(stream);
+
+                MPTNode n = (MPTNode)ReflectionCache<NodeType>.CreateInstance((NodeType)reader.ReadByte());
+                if (n is null) throw new InvalidOperationException();
+
+                n.DecodeSpecific(reader);
+                return n;
+            }
+        }
+
+        internal abstract void DecodeSpecific(BinaryReader reader);
+
+        protected void WriteHash(BinaryWriter writer, UInt256 hash)
+        {
+            if (hash is null)
+                writer.Write((byte)0);
+            else
+                writer.WriteVarBytes(hash.ToArray());
+        }
+    }
+}

src/neo/Cryptography/MPT/MPTNodeType.cs

@@ -0,0 +1,16 @@
+using Neo.IO.Caching;
+
+namespace Neo.Cryptography.MPT
+{
+    public enum NodeType : byte
+    {
+        [ReflectionCache(typeof(BranchNode))]
+        BranchNode = 0x00,
+        [ReflectionCache(typeof(ExtensionNode))]
+        ExtensionNode = 0x01,
+        [ReflectionCache(typeof(HashNode))]
+        HashNode = 0x02,
+        [ReflectionCache(typeof(LeafNode))]
+        LeafNode = 0x03,
+    }
+}

src/neo/Cryptography/MPT/MPTTrie.Delete.cs

@@ -0,0 +1,120 @@
+using Neo.IO;
+using System;
+using System.Collections.Generic;
+using static Neo.Helper;
+
+namespace Neo.Cryptography.MPT
+{
+    partial class MPTTrie<TKey, TValue>
+    {
+        public bool Delete(TKey key)
+        {
+            var path = ToNibbles(key.ToArray());
+            if (path.Length == 0) return false;
+            return TryDelete(ref root, path);
+        }
+
+        private bool TryDelete(ref MPTNode node, ReadOnlySpan<byte> path)
+        {
+            switch (node)
+            {
+                case LeafNode _:
+                    {
+                        if (path.IsEmpty)
+                        {
+                            node = HashNode.EmptyNode;
+                            return true;
+                        }
+                        return false;
+                    }
+                case ExtensionNode extensionNode:
+                    {
+                        if (path.StartsWith(extensionNode.Key))
+                        {
+                            var result = TryDelete(ref extensionNode.Next, path[extensionNode.Key.Length..]);
+                            if (!result) return false;
+                            if (extensionNode.Next is HashNode hashNode && hashNode.IsEmpty)
+                            {
+                                node = extensionNode.Next;
+                                return true;
+                            }
+                            if (extensionNode.Next is ExtensionNode sn)
+                            {
+                                extensionNode.Key = Concat(extensionNode.Key, sn.Key);
+                                extensionNode.Next = sn.Next;
+                            }
+                            extensionNode.SetDirty();
+                            PutToStore(extensionNode);
+                            return true;
+                        }
+                        return false;
+                    }
+                case BranchNode branchNode:
+                    {
+                        bool result;
+                        if (path.IsEmpty)
+                        {
+                            result = TryDelete(ref branchNode.Children[BranchNode.ChildCount - 1], path);
+                        }
+                        else
+                        {
+                            result = TryDelete(ref branchNode.Children[path[0]], path[1..]);
+                        }
+                        if (!result) return false;
+                        List<byte> childrenIndexes = new List<byte>(BranchNode.ChildCount);
+                        for (int i = 0; i < BranchNode.ChildCount; i++)
+                        {
+                            if (branchNode.Children[i] is HashNode hn && hn.IsEmpty) continue;
+                            childrenIndexes.Add((byte)i);
+                        }
+                        if (childrenIndexes.Count > 1)
+                        {
+                            branchNode.SetDirty();
+                            PutToStore(branchNode);
+                            return true;
+                        }
+                        var lastChildIndex = childrenIndexes[0];
+                        var lastChild = branchNode.Children[lastChildIndex];
+                        if (lastChildIndex == BranchNode.ChildCount - 1)
+                        {
+                            node = lastChild;
+                            return true;
+                        }
+                        if (lastChild is HashNode hashNode)
+                        {
+                            lastChild = Resolve(hashNode);
+                            if (lastChild is null) return false;
+                        }
+                        if (lastChild is ExtensionNode exNode)
+                        {
+                            exNode.Key = Concat(childrenIndexes.ToArray(), exNode.Key);
+                            exNode.SetDirty();
+                            PutToStore(exNode);
+                            node = exNode;
+                            return true;
+                        }
+                        node = new ExtensionNode()
+                        {
+                            Key = childrenIndexes.ToArray(),
+                            Next = lastChild,
+                        };
+                        PutToStore(node);
+                        return true;
+                    }
+                case HashNode hashNode:
+                    {
+                        if (hashNode.IsEmpty)
+                        {
+                            return true;
+                        }
+                        var newNode = Resolve(hashNode);
+                        if (newNode is null) return false;
+                        node = newNode;
+                        return TryDelete(ref node, path);
+                    }
+                default:
+                    return false;
+            }
+        }
+    }
+}