Updated KeyBuilder Fast Access

benchmarks/Neo.Benchmarks/SmartContract/Benchmarks.StorageKey.cs

@@ -34,19 +34,6 @@ public void KeyBuilder_AddInt()
                 throw new InvalidOperationException();
         }
 
-        [Benchmark]
-        public void KeyBuilder_AddIntWithoutPrealloc()
-        {
-            var key = new KeyBuilder(1, 0, 0)
-                .AddBigEndian(1)
-                .AddBigEndian(2)
-                .AddBigEndian(3);
-
-            var bytes = key.ToArray();
-            if (bytes.Length != prefixSize + 3 * sizeof(int))
-                throw new InvalidOperationException();
-        }
-
         [Benchmark]
         public void KeyBuilder_AddBytes()
         {

src/Neo/SmartContract/KeyBuilder.cs

@@ -14,7 +14,6 @@
 using Neo.IO;
 using System;
 using System.Buffers.Binary;
-using System.IO;
 using System.Runtime.CompilerServices;
 
 namespace Neo.SmartContract
@@ -24,22 +23,34 @@ namespace Neo.SmartContract
     /// </summary>
     public class KeyBuilder
     {
-        private readonly MemoryStream _stream;
+        public const int PrefixLength = sizeof(int) + sizeof(byte);
+
+        private readonly Memory<byte> _cacheData;
+        private int _keyLength = 0;
+
 
         /// <summary>
         /// Initializes a new instance of the <see cref="KeyBuilder"/> class.
         /// </summary>
         /// <param name="id">The id of the contract.</param>
         /// <param name="prefix">The prefix of the key.</param>
-        /// <param name="keySizeHint">The hint of the storage key size(including the id and prefix).</param>
-        public KeyBuilder(int id, byte prefix, int keySizeHint = ApplicationEngine.MaxStorageKeySize)
+        /// <param name="maxLength">The hint of the storage key size(Not including the id and prefix).</param>
+        public KeyBuilder(int id, byte prefix, int maxLength = ApplicationEngine.MaxStorageKeySize)
         {
-            Span<byte> data = stackalloc byte[sizeof(int)];
-            BinaryPrimitives.WriteInt32LittleEndian(data, id);
+            if (maxLength < 0)
+                throw new ArgumentOutOfRangeException(nameof(maxLength), "Must be greater than or equal to zero.");
 
-            _stream = new(keySizeHint);
-            _stream.Write(data);
-            _stream.WriteByte(prefix);
+            _cacheData = new byte[maxLength + PrefixLength];
+            BinaryPrimitives.WriteInt32LittleEndian(_cacheData.Span, id);
+
+            _keyLength = sizeof(int);
+            _cacheData.Span[_keyLength++] = prefix;
+        }
+
+        private void CheckLength(int length)
+        {
+            if ((length + _keyLength) > _cacheData.Length)
+                throw new OverflowException("Input data too Large!");
         }
 
         /// <summary>
@@ -50,7 +61,8 @@ public KeyBuilder(int id, byte prefix, int keySizeHint = ApplicationEngine.MaxSt
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public KeyBuilder Add(byte key)
         {
-            _stream.WriteByte(key);
+            CheckLength(1);
+            _cacheData.Span[_keyLength++] = key;
             return this;
         }
 
@@ -62,7 +74,9 @@ public KeyBuilder Add(byte key)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public KeyBuilder Add(ReadOnlySpan<byte> key)
         {
-            _stream.Write(key);
+            CheckLength(key.Length);
+            key.CopyTo(_cacheData.Span[_keyLength..]);
+            _keyLength += key.Length;
             return this;
         }
 
@@ -72,14 +86,16 @@ public KeyBuilder Add(ReadOnlySpan<byte> key)
         /// <param name="key">Part of the key represented by a byte array.</param>
         /// <returns>A reference to this instance after the add operation has completed.</returns>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public KeyBuilder Add(byte[] key) => Add(key.AsSpan());
+        public KeyBuilder Add(byte[] key) =>
+            Add(key.AsSpan());
 
         /// <summary>
         /// Adds part of the key to the builder.
         /// </summary>
         /// <param name="key">Part of the key.</param>
         /// <returns>A reference to this instance after the add operation has completed.</returns>
-        public KeyBuilder Add(ISerializableSpan key) => Add(key.GetSpan());
+        public KeyBuilder Add(ISerializableSpan key) =>
+            Add(key.GetSpan());
 
         /// <summary>
         /// Adds part of the key to the builder in BigEndian.
@@ -141,18 +157,11 @@ public KeyBuilder AddBigEndian(ulong key)
         /// Gets the storage key generated by the builder.
         /// </summary>
         /// <returns>The storage key.</returns>
-        public byte[] ToArray()
-        {
-            using (_stream)
-            {
-                return _stream.ToArray();
-            }
-        }
+        public byte[] ToArray() =>
+            _cacheData[.._keyLength].ToArray();
 
-        public static implicit operator StorageKey(KeyBuilder builder)
-        {
-            return new StorageKey(builder.ToArray());
-        }
+        public static implicit operator StorageKey(KeyBuilder builder) =>
+            new(builder._cacheData[..builder._keyLength].ToArray());
     }
 }