Create KeyBuilder for native contracts

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()
+ };
+ }
+ }
+ }
+}

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)

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>();

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);

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];