Optimizations for KeyBuilder

benchmarks/Neo.Benchmarks/Program.cs

@@ -11,7 +11,9 @@
 
 using BenchmarkDotNet.Running;
 using Neo.Benchmark;
+using Neo.SmartContract.Benchmark;
 
 // BenchmarkRunner.Run<Benchmarks_PoCs>();
 BenchmarkRunner.Run<Benchmarks_UInt160>();
 BenchmarkRunner.Run<Benchmarks_Hash>();
+BenchmarkRunner.Run<Benchmarks_StorageKey>();

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

@@ -0,0 +1,50 @@
+// Copyright (C) 2015-2024 The Neo Project.
+//
+// Benchmarks.StorageKey.cs file belongs to the neo project and is free
+// software distributed under the MIT software license, see the
+// accompanying file LICENSE in the main directory of the
+// repository or http://www.opensource.org/licenses/mit-license.php
+// for more details.
+//
+// Redistribution and use in source and binary forms with or without
+// modifications are permitted.
+
+using BenchmarkDotNet.Attributes;
+using System.Text;
+
+namespace Neo.SmartContract.Benchmark
+{
+    public class Benchmarks_StorageKey
+    {
+        // for avoiding overhead of encoding
+        private static readonly byte[] testBytes = Encoding.ASCII.GetBytes("StorageKey");
+
+        private const int prefixSize = sizeof(int) + sizeof(byte);
+
+        [Benchmark]
+        public void KeyBuilder_AddInt()
+        {
+            var key = new KeyBuilder(1, 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()
+        {
+            var key = new KeyBuilder(1, 0)
+                .Add(testBytes)
+                .Add(testBytes)
+                .Add(testBytes);
+
+            var bytes = key.ToArray();
+            if (bytes.Length != prefixSize + 3 * testBytes.Length)
+                throw new InvalidOperationException();
+        }
+    }
+}

src/Neo/SmartContract/KeyBuilder.cs

@@ -13,6 +13,7 @@
 using System;
 using System.Buffers.Binary;
 using System.IO;
+using System.Runtime.CompilerServices;
 
 namespace Neo.SmartContract
 {
@@ -21,18 +22,20 @@ namespace Neo.SmartContract
     /// </summary>
     public class KeyBuilder
     {
-        private readonly MemoryStream stream = new();
+        private readonly MemoryStream stream;
 
         /// <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>
-        public KeyBuilder(int id, byte prefix)
+        /// <param name="keySizeHint">The hint of the storage key size.</param>
+        public KeyBuilder(int id, byte prefix, int keySizeHint = ApplicationEngine.MaxStorageKeySize)
         {
-            var data = new byte[sizeof(int)];
+            Span<byte> data = stackalloc byte[sizeof(int)];
             BinaryPrimitives.WriteInt32LittleEndian(data, id);
 
+            stream = new(keySizeHint);
             stream.Write(data);
             stream.WriteByte(prefix);
         }
@@ -42,6 +45,7 @@ public KeyBuilder(int id, byte prefix)
         /// </summary>
         /// <param name="key">Part of the key.</param>
         /// <returns>A reference to this instance after the add operation has completed.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public KeyBuilder Add(byte key)
         {
             stream.WriteByte(key);
@@ -53,6 +57,7 @@ public KeyBuilder Add(byte key)
         /// </summary>
         /// <param name="key">Part of the key.</param>
         /// <returns>A reference to this instance after the add operation has completed.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public KeyBuilder Add(ReadOnlySpan<byte> key)
         {
             stream.Write(key);
@@ -79,9 +84,10 @@ public KeyBuilder Add(ISerializable key)
         /// </summary>
         /// <param name="key">Part of the key.</param>
         /// <returns>A reference to this instance after the add operation has completed.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public KeyBuilder AddBigEndian(int key)
         {
-            var data = new byte[sizeof(int)];
+            Span<byte> data = stackalloc byte[sizeof(int)];
             BinaryPrimitives.WriteInt32BigEndian(data, key);
 
             return Add(data);
@@ -92,9 +98,10 @@ public KeyBuilder AddBigEndian(int key)
         /// </summary>
         /// <param name="key">Part of the key.</param>
         /// <returns>A reference to this instance after the add operation has completed.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public KeyBuilder AddBigEndian(uint key)
         {
-            var data = new byte[sizeof(uint)];
+            Span<byte> data = stackalloc byte[sizeof(uint)];
             BinaryPrimitives.WriteUInt32BigEndian(data, key);
 
             return Add(data);
@@ -105,9 +112,10 @@ public KeyBuilder AddBigEndian(uint key)
         /// </summary>
         /// <param name="key">Part of the key.</param>
         /// <returns>A reference to this instance after the add operation has completed.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public KeyBuilder AddBigEndian(long key)
         {
-            var data = new byte[sizeof(long)];
+            Span<byte> data = stackalloc byte[sizeof(long)];
             BinaryPrimitives.WriteInt64BigEndian(data, key);
 
             return Add(data);
@@ -118,9 +126,10 @@ public KeyBuilder AddBigEndian(long key)
         /// </summary>
         /// <param name="key">Part of the key.</param>
         /// <returns>A reference to this instance after the add operation has completed.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public KeyBuilder AddBigEndian(ulong key)
         {
-            var data = new byte[sizeof(ulong)];
+            Span<byte> data = stackalloc byte[sizeof(ulong)];
             BinaryPrimitives.WriteUInt64BigEndian(data, key);
 
             return Add(data);

tests/Neo.UnitTests/SmartContract/UT_KeyBuilder.cs

@@ -42,5 +42,40 @@ public void Test()
             key = key.AddBigEndian(1);
             Assert.AreEqual("010000000000000001", key.ToArray().ToHexString());
         }
+
+        [TestMethod]
+        public void TestAdd()
+        {
+            var key = new KeyBuilder(1, 2);
+            Assert.AreEqual("0100000002", key.ToArray().ToHexString());
+
+            // add int
+            key = new KeyBuilder(1, 2);
+            key = key.AddBigEndian(-1);
+            key = key.AddBigEndian(2);
+            key = key.AddBigEndian(3);
+            Assert.AreEqual("0100000002ffffffff0000000200000003", key.ToArray().ToHexString());
+
+            // add ulong
+            key = new KeyBuilder(1, 2);
+            key = key.AddBigEndian(1ul);
+            key = key.AddBigEndian(2ul);
+            key = key.AddBigEndian(ulong.MaxValue);
+            Assert.AreEqual("010000000200000000000000010000000000000002ffffffffffffffff", key.ToArray().ToHexString());
+
+            // add uint
+            key = new KeyBuilder(1, 2);
+            key = key.AddBigEndian(1u);
+            key = key.AddBigEndian(2u);
+            key = key.AddBigEndian(uint.MaxValue);
+            Assert.AreEqual("01000000020000000100000002ffffffff", key.ToArray().ToHexString());
+
+            // add byte
+            key = new KeyBuilder(1, 2);
+            key = key.Add((byte)1);
+            key = key.Add((byte)2);
+            key = key.Add((byte)3);
+            Assert.AreEqual("0100000002010203", key.ToArray().ToHexString());
+        }
     }
 }