Dynamic state access

packages/common/src/utils.ts

@@ -295,3 +295,21 @@ export function assertDefined<T>(
     throw new Error(msg ?? "Value is undefined");
   }
 }
+
+export function compareStrings(a: string, b: string): number {
+  // eslint-disable-next-line no-nested-ternary
+  return a === b ? 0 : a > b ? 1 : -1;
+}
+
+export function recordByKey<Object>(
+  array: Object[],
+  keyMapper: (obj: Object) => string
+): Record<string, Object> {
+  return Object.fromEntries(
+    array.map((obj) => {
+      const key = keyMapper(obj);
+      const copy = { ...obj };
+      return [key, copy];
+    })
+  );
+}

packages/library/src/hooks/RuntimeFeeAnalyzerService.ts

@@ -2,6 +2,7 @@ import {
   ConfigurableModule,
   createMerkleTree,
   InMemoryMerkleTreeStorage,
+  recordByKey,
 } from "@proto-kit/common";
 import { Runtime, RuntimeModulesRecord } from "@proto-kit/module";
 import { container, inject } from "tsyringe";
@@ -10,6 +11,7 @@ import {
   RuntimeTransaction,
   NetworkState,
 } from "@proto-kit/protocol";
+import { RuntimeAnalyzerService } from "@proto-kit/sequencer";
 import { Field, Poseidon, Struct } from "o1js";
 
 import { UInt64 } from "../math/UInt64";
@@ -65,6 +67,7 @@ export class RuntimeFeeAnalyzerService extends ConfigurableModule<RuntimeFeeAnal
   };
 
   public constructor(
+    public runtimeAnalyzerService: RuntimeAnalyzerService,
     @inject("Runtime") public runtime: Runtime<RuntimeModulesRecord>
   ) {
     super();
@@ -81,86 +84,61 @@ export class RuntimeFeeAnalyzerService extends ConfigurableModule<RuntimeFeeAnal
     });
 
     container.resolve(RuntimeMethodExecutionContext).clear();
-    let methodCounter = 0;
-    const [values, indexes] =
-      await this.runtime.zkProgrammable.zkProgram.reduce<
-        Promise<[FeeTreeValues, FeeIndexes]>
-      >(
-        async (accum, program) => {
-          const [valuesProg, indexesProg] = await accum;
-          const analyzedMethods = await program.analyzeMethods();
-          const [valuesMeth, indexesMeth] = Object.keys(program.methods).reduce<
-            [FeeTreeValues, FeeIndexes]
-          >(
-            // eslint-disable-next-line @typescript-eslint/no-shadow
-            ([values, indexes], combinedMethodName) => {
-              const { rows } = analyzedMethods[combinedMethodName];
-              // const rows = 1000;
-              const [moduleName, methodName] = combinedMethodName.split(".");
-              const methodId = this.runtime.methodIdResolver.getMethodId(
-                moduleName,
-                methodName
-              );
-
-              /**
-               * Determine the fee config for the given method id, and merge it with
-               * the default fee config.
-               */
-              return [
-                {
-                  ...values,
-
-                  [methodId.toString()]: {
-                    methodId,
-
-                    baseFee:
-                      this.config.methods[combinedMethodName]?.baseFee ??
-                      this.config.baseFee,
-
-                    perWeightUnitFee:
-                      this.config.methods[combinedMethodName]
-                        ?.perWeightUnitFee ?? this.config.perWeightUnitFee,
-
-                    weight:
-                      this.config.methods[combinedMethodName]?.weight ??
-                      BigInt(rows),
-                  },
-                },
-                {
-                  ...indexes,
-                  // eslint-disable-next-line no-plusplus
-                  [methodId.toString()]: BigInt(methodCounter++),
-                },
-              ];
-            },
-            [{}, {}]
-          );
-          return [
-            { ...valuesProg, ...valuesMeth },
-            { ...indexesProg, ...indexesMeth },
-          ];
-        },
-        Promise.resolve([{}, {}])
-      );
 
+    const runtimeInfo = await this.runtimeAnalyzerService.getRuntimeInfo();
+
+    const values = Object.entries(runtimeInfo).map(
+      ([combinedMethodName, { rows }]) => {
+        // const rows = 1000;
+        const [moduleName, methodName] = combinedMethodName.split(".");
+        const methodId = this.runtime.methodIdResolver.getMethodId(
+          moduleName,
+          methodName
+        );
+
+        /**
+         * Determine the fee config for the given method id, and merge it with
+         * the default fee config.
+         */
+        return {
+          methodId,
+
+          baseFee:
+            this.config.methods[combinedMethodName]?.baseFee ??
+            this.config.baseFee,
+
+          perWeightUnitFee:
+            this.config.methods[combinedMethodName]?.perWeightUnitFee ??
+            this.config.perWeightUnitFee,
+
+          weight:
+            this.config.methods[combinedMethodName]?.weight ?? BigInt(rows),
+        };
+      }
+    );
     const tree = new FeeTree(new InMemoryMerkleTreeStorage());
 
-    Object.values(values).forEach((value, index) => {
+    const indexes = values.map((value, index) => {
       const feeConfig = new MethodFeeConfigData({
         methodId: Field(value.methodId),
         baseFee: UInt64.from(value.baseFee),
         weight: UInt64.from(value.weight),
         perWeightUnitFee: UInt64.from(value.perWeightUnitFee),
       });
       tree.setLeaf(BigInt(index), feeConfig.hash());
+      return [value.methodId.toString(), BigInt(index)] as const;
     });
 
-    this.persistedFeeTree = { tree, values, indexes };
+    this.persistedFeeTree = {
+      tree,
+      values: recordByKey(values, (v) => v.methodId.toString()),
+      indexes: Object.fromEntries<bigint>(indexes),
+    };
   }
 
   public getFeeTree() {
     if (this.persistedFeeTree === undefined) {
-      throw new Error("Fee Tree not intialized");
+      throw new Error("Fee Tree not initialized");
     }
 
     return this.persistedFeeTree;

packages/library/src/hooks/TransactionFeeHook.ts

@@ -12,6 +12,7 @@ import {
 } from "@proto-kit/protocol";
 import { Field, Provable, PublicKey } from "o1js";
 import { noop } from "@proto-kit/common";
+import { RuntimeAnalyzerService } from "@proto-kit/sequencer";
 
 import { UInt64 } from "../math/UInt64";
 import { Balance, TokenId } from "../runtime/Balances";
@@ -53,7 +54,9 @@ export class TransactionFeeHook extends ProvableTransactionHook<TransactionFeeHo
   public constructor(
     // dependency on runtime, since balances are part of runtime logic
     @inject("Runtime") public runtime: Runtime<RuntimeModulesRecord>,
-    @inject("Balances") public balances: Balances
+    @inject("Balances") public balances: Balances,
+    // TODO Check that the container is the right one here
+    public runtimeAnalyzerService: RuntimeAnalyzerService
   ) {
     super();
   }
@@ -79,7 +82,10 @@ export class TransactionFeeHook extends ProvableTransactionHook<TransactionFeeHo
   }
 
   public async start() {
-    this.persistedFeeAnalyzer = new RuntimeFeeAnalyzerService(this.runtime);
+    this.persistedFeeAnalyzer = new RuntimeFeeAnalyzerService(
+      this.runtimeAnalyzerService,
+      this.runtime
+    );
     this.verifyConfig();
     this.persistedFeeAnalyzer.config = this.config;
     await this.persistedFeeAnalyzer.initializeFeeTree();

packages/module/src/index.ts

@@ -7,3 +7,4 @@ export * from "./state/InMemoryStateService";
 export * from "./method/MethodParameterEncoder";
 export * from "./runtime/MethodIdResolver";
 export * from "./factories/MethodIdFactory";
+export * from "./testing/TestingRuntime";