Don't emit STs when accessing STs outside provable code

packages/protocol/src/state/State.ts

@@ -9,6 +9,7 @@ import { StateTransition } from "../model/StateTransition";
 
 import { StateServiceProvider } from "./StateServiceProvider";
 import { RuntimeMethodExecutionContext } from "./context/RuntimeMethodExecutionContext";
+import { WitnessBlockContext } from "./WitnessBlockContext";
 
 export class WithPath {
   public path?: Field;
@@ -137,11 +138,17 @@ export class State<Value> extends Mixin(WithPath, WithStateServiceProvider) {
 
     this.hasPathOrFail();
 
-    const stateTransition = StateTransition.from(this.path, option);
+    const { isInWitnessBlock } = container.resolve(WitnessBlockContext);
 
-    container
-      .resolve(RuntimeMethodExecutionContext)
-      .addStateTransition(stateTransition);
+    // If we're inside a witness block, we only want to retrieve the state
+    // to use as a witness but not emit an ST
+    if (!isInWitnessBlock) {
+      const stateTransition = StateTransition.from(this.path, option);
+
+      container
+        .resolve(RuntimeMethodExecutionContext)
+        .addStateTransition(stateTransition);
+    }
 
     return option;
   }
@@ -170,6 +177,12 @@ export class State<Value> extends Mixin(WithPath, WithStateServiceProvider) {
       toOption
     );
 
+    const { isInWitnessBlock } = container.resolve(WitnessBlockContext);
+
+    if (isInWitnessBlock) {
+      throw new Error("Cannot set state inside of provable block.");
+    }
+
     container
       .resolve(RuntimeMethodExecutionContext)
       .addStateTransition(stateTransition);

packages/protocol/src/state/WitnessBlockContext.ts

@@ -0,0 +1,50 @@
+import { container, singleton } from "tsyringe";
+import { Provable } from "o1js";
+
+@singleton()
+export class WitnessBlockContext {
+  public witnessBlockDepth: number = 0;
+
+  public get isInWitnessBlock() {
+    return this.witnessBlockDepth > 0;
+  }
+}
+
+const asyncProxyWitnessFunction = <
+  Ret,
+  F extends (...args: any[]) => Promise<Ret>,
+>(
+  originalFuncDef: F
+) => {
+  return async (...args: Parameters<F>) => {
+    const context = container.resolve(WitnessBlockContext);
+    context.witnessBlockDepth += 1;
+    const ret = await originalFuncDef(...args);
+    context.witnessBlockDepth -= 1;
+    return ret;
+  };
+};
+
+const proxySyncWitnessFunction = <
+  Params extends any[],
+  Ret,
+  F extends (...args: Params) => Ret,
+>(
+  originalFuncDef: F
+) => {
+  return (...args: Params): Ret => {
+    const context = container.resolve(WitnessBlockContext);
+    context.witnessBlockDepth += 1;
+    const ret = originalFuncDef(...args);
+    context.witnessBlockDepth -= 1;
+    return ret;
+  };
+};
+
+Provable.witnessAsync = asyncProxyWitnessFunction(Provable.witnessAsync);
+
+Provable.witness = proxySyncWitnessFunction(Provable.witness);
+
+Provable.witnessFields = proxySyncWitnessFunction(Provable.witnessFields);
+
+Provable.asProver = proxySyncWitnessFunction(Provable.asProver);

packages/sdk/test/stprover-emit-sts.test.ts

@@ -0,0 +1,130 @@
+import "reflect-metadata";
+
+import { UInt64 } from "@proto-kit/library";
+import { runtimeMethod, runtimeModule, RuntimeModule } from "@proto-kit/module";
+import { Field, PrivateKey, Provable } from "o1js";
+import {
+  RuntimeMethodExecutionContext,
+  State,
+  state,
+} from "@proto-kit/protocol";
+import { container } from "tsyringe";
+
+import { TestingAppChain } from "../src";
+
+@runtimeModule()
+class StateTester extends RuntimeModule<unknown> {
+  @state() public state1 = State.from<UInt64>(UInt64);
+
+  @runtimeMethod()
+  public async setFail() {
+    await Provable.witnessAsync(Field, async () => {
+      await this.state1.set(UInt64.from(10));
+      return Field(0);
+    });
+  }
+
+  @runtimeMethod()
+  public async setPass() {
+    await this.state1.set(UInt64.from(10));
+  }
+
+  @runtimeMethod()
+  public async getSTs() {
+    await this.state1.get();
+  }
+
+  @runtimeMethod()
+  public async getNoSTs() {
+    await Provable.witnessAsync(Field, async () => {
+      const stateReturned = await this.state1.get();
+      return Field.from(stateReturned.value.toBigInt());
+    });
+  }
+}
+
+describe("StateTransition", () => {
+  const senderKey = PrivateKey.random();
+
+  const appChain = TestingAppChain.fromRuntime({
+    StateTester,
+  });
+
+  beforeEach(async () => {
+    appChain.configurePartial({
+      Runtime: {
+        StateTester: {},
+        Balances: {},
+      },
+
+      Protocol: {
+        ...appChain.config.Protocol!,
+      },
+    });
+
+    await appChain.start();
+    appChain.setSigner(senderKey);
+  });
+
+  afterEach(async () => {
+    await appChain.close();
+  });
+
+  it("should emit no sts for get", async () => {
+    const stateTester = appChain.runtime.resolve("StateTester");
+    const context = container.resolve(RuntimeMethodExecutionContext);
+
+    // We set the state so when we fetch it it won't error.
+    const tx0 = await appChain.transaction(
+      senderKey.toPublicKey(),
+      async () => {
+        await stateTester.setPass();
+      }
+    );
+    await tx0.sign();
+    await tx0.send();
+    await appChain.produceBlock();
+
+    const tx1 = await appChain.transaction(
+      senderKey.toPublicKey(),
+      async () => {
+        await stateTester.getSTs();
+      }
+    );
+    await tx1.sign();
+    await tx1.send();
+    const STs = context.current().result.stateTransitions;
+
+    expect(STs.length).not.toBe(0);
+
+    const tx2 = await appChain.transaction(
+      senderKey.toPublicKey(),
+      async () => {
+        await stateTester.getNoSTs();
+      }
+    );
+    await tx2.sign();
+    await tx2.send();
+    const STs2 = context.current().result.stateTransitions;
+    expect(STs2.length).toBe(0);
+  });
+
+  it("should fail outside provable code for set", async () => {
+    const stateTester = appChain.runtime.resolve("StateTester");
+    const tx1 = await appChain.transaction(
+      senderKey.toPublicKey(),
+      async () => {
+        await stateTester.setPass();
+      }
+    );
+    await tx1.sign();
+    await tx1.send();
+    await appChain.produceBlock();
+
+    await expect(() =>
+      appChain.transaction(senderKey.toPublicKey(), async () => {
+        await stateTester.setFail();
+      })
+    ).rejects.toThrow(new Error("Cannot set state inside of provable block."));
+  });
+});