Create KeyBuilder (#1748)

neo-project · Jul 4, 2020 · 54095fc · 54095fc
1 parent ad5816a
commit 54095fc
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 32 deletions.
diff --git a/src/neo/SmartContract/Native/KeyBuilder.cs b/src/neo/SmartContract/Native/KeyBuilder.cs
@@ -0,0 +1,52 @@
+using Neo.IO;
+using Neo.Ledger;
+using System;
+using System.IO;
+
+namespace Neo.SmartContract.Native
+{
+    internal class KeyBuilder
+    {
+        private readonly int id;
+        private readonly MemoryStream stream = new MemoryStream();
+
+        public KeyBuilder(int id, byte prefix)
+        {
+            this.id = id;
+            this.stream.WriteByte(prefix);
+        }
+
+        public KeyBuilder Add(ReadOnlySpan<byte> key)
+        {
+            stream.Write(key);
+            return this;
+        }
+
+        public KeyBuilder Add(ISerializable key)
+        {
+            using (BinaryWriter writer = new BinaryWriter(stream, Utility.StrictUTF8, true))
+            {
+                key.Serialize(writer);
+                writer.Flush();
+            }
+            return this;
+        }
+
+        unsafe public KeyBuilder Add<T>(T key) where T : unmanaged
+        {
+            return Add(new ReadOnlySpan<byte>(&key, sizeof(T)));
+        }
+
+        public static implicit operator StorageKey(KeyBuilder builder)
+        {
+            using (builder.stream)
+            {
+                return new StorageKey
+                {
+                    Id = builder.id,
+                    Key = builder.stream.ToArray()
+                };
+            }
+        }
+    }
+}
diff --git a/src/neo/SmartContract/Native/NativeContract.cs b/src/neo/SmartContract/Native/NativeContract.cs
@@ -1,5 +1,4 @@
 using Neo.IO;
-using Neo.Ledger;
 using Neo.SmartContract.Manifest;
 using Neo.SmartContract.Native.Tokens;
 using Neo.VM;
@@ -79,21 +78,9 @@ protected NativeContract()
             contractsHashDictionary.Add(Hash, this);
         }
 
-        protected StorageKey CreateStorageKey(byte prefix, byte[] key = null)
+        private protected KeyBuilder CreateStorageKey(byte prefix)
         {
-            StorageKey storageKey = new StorageKey
-            {
-                Id = Id,
-                Key = new byte[sizeof(byte) + (key?.Length ?? 0)]
-            };
-            storageKey.Key[0] = prefix;
-            key?.CopyTo(storageKey.Key.AsSpan(1));
-            return storageKey;
-        }
-
-        internal protected StorageKey CreateStorageKey(byte prefix, ISerializable key)
-        {
-            return CreateStorageKey(prefix, key.ToArray());
+            return new KeyBuilder(Id, prefix);
         }
 
         public static NativeContract GetContract(UInt160 hash)

diff --git a/src/neo/SmartContract/Native/Tokens/NeoToken.cs b/src/neo/SmartContract/Native/Tokens/NeoToken.cs
@@ -40,7 +40,7 @@ protected override void OnBalanceChanging(ApplicationEngine engine, UInt160 acco
             if (amount.IsZero) return;
             if (state.VoteTo != null)
             {
-                StorageItem storage_validator = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Candidate, state.VoteTo.ToArray()));
+                StorageItem storage_validator = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Candidate).Add(state.VoteTo));
                 CandidateState state_validator = storage_validator.GetInteroperable<CandidateState>();
                 state_validator.Votes += amount;
             }
@@ -112,7 +112,7 @@ protected override void OnPersist(ApplicationEngine engine)
         [ContractMethod(0_03000000, CallFlags.AllowStates)]
         public BigInteger UnclaimedGas(StoreView snapshot, UInt160 account, uint end)
         {
-            StorageItem storage = snapshot.Storages.TryGet(CreateAccountKey(account));
+            StorageItem storage = snapshot.Storages.TryGet(CreateStorageKey(Prefix_Account).Add(account));
             if (storage is null) return BigInteger.Zero;
             NeoAccountState state = storage.GetInteroperable<NeoAccountState>();
             return CalculateBonus(state.Balance, state.BalanceHeight, end);
@@ -129,7 +129,7 @@ private bool RegisterCandidate(ApplicationEngine engine, ECPoint pubkey)
 
         private void RegisterCandidateInternal(StoreView snapshot, ECPoint pubkey)
         {
-            StorageKey key = CreateStorageKey(Prefix_Candidate, pubkey);
+            StorageKey key = CreateStorageKey(Prefix_Candidate).Add(pubkey);
             StorageItem item = snapshot.Storages.GetAndChange(key, () => new StorageItem(new CandidateState()));
             CandidateState state = item.GetInteroperable<CandidateState>();
             state.Registered = true;
@@ -140,7 +140,7 @@ private bool UnregisterCandidate(ApplicationEngine engine, ECPoint pubkey)
         {
             if (!engine.CheckWitnessInternal(Contract.CreateSignatureRedeemScript(pubkey).ToScriptHash()))
                 return false;
-            StorageKey key = CreateStorageKey(Prefix_Candidate, pubkey);
+            StorageKey key = CreateStorageKey(Prefix_Candidate).Add(pubkey);
             if (engine.Snapshot.Storages.TryGet(key) is null) return true;
             StorageItem item = engine.Snapshot.Storages.GetAndChange(key);
             CandidateState state = item.GetInteroperable<CandidateState>();
@@ -160,13 +160,13 @@ private bool Vote(ApplicationEngine engine, UInt160 account, ECPoint voteTo)
 
         private bool VoteInternal(StoreView snapshot, UInt160 account, ECPoint voteTo)
         {
-            StorageKey key_account = CreateAccountKey(account);
+            StorageKey key_account = CreateStorageKey(Prefix_Account).Add(account);
             if (snapshot.Storages.TryGet(key_account) is null) return false;
             StorageItem storage_account = snapshot.Storages.GetAndChange(key_account);
             NeoAccountState state_account = storage_account.GetInteroperable<NeoAccountState>();
             if (state_account.VoteTo != null)
             {
-                StorageKey key = CreateStorageKey(Prefix_Candidate, state_account.VoteTo.ToArray());
+                StorageKey key = CreateStorageKey(Prefix_Candidate).Add(state_account.VoteTo);
                 StorageItem storage_validator = snapshot.Storages.GetAndChange(key);
                 CandidateState state_validator = storage_validator.GetInteroperable<CandidateState>();
                 state_validator.Votes -= state_account.Balance;
@@ -176,7 +176,7 @@ private bool VoteInternal(StoreView snapshot, UInt160 account, ECPoint voteTo)
             state_account.VoteTo = voteTo;
             if (voteTo != null)
             {
-                StorageKey key = CreateStorageKey(Prefix_Candidate, voteTo.ToArray());
+                StorageKey key = CreateStorageKey(Prefix_Candidate).Add(voteTo);
                 if (snapshot.Storages.TryGet(key) is null) return false;
                 StorageItem storage_validator = snapshot.Storages.GetAndChange(key);
                 CandidateState state_validator = storage_validator.GetInteroperable<CandidateState>();

diff --git a/src/neo/SmartContract/Native/Tokens/Nep5Token.cs b/src/neo/SmartContract/Native/Tokens/Nep5Token.cs
@@ -58,16 +58,11 @@ protected Nep5Token()
             Manifest.Abi.Events = events.ToArray();
         }
 
-        protected StorageKey CreateAccountKey(UInt160 account)
-        {
-            return CreateStorageKey(Prefix_Account, account);
-        }
-
         internal protected virtual void Mint(ApplicationEngine engine, UInt160 account, BigInteger amount)
         {
             if (amount.Sign < 0) throw new ArgumentOutOfRangeException(nameof(amount));
             if (amount.IsZero) return;
-            StorageItem storage = engine.Snapshot.Storages.GetAndChange(CreateAccountKey(account), () => new StorageItem(new TState()));
+            StorageItem storage = engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Account).Add(account), () => new StorageItem(new TState()));
             TState state = storage.GetInteroperable<TState>();
             OnBalanceChanging(engine, account, state, amount);
             state.Balance += amount;
@@ -85,7 +80,7 @@ internal protected virtual void Burn(ApplicationEngine engine, UInt160 account,
         {
             if (amount.Sign < 0) throw new ArgumentOutOfRangeException(nameof(amount));
             if (amount.IsZero) return;
-            StorageKey key = CreateAccountKey(account);
+            StorageKey key = CreateStorageKey(Prefix_Account).Add(account);
             StorageItem storage = engine.Snapshot.Storages.GetAndChange(key);
             TState state = storage.GetInteroperable<TState>();
             if (state.Balance < amount) throw new InvalidOperationException();
@@ -112,7 +107,7 @@ public virtual BigInteger TotalSupply(StoreView snapshot)
         [ContractMethod(0_01000000, CallFlags.AllowStates)]
         public virtual BigInteger BalanceOf(StoreView snapshot, UInt160 account)
         {
-            StorageItem storage = snapshot.Storages.TryGet(CreateAccountKey(account));
+            StorageItem storage = snapshot.Storages.TryGet(CreateStorageKey(Prefix_Account).Add(account));
             if (storage is null) return BigInteger.Zero;
             return storage.GetInteroperable<TState>().Balance;
         }
@@ -125,7 +120,7 @@ protected virtual bool Transfer(ApplicationEngine engine, UInt160 from, UInt160
                 return false;
             ContractState contract_to = engine.Snapshot.Contracts.TryGet(to);
             if (contract_to?.Payable == false) return false;
-            StorageKey key_from = CreateAccountKey(from);
+            StorageKey key_from = CreateStorageKey(Prefix_Account).Add(from);
             StorageItem storage_from = engine.Snapshot.Storages.GetAndChange(key_from);
             if (amount.IsZero)
             {
@@ -151,7 +146,7 @@ protected virtual bool Transfer(ApplicationEngine engine, UInt160 from, UInt160
                         engine.Snapshot.Storages.Delete(key_from);
                     else
                         state_from.Balance -= amount;
-                    StorageKey key_to = CreateAccountKey(to);
+                    StorageKey key_to = CreateStorageKey(Prefix_Account).Add(to);
                     StorageItem storage_to = engine.Snapshot.Storages.GetAndChange(key_to, () => new StorageItem(new TState()));
                     TState state_to = storage_to.GetInteroperable<TState>();
                     OnBalanceChanging(engine, to, state_to, amount);

diff --git a/tests/neo.UnitTests/TestUtils.cs b/tests/neo.UnitTests/TestUtils.cs
@@ -5,6 +5,7 @@
 using Neo.Network.P2P.Payloads;
 using Neo.SmartContract;
 using Neo.SmartContract.Manifest;
+using Neo.SmartContract.Native;
 using Neo.VM;
 using Neo.Wallets.NEP6;
 using System;
@@ -55,6 +56,15 @@ public static ContractManifest CreateManifest(UInt160 hash, string method, Contr
             return manifest;
         }
 
+        public static StorageKey CreateStorageKey(this NativeContract contract, byte prefix, ISerializable key)
+        {
+            return new StorageKey
+            {
+                Id = contract.Id,
+                Key = key.ToArray().Prepend(prefix).ToArray()
+            };
+        }
+
         public static byte[] GetByteArray(int length, byte firstByte)
         {
             byte[] array = new byte[length];