Make RuntimeModules Injectable

packages/library/src/hooks/TransactionFeeHook.ts

@@ -51,7 +51,8 @@ const errors = {
 export class TransactionFeeHook extends ProvableTransactionHook<TransactionFeeHookConfig> {
   public constructor(
     // dependency on runtime, since balances are part of runtime logic
-    @inject("Runtime") public runtime: Runtime<RuntimeModulesRecord>
+    @inject("Runtime") public runtime: Runtime<RuntimeModulesRecord>,
+    @inject("Balances") public balances: Balances
   ) {
     super();
   }
@@ -91,10 +92,6 @@ export class TransactionFeeHook extends ProvableTransactionHook<TransactionFeeHo
     super.config = value;
   }
 
-  public get balances() {
-    return this.runtime.dependencyContainer.resolve<Balances>("Balances");
-  }
-
   public get feeAnalyzer() {
     if (this.persistedFeeAnalyzer === undefined) {
       throw new Error("TransactionFeeHook.start not called by protocol");

packages/protocol/src/protocol/Protocol.ts

@@ -187,6 +187,28 @@ export class Protocol<
         );
       }
     });
+
+    // Cross-register all runtime modules to the protocol container for easier
+    // access of runtime modules inside protocol hooks
+    if (this.container.isRegistered("Runtime", true)) {
+      const runtimeContainer: ModuleContainer<any> =
+        this.container.resolve("Runtime");
+
+      runtimeContainer.moduleNames.forEach((runtimeModuleName) => {
+        this.container.register(runtimeModuleName, {
+          useFactory: (dependencyContainer) => {
+            // Prevents creation of closure
+            const runtime: ModuleContainer<any> =
+              dependencyContainer.resolve("Runtime");
+            return runtime.resolve(runtimeModuleName);
+          },
+        });
+      });
+    } else {
+      log.warn(
+        "Couldn't resolve Runtime reference in Protocol, resolving RuntimeModules in hooks won't be available"
+      );
+    }
   }
 
   public async start() {

packages/protocol/test/TestingProtocol.ts

@@ -1,21 +1,17 @@
-import { WithZkProgrammable, ZkProgrammable } from "@proto-kit/common";
 import { container } from "tsyringe";
+import { Runtime } from "@proto-kit/module";
+import { Balance } from "@proto-kit/sequencer/test/integration/mocks/Balance";
+import { NoopRuntime } from "@proto-kit/sequencer/test/integration/mocks/NoopRuntime";
 
 import {
   AccountStateHook,
   BlockHeightHook,
   BlockProver,
   LastStateRootBlockHook,
-  MethodPublicOutput,
   Protocol,
   StateTransitionProver,
 } from "../src";
 
-class RuntimeMock implements WithZkProgrammable<undefined, MethodPublicOutput> {
-  zkProgrammable: ZkProgrammable<undefined, MethodPublicOutput> =
-    undefined as unknown as ZkProgrammable<undefined, MethodPublicOutput>;
-}
-
 export function createAndInitTestingProtocol() {
   const ProtocolClass = Protocol.from({
     modules: {
@@ -35,11 +31,22 @@ export function createAndInitTestingProtocol() {
     StateTransitionProver: {},
     LastStateRoot: {},
   });
-  protocol.create(() => container.createChildContainer());
 
-  protocol.registerValue({
-    Runtime: new RuntimeMock(),
+  const appChain = container.createChildContainer();
+
+  appChain.register("Runtime", {
+    useClass: Runtime.from({
+      modules: {
+        Balance,
+        NoopRuntime,
+      },
+      config: {
+        Balance: {},
+        NoopRuntime: {},
+      },
+    }),
   });
+  protocol.create(() => appChain.createChildContainer());
 
   return protocol;
 }

packages/sdk/test/modularization.test.ts

@@ -2,7 +2,10 @@ import "reflect-metadata";
 import { MethodIdResolver, Runtime, RuntimeModule } from "@proto-kit/module";
 import { ChildContainerProvider } from "@proto-kit/common";
 import { Protocol, ProtocolModule } from "@proto-kit/protocol";
-import { VanillaProtocolModules } from "@proto-kit/library";
+import {
+  VanillaProtocolModules,
+  VanillaRuntimeModules,
+} from "@proto-kit/library";
 import { Sequencer, SequencerModule } from "@proto-kit/sequencer";
 import { PrivateKey } from "o1js";
 
@@ -13,7 +16,6 @@ class TestRuntimeModule extends RuntimeModule<object> {
 
   public create(childContainerProvider: ChildContainerProvider) {
     super.create(childContainerProvider);
-
     // Just to test if it doesn't throw
     childContainerProvider();
 
@@ -47,9 +49,9 @@ describe("modularization", () => {
   it("should initialize all modules correctly", async () => {
     const appChain = AppChain.from({
       Runtime: Runtime.from({
-        modules: {
+        modules: VanillaRuntimeModules.with({
           TestRuntimeModule,
-        },
+        }),
       }),
       Protocol: Protocol.from({
         modules: VanillaProtocolModules.with({
@@ -66,6 +68,7 @@ describe("modularization", () => {
 
     appChain.configurePartial({
       Runtime: {
+        Balances: {},
         TestRuntimeModule: {},
       },
       Protocol: {