Optimizations for KeyBuilder (#3598)

* Optimize KeyBuilder and Add some UTs for not-covered methods

* optimization for adding UInt160 and UInt256 to KeyBuilder

* fix stream capacity calculation

* Update src/Neo/SmartContract/KeyBuilder.cs

Co-authored-by: Shargon <shargon@gmail.com>

* Update src/Neo/SmartContract/KeyBuilder.cs

Co-authored-by: Shargon <shargon@gmail.com>

* add more benchmark

* check IsLittleEndian or not

* add UTs for UInt160 and UInt256

* fix code format

---------

Co-authored-by: Shargon <shargon@gmail.com>
Co-authored-by: Jimmy <jinghui@wayne.edu>

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,75 @@
+// 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_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()
+        {
+            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();
+        }
+
+        [Benchmark]
+        public void KeyBuilder_AddUInt160()
+        {
+            Span<byte> value = stackalloc byte[UInt160.Length];
+            var key = new KeyBuilder(1, 0)
+                .Add(new UInt160(value));
+
+            var bytes = key.ToArray();
+            if (bytes.Length != prefixSize + UInt160.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(including the id and prefix).</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,12 +57,37 @@ 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);
             return this;
         }
 
+        /// <summary>
+        /// Adds part of the key to the builder.
+        /// </summary>
+        /// <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());
+
+        /// <summary>
+        /// Adds part of the key to the builder.
+        /// </summary>
+        /// <param name="key">Part of the key represented by a <see cref="UInt160"/>.</param>
+        /// <returns>A reference to this instance after the add operation has completed.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public KeyBuilder Add(UInt160 key) => Add(key.GetSpan());
+
+        /// <summary>
+        /// Adds part of the key to the builder.
+        /// </summary>
+        /// <param name="key">Part of the key represented by a <see cref="UInt256"/>.</param>
+        /// <returns>A reference to this instance after the add operation has completed.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public KeyBuilder Add(UInt256 key) => Add(key.GetSpan());
+
         /// <summary>
         /// Adds part of the key to the builder.
         /// </summary>
@@ -79,9 +108,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 +122,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 +136,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 +150,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);

src/Neo/UInt160.cs

@@ -101,6 +101,18 @@ public override int GetHashCode()
             return HashCode.Combine(_value1, _value2, _value3);
         }
 
+        /// <summary>
+        /// Gets a ReadOnlySpan that represents the current value in little-endian.
+        /// </summary>
+        /// <returns>A ReadOnlySpan that represents the current value in little-endian.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public ReadOnlySpan<byte> GetSpan()
+        {
+            if (BitConverter.IsLittleEndian)
+                return MemoryMarshal.CreateReadOnlySpan(ref Unsafe.As<ulong, byte>(ref _value1), Length);
+            return this.ToArray().AsSpan(); // Keep the same output as Serialize when BigEndian
+        }
+
         /// <summary>
         /// Parses an <see cref="UInt160"/> from the specified <see cref="string"/>.
         /// </summary>

src/Neo/UInt256.cs

@@ -14,6 +14,7 @@
 using System;
 using System.Globalization;
 using System.IO;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 
 namespace Neo
@@ -101,6 +102,18 @@ public override int GetHashCode()
             return (int)value1;
         }
 
+        /// <summary>
+        /// Gets a ReadOnlySpan that represents the current value in little-endian.
+        /// </summary>
+        /// <returns>A ReadOnlySpan that represents the current value in little-endian.</returns>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public ReadOnlySpan<byte> GetSpan()
+        {
+            if (BitConverter.IsLittleEndian)
+                return MemoryMarshal.CreateReadOnlySpan(ref Unsafe.As<ulong, byte>(ref value1), Length);
+            return this.ToArray().AsSpan(); // Keep the same output as Serialize when BigEndian
+        }
+
         /// <summary>
         /// Parses an <see cref="UInt256"/> from the specified <see cref="string"/>.
         /// </summary>

tests/Neo.UnitTests/SmartContract/UT_KeyBuilder.cs

@@ -11,6 +11,7 @@
 
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using Neo.Extensions;
+using Neo.IO;
 using Neo.SmartContract;
 
 namespace Neo.UnitTests.SmartContract
@@ -42,5 +43,71 @@ public void Test()
             key = key.AddBigEndian(1);
             Assert.AreEqual("010000000000000001", key.ToArray().ToHexString());
         }
+
+        [TestMethod]
+        public void TestAddInt()
+        {
+            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());
+        }
+
+        [TestMethod]
+        public void TestAddUInt()
+        {
+            var key = new KeyBuilder(1, 2);
+            var value = new byte[UInt160.Length];
+            for (int i = 0; i < value.Length; i++)
+                value[i] = (byte)i;
+
+            key = key.Add(new UInt160(value));
+            Assert.AreEqual("0100000002000102030405060708090a0b0c0d0e0f10111213", key.ToArray().ToHexString());
+
+            var key2 = new KeyBuilder(1, 2);
+            key2 = key2.Add((ISerializable)(new UInt160(value)));
+
+            // It must be same before and after optimization.
+            Assert.AreEqual(key.ToArray().ToHexString(), key2.ToArray().ToHexString());
+
+            key = new KeyBuilder(1, 2);
+            value = new byte[UInt256.Length];
+            for (int i = 0; i < value.Length; i++)
+                value[i] = (byte)i;
+            key = key.Add(new UInt256(value));
+            Assert.AreEqual("0100000002000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f", key.ToArray().ToHexString());
+
+            key2 = new KeyBuilder(1, 2);
+            key2 = key2.Add((ISerializable)(new UInt256(value)));
+
+            // It must be same before and after optimization.
+            Assert.AreEqual(key.ToArray().ToHexString(), key2.ToArray().ToHexString());
+        }
     }
 }

tests/Neo.UnitTests/UT_UInt160.cs

@@ -11,6 +11,7 @@
 
 using FluentAssertions;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Neo.Extensions;
 using System;
 using System.Security.Cryptography;
 
@@ -126,5 +127,18 @@ public void TestOperatorSmallerAndEqual()
             Assert.AreEqual(true, UInt160.Zero <= UInt160.Zero);
             Assert.IsTrue(UInt160.Zero >= "0x0000000000000000000000000000000000000000");
         }
+
+        [TestMethod]
+        public void TestSpanAndSerialize()
+        {
+            // random data
+            var random = new Random();
+            var data = new byte[UInt160.Length];
+            random.NextBytes(data);
+
+            var value = new UInt160(data);
+            var span = value.GetSpan();
+            Assert.IsTrue(span.SequenceEqual(value.ToArray()));
+        }
     }
 }

tests/Neo.UnitTests/UT_UInt256.cs

@@ -13,6 +13,7 @@
 
 using FluentAssertions;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Neo.Extensions;
 using Neo.IO;
 using System;
 using System.IO;
@@ -155,5 +156,17 @@ public void TestOperatorSmallerAndEqual()
         {
             Assert.AreEqual(true, UInt256.Zero <= UInt256.Zero);
         }
+
+        [TestMethod]
+        public void TestSpanAndSerialize()
+        {
+            var random = new Random();
+            var data = new byte[UInt256.Length];
+            random.NextBytes(data);
+
+            var value = new UInt256(data);
+            var span = value.GetSpan();
+            Assert.IsTrue(span.SequenceEqual(value.ToArray()));
+        }
     }
 }