From 15e15857e36bf539334da37c96cc4b63fc76734c Mon Sep 17 00:00:00 2001
From: Pierre Seznec <petterseznec@gmail.com>
Date: Wed, 25 Sep 2024 16:30:07 +0200
Subject: [PATCH] add read storage for provider

---
 src/client/publicAPI.ts                     | 17 +++++++------
 src/client/types.ts                         |  2 +-
 src/provider/interface.ts                   | 12 +++++++++
 src/provider/web3Provider/smartContracts.ts | 27 ++++++++++++++++++++-
 4 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/src/client/publicAPI.ts b/src/client/publicAPI.ts
index 986473b8..ee98210c 100644
--- a/src/client/publicAPI.ts
+++ b/src/client/publicAPI.ts
@@ -219,10 +219,14 @@ export class PublicAPI {
     inputs: DatastoreEntry[],
     final = true
   ): Promise<Uint8Array[]> {
-    const entriesQuery = inputs.map((entry) => ({
-      key: Array.from(entry.key),
-      address: entry.address,
-    }))
+    const entriesQuery = inputs.map((entry) => {
+      const byteKey: Uint8Array =
+        typeof entry.key === 'string' ? strToBytes(entry.key) : entry.key
+      return {
+        key: Array.from(byteKey),
+        address: entry.address,
+      }
+    })
     const res = await withRetry(
       () => this.connector.get_datastore_entries(entriesQuery),
       this.options.retry!
@@ -238,10 +242,7 @@ export class PublicAPI {
     address: string,
     final = true
   ): Promise<Uint8Array> {
-    const byteKey: Uint8Array = typeof key === 'string' ? strToBytes(key) : key
-    return this.getDatastoreEntries([{ key: byteKey, address }], final).then(
-      (r) => r[0]
-    )
+    return this.getDatastoreEntries([{ key, address }], final).then((r) => r[0])
   }
 
   async getSlotTransfers(slot: Slot): Promise<Transfer[]> {
diff --git a/src/client/types.ts b/src/client/types.ts
index ea4f5ae5..7b793300 100644
--- a/src/client/types.ts
+++ b/src/client/types.ts
@@ -62,5 +62,5 @@ export type ReadOnlyCallResult = {
 
 export type DatastoreEntry = {
   address: string
-  key: Uint8Array
+  key: Uint8Array | string
 }
diff --git a/src/provider/interface.ts b/src/provider/interface.ts
index f0845f98..d5bb89cf 100644
--- a/src/provider/interface.ts
+++ b/src/provider/interface.ts
@@ -40,4 +40,16 @@ export type Provider = {
   getOperationStatus(opId: string): Promise<OperationStatus>
   getEvents(filter: EventFilter): Promise<SCEvent[]>
   getNodeStatus(): Promise<NodeStatusInfo>
+
+  /** Storage */
+  getStorageKeys(
+    address: string,
+    filter: Uint8Array | string,
+    final?: boolean
+  ): Promise<Uint8Array[]>
+  readStorage(
+    address: string,
+    keys: Uint8Array[] | string[],
+    final?: boolean
+  ): Promise<Uint8Array[]>
 }
diff --git a/src/provider/web3Provider/smartContracts.ts b/src/provider/web3Provider/smartContracts.ts
index 73ba982c..71d0fed9 100644
--- a/src/provider/web3Provider/smartContracts.ts
+++ b/src/provider/web3Provider/smartContracts.ts
@@ -1,4 +1,10 @@
-import { Account, minBigInt, PublicAPI } from '../..'
+import {
+  Account,
+  DatastoreEntry,
+  minBigInt,
+  PublicAPI,
+  strToBytes,
+} from '../..'
 import { U64 } from '../../basicElements/serializers/number/u64'
 import { ErrorInsufficientBalance, ErrorMaxGas } from '../../errors'
 import { MAX_GAS_CALL, MIN_GAS_CALL } from '../../smartContracts'
@@ -168,4 +174,23 @@ export class SCProvider {
       datastore,
     })
   }
+
+  public async getStorageKeys(
+    address: string,
+    filter: Uint8Array | string = new Uint8Array(),
+    final = true
+  ): Promise<Uint8Array[]> {
+    const filterBytes: Uint8Array =
+      typeof filter === 'string' ? strToBytes(filter) : filter
+    return this.client.getDataStoreKeys(address, filterBytes, final)
+  }
+
+  public async readStorage(
+    address: string,
+    keys: Uint8Array[] | string[],
+    final = true
+  ): Promise<Uint8Array[]> {
+    const entries: DatastoreEntry[] = keys.map((key) => ({ address, key }))
+    return this.client.getDatastoreEntries(entries, final)
+  }
 }