Skip to content

Commit

Permalink
Create KeyBuilder (neo-project#1748)
Browse files Browse the repository at this point in the history
  • Loading branch information
@erikzhang @ShawnYun
erikzhang authored and ShawnYun committed Jul 6, 2020
1 parent 58e62a1 commit d7b8725
Show file tree
Hide file tree
Showing 5 changed files with 76 additions and 32 deletions.
52 changes: 52 additions & 0 deletions src/neo/SmartContract/Native/KeyBuilder.cs
View file
Original file line number Diff line number Diff line change
@@ -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()
};
}
}
}
}
17 changes: 2 additions & 15 deletions src/neo/SmartContract/Native/NativeContract.cs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
using Neo.IO;
using Neo.Ledger;
using Neo.SmartContract.Manifest;
using Neo.SmartContract.Native.Tokens;
using Neo.VM;
Expand Down Expand Up @@ -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)
Expand Down
14 changes: 7 additions & 7 deletions src/neo/SmartContract/Native/Tokens/NeoToken.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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);
Expand All @@ -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;
Expand All @@ -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>();
Expand All @@ -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;
Expand All @@ -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>();
Expand Down
15 changes: 5 additions & 10 deletions src/neo/SmartContract/Native/Tokens/Nep5Token.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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();
Expand All @@ -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;
}
Expand All @@ -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)
{
Expand All @@ -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);
Expand Down
10 changes: 10 additions & 0 deletions tests/neo.UnitTests/TestUtils.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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];
Expand Down

0 comments on commit d7b8725

Please sign in to comment.