From 7f1d886a4aff7c04d0566bc04b8232af66039305 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 22 Aug 2022 22:47:56 +0300
Subject: [PATCH] Native: store ID->hash map in ContractManagement, add
 iterator

That probably is the lighest design possible, store and retrieve hashes
only. But it can be expanded to return full contract state (getContractByID,
for example) and probably even iterate over contract states.

Fixes #2803.
---
 .../Native/ContractManagement.cs              | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/src/Neo/SmartContract/Native/ContractManagement.cs b/src/Neo/SmartContract/Native/ContractManagement.cs
index 6b450df8e7d..94c9ddd9e89 100644
--- a/src/Neo/SmartContract/Native/ContractManagement.cs
+++ b/src/Neo/SmartContract/Native/ContractManagement.cs
@@ -13,9 +13,11 @@
 using Neo.IO;
 using Neo.Network.P2P.Payloads;
 using Neo.Persistence;
+using Neo.SmartContract.Iterators;
 using Neo.SmartContract.Manifest;
 using Neo.VM.Types;
 using System;
+using System.Buffers.Binary;
 using System.Collections.Generic;
 using System.Linq;
 using System.Numerics;
@@ -28,6 +30,7 @@ namespace Neo.SmartContract.Native
     public sealed class ContractManagement : NativeContract
     {
         private const byte Prefix_MinimumDeploymentFee = 20;
+        private const byte Prefix_ContractHash = 17;
         private const byte Prefix_NextAvailableId = 15;
         private const byte Prefix_Contract = 8;
 
@@ -143,6 +146,43 @@ public ContractState GetContract(DataCache snapshot, UInt160 hash)
             return snapshot.TryGet(CreateStorageKey(Prefix_Contract).Add(hash))?.GetInteroperable<ContractState>();
         }
 
+        /// <summary>
+        /// Maps specified ID to deployed contract hash if it's available.
+        /// </summary>
+        /// <param name="snapshot">The snapshot used to read data.</param>
+        /// <param name="id">Contract ID.</param>
+        /// <returns>Contract hash.</returns>
+        [ContractMethod(CpuFee = 1 << 15, RequiredCallFlags = CallFlags.ReadStates)]
+        public UInt256 GetContractHash(DataCache snapshot, int id)
+        {
+            StorageItem item = snapshot.TryGet(CreateStorageKey(Prefix_ContractHash).AddBigEndian(id));
+            if (item is null) return null;
+            return new UInt256(item.Value.Span);
+        }
+
+        /// <summary>
+        /// Gets hashes of all deployed contracts.
+        /// </summary>
+        /// <param name="snapshot">The snapshot used to read data.</param>
+        /// <returns>Iterator with hashes of all deployed contracts.</returns>
+        [ContractMethod(CpuFee = 1 << 15, RequiredCallFlags = CallFlags.ReadStates)]
+        private IIterator GetContractHashes(DataCache snapshot)
+        {
+            const FindOptions options = FindOptions.RemovePrefix;
+            var enumerator = GetContractHashesInternal(snapshot)
+                .Select(p => (p.Key, p.Value))
+                .GetEnumerator();
+            return new StorageIterator(enumerator, 1, options);
+        }
+
+        internal IEnumerable<(StorageKey Key, StorageItem Value, int Id, UInt256 Hash)> GetContractHashesInternal(DataCache snapshot)
+        {
+            byte[] prefix_key = CreateStorageKey(Prefix_ContractHash).ToArray();
+            return snapshot.Find(prefix_key)
+                .Select(p => (p.Key, p.Value, Id: BinaryPrimitives.ReadInt32BigEndian(p.Key.Key.Span[1..]), Hash: new UInt256(p.Value.Value.Span)))
+                .Where(p => p.Id >= 0);
+        }
+
         /// <summary>
         /// Check if a method exists in a contract.
         /// </summary>
@@ -215,6 +255,7 @@ private async ContractTask<ContractState> Deploy(ApplicationEngine engine, byte[
             if (!contract.Manifest.IsValid(hash)) throw new InvalidOperationException($"Invalid Manifest Hash: {hash}");
 
             engine.Snapshot.Add(key, new StorageItem(contract));
+            engine.Snapshot.Add(CreateStorageKey(Prefix_ContractHash).AddBigEndian(contract.Id), new StorageItem(hash.ToArray()));
 
             await OnDeploy(engine, contract, data, false);