From c846f1edc4c24507fe57f9eb709e400da8cddd08 Mon Sep 17 00:00:00 2001
From: Harry Solovay <harrysolovay@gmail.com>
Date: Mon, 27 Feb 2023 10:51:12 -0500
Subject: [PATCH] feat: `storageRune.size` method + `sizeTree` pattern (#657)

---
 _tasks/dnt.ts         |  3 +++
 examples/size_tree.ts |  6 ++++++
 fluent/StorageRune.ts | 11 +++++++++++
 patterns/sizeTree.ts  | 23 +++++++++++++++++++++++
 4 files changed, 43 insertions(+)
 create mode 100644 examples/size_tree.ts
 create mode 100644 patterns/sizeTree.ts

diff --git a/_tasks/dnt.ts b/_tasks/dnt.ts
index 4dd95ec14..bde432998 100755
--- a/_tasks/dnt.ts
+++ b/_tasks/dnt.ts
@@ -36,6 +36,9 @@ await Promise.all([
     }, {
       name: "./patterns/identity",
       path: "./patterns/identity.ts",
+    }, {
+      name: "./patterns/size_tree",
+      path: "./patterns/sizeTree.ts",
     }, {
       name: "./server",
       path: "./server/mod.ts",
diff --git a/examples/size_tree.ts b/examples/size_tree.ts
new file mode 100644
index 000000000..136bcec7a
--- /dev/null
+++ b/examples/size_tree.ts
@@ -0,0 +1,6 @@
+import { sizeTree } from "capi/patterns/sizeTree.ts"
+import { chain } from "polkadot_dev/mod.ts"
+
+const result = await sizeTree(chain).run()
+
+console.log(result)
diff --git a/fluent/StorageRune.ts b/fluent/StorageRune.ts
index 8048e2b49..abeee1bf4 100644
--- a/fluent/StorageRune.ts
+++ b/fluent/StorageRune.ts
@@ -19,6 +19,17 @@ export class StorageRune<in out K extends unknown[], out V, out U> extends Rune<
     this.$value = this.pallet.metadata.codec(this.into(ValueRune).access("value"))
   }
 
+  size<X>(...[partialKey, blockHash]: RunicArgs<X, [partialKey?: unknown[], blockHash?: HexHash]>) {
+    return this.pallet.metadata.chain.connection
+      .call(
+        "state_getStorageSize",
+        this.$key.encoded(Rune.resolve(partialKey).map((x) => x ?? [])).map(hex.encode),
+        blockHash,
+      )
+      .unhandle(null)
+      .rehandle(null, () => Rune.constant(undefined))
+  }
+
   entryPageRaw<X>(
     ...[count, partialKey, start, blockHash]: RunicArgs<X, [
       count: number,
diff --git a/patterns/sizeTree.ts b/patterns/sizeTree.ts
new file mode 100644
index 000000000..e3ddbfc70
--- /dev/null
+++ b/patterns/sizeTree.ts
@@ -0,0 +1,23 @@
+import { Chain, ChainRune, HexHash, MetaRune, Rune, RunicArgs, ValueRune } from "../mod.ts"
+
+export function sizeTree<U, C extends Chain, X>(
+  chain: ChainRune<U, C>,
+  ...[blockHash]: RunicArgs<X, [blockHash?: HexHash]>
+) {
+  const metadata = chain.metadata(blockHash)
+  return metadata
+    .into(ValueRune)
+    .map(({ pallets }) =>
+      Rune.rec(Object.fromEntries(pallets.map((pallet) => [
+        pallet.name,
+        Rune.rec(Object.fromEntries(
+          pallet.storage?.entries.map((entry) => [
+            entry.name,
+            metadata.pallet(pallet.name).storage(entry.name).size([], blockHash),
+          ]) || [],
+        )),
+      ])))
+    )
+    .into(MetaRune)
+    .flat()
+}