Refactor offchain state #1834

45930 · 2024-09-24T20:54:52Z

Summary

The experimental offchain state API currently uses a singleton class to manage storage on behalf of a smart contract. This blocks the use case where many instances of the same smart contract need to be instantiated at the same time, because they will all use the same offchain state under the hood, leading to data mismatches.

This PR splits the offchain state API into OffchainState and OffchainStateInstance such that the OffchainState can be compiled into a circuit definition, and the OffchainStateInstance stores the internal data like which contract instance it is associated with and the merkle trees of data.

Note

I called it out in line, but please note that there is still an ongoing DevX problem where the smart contract class definition has to reference a specific offchain state instance in order to compile. That leaves us in the same problem, where unless handled by the developer, multiple contracts will point to the same storage. There is now a viable workaround, but it doesn't work as nicely as we'd like. I'd appreciate any input to overcome this issue.

Closes: #1832

45930 · 2024-10-01T02:32:25Z

src/lib/mina/actions/offchain-contract.unit-test.ts

+    // We also can't rely on different instances of the contract having different offchain state instances, because this specific instance of
+    // offchain state is referenced by the class definition of the smart contract.
+    // By using a closure and exporting the offchain state and the contract together, we can always refer to the correct instance of the offchain state
+    offchainState = offchainStateInstance;


This line is the source of a lot of frustration. Here are some alternatives attempted:

offchainState = offchainState.init()

Re-initialized the state every time the prover is invoked. The smart contract is not capable of staying up to date with the contract state.

OffchainState#memoizedInstances

Attempted to store previously initialized instances in a map in global state to get around the re-init problem, but ran into issues trying to read variables in provable code. E.g. if we use the contract address as a key, the prover blows up.

offchainState: typeof offchainStateInstance

The contract can't compile because provable methods rely on offchainState being defined at compile time

That leaves the implementation as implemented here. The instance offchainStateInstance is tied to the class definition of the smart contract. That means 2 or more smart contract instances will all have the same offchain state instance. Even if we reset the instance state after initialization, the prover will re-run this line and set the offchain state back to the orifinally-defined value.

I found the best/only way to manage this is to turn the whole smart contract class into a closure so that the definition of each smart contract can be paired with its offchain state instance. This doesn't feel amazing, but it does unblock the use case.

45930 · 2024-10-01T02:34:54Z

src/lib/mina/actions/offchain-contract.unit-test.ts

+  ExampleContract: ExampleContractB,
+} = ExampleContractWithState();
+
+const Local = await Mina.LocalBlockchain({ proofsEnabled: true });


testLocal accepts a single contract instance only, so I re-implemented the behavior here for the specific case I needed.

45930 · 2024-10-01T02:36:14Z

src/lib/mina/actions/offchain-contract.unit-test.ts

+
+console.time('compile contract');
+await ExampleContractA.compile();
+await ExampleContractB.compile();


Contracts have the same _verificationKey.hash, but _provers need to be recompiled, or else all the contracts will point to the same instance of offchain state again.

kadirchan · 2024-10-01T11:45:20Z

Love it

Trivo25 · 2024-10-01T12:51:29Z

Given that these modifications affect both the public API and developer usage patterns, we should evaluate options for maintaining backwards compatibility. If backwards compatibility isn't feasible, we may have to consider targeting V2 for these changes

45930 · 2024-10-01T13:36:24Z

@Trivo25 this is in the experimental namespace, so do we need to be worried about compatibility?

I do think an accompanying PR in the docs is required to match this change before release.

Trivo25 · 2024-10-01T13:43:21Z

@Trivo25 this is in the experimental namespace, so do we need to be worried about compatibility?

That's a good point, then probably not

mitschabaude · 2024-10-01T15:06:10Z

Attempted to store previously initialized instances in a map in global state to get around the re-init problem, but ran into issues trying to read variables in provable code. E.g. if we use the contract address as a key, the prover blows up.

@45930 sounds like that could be solved with Provable.asProver()

45930 · 2024-10-01T15:44:24Z

sounds like that could be solved with Provable.asProver()

Ok, I think I have a working branch locally that works as expected with the memoization happening "behind the scenes". I don't love that this is hidden from the developer, but I like it better than forcing them to use a closure. WDYT?

// Closure ensures that different contract states never interfere
function ExampleContractWithState() {
  const offchainStateInstance = offchainState.init();
  class ExampleContract extends SmartContract { ]
  
  return { offchainStateInstance, ExampleContract }
}

OR

class ExampleContract extends SmartContract {
  @state(OffchainState.Commitments) offchainStateCommitments =
    offchainState.emptyCommitments();
  
  // init method handles the instance management and doesn't expose it to the user at all
  offchainState = offchainState.init(this);
}

45930 · 2024-10-01T16:15:42Z

@mitschabaude , bf0e09e is pretty clean!

I think this API is a lot less error-prone and requires less deep knowledge to use. But it does add more secret global state, which tends to cause its own issues down the line.

mitschabaude · 2024-10-02T06:15:47Z

src/lib/mina/actions/offchain-contract.unit-test.ts

+  init(): void {
+    super.init();
+    this.offchainState.setContractInstance(this);
+  }


in general, this will be the wrong place for connecting the offchain state. init() is only called on initial deployment. can't we move this into offchainState.init()?

@Trivo25

can't we move this into offchainState.init()?

No, we can't.

We can just remove the init method here to make it more clear that offchainState.setContractInstance needs to be called outside of the circuit. The way it works in my example is:

// first time deploying the contract const c = Contract(address); Mina.transaction({ c.deploy }); // contract instance is set automatically

// contract is already deployed const c = Contract(address); c.offchainState.setActiveInstance(c); // set contract instance manually Mina.transaction({ c.method(...params) });

Again, this is a unit test, so we can't stop a developer from using it like I have here. I'm open to the idea that it's confusing and we should focus on demonstrating one canonical way to set the contract instance.

fair enough, just keep in mind that most devs will use our unit tests as example

mitschabaude · 2024-10-02T06:21:53Z

src/lib/mina/actions/offchain-state.ts

+        instance.setContractClass(
+          contractInstance.constructor as OffchainStateContractClass<Config>
+        );
+        memoizedInstances.set(key, instance);
+      }


any reason not to do instance.setContractInstance(contractInstance) here?

That was my first instinct, but it continues to cause provable errors:

Error: variables only exist inside checked computations This shouldn't have happened and indicates an internal bug. at Bug (file:///Users/coby/Projects/o1js/dist/node/lib/util/errors.js:226:12) at assert (file:///Users/coby/Projects/o1js/dist/node/lib/util/errors.js:233:15) at toConstantField (file:///Users/coby/Projects/o1js/dist/node/lib/provable/field.js:1083:5) at toConstant (file:///Users/coby/Projects/o1js/dist/node/lib/provable/field.js:1076:12) at Field.toConstant (file:///Users/coby/Projects/o1js/dist/node/lib/provable/field.js:118:16) at Object.fromPublicKey (file:///Users/coby/Projects/o1js/dist/node/lib/ml/conversion.js:64:24) at Object.getAccount (file:///Users/coby/Projects/o1js/dist/node/lib/mina/local-blockchain.js:77:52) at Module.getAccount (file:///Users/coby/Projects/o1js/dist/node/lib/mina/mina-instance.js:43:27) at Object.fetch (file:///Users/coby/Projects/o1js/dist/node/lib/mina/state.js:223:32) at onchainActionState (file:///Users/coby/Projects/o1js/dist/node/lib/mina/actions/offchain-state.js:77:81) at merkleMaps (file:///Users/coby/Projects/o1js/dist/node/lib/mina/actions/offchain-state.js:86:37) at Object.createSettlementProof (file:///Users/coby/Projects/o1js/dist/node/lib/mina/actions/offchain-state.js:222:43) at file:///Users/coby/Projects/o1js/dist/node/lib/mina/actions/offchain-contract.unit-test.js:114:50 at runInstruction (file:///Users/coby/Projects/o1js/dist/node/lib/mina/test/test-contract.js:101:27) at execute (file:///Users/coby/Projects/o1js/dist/node/lib/mina/test/test-contract.js:77:19) at process.processTicksAndRejections (node:internal/process/task_queues:95:5) at async testLocal (file:///Users/coby/Projects/o1js/dist/node/lib/mina/test/test-contract.js:93:9) at async file:///Users/coby/Projects/o1js/dist/node/lib/mina/actions/offchain-contract.unit-test.js:104:1

This error throws after compiling and deploying successfully, when I try to call createAccount in the first transaction to the zkapp after deploying.

I landed on the init method because it's the only place within the zkapp itself that will only be executed when I want it to be executed.

Setting the contract instance outside of the zkapp is still possible, e.g.

// already-deployed contract const zkapp = Contract(deployedAddress); zkapp.offchainState.setContractInstance(zkapp);

Can you please export IndexedMerkleMapBase as part of this refactoring? It is needed for the serialization of IndexedMerkleMap of any height. Now, only the function IndexedMerkleMap is exported, and I need to get this type using syntax like

export function serializeIndexedMap( map: InstanceType<ReturnType<typeof IndexedMerkleMap>> )

https://github.com/zkcloudworker/zkcloudworker-tests/blob/main/src/indexed-map.ts

45930 · 2024-10-07T14:23:26Z

@dfstio do you want to open a separate PR to expose that API? This PR may take some more time to fully review, but something simpler may get through faster on its own.

dfstio · 2024-10-07T14:39:00Z

@dfstio do you want to open a separate PR to expose that API? This PR may take some more time to fully review, but something simpler may get through faster on its own.

Opened issue: #1855
If it will have positive feedback, I can create a PR.

boray · 2024-10-07T17:13:57Z

src/lib/mina/actions/offchain-state.ts

-      return getContract();
-    } catch {
-      return internal.contract;
+    type InitializedInternalState = {


This type seems to be unused. Do you plan to use it?

No, it can be removed, I'll do that in my next commit

45930 · 2024-10-09T19:21:30Z

@Trivo25 just to add some visual context and reasoning behind this change:

This diagram shows how the feature works today. The instance of the offchain state, and the definition of the smart contract are directly linked. They need to be linked somehow because offchain state described a zk circuit which is a dependency of the smart contract. Unfortunately, this way doesn't work because a 2nd contract initiialized in the same js session will reference the same offchain state, and when the states differ, you won't be able to generate a valid proof for at least one of them.

My changes instantiate a separate instance of the offchain state for each smart contract, but they each reference the same one zk circuit described by the offchain state definition. This leaves 2 "fake classes": OffchainState which has the methods compile, emptyCommitments and init, which compile the circuit, return the commitments of an empty copy of the state, and intantiate a new copy of the state respectively. The other is OffchainStateInstance which has all the methods and properties related to an actual copy of the offchain state, e.g. settle, and commitments.

45930 · 2024-10-09T19:22:19Z

Trying to find the right API that a) works and b) is simple to use has been the challenging piece, and the last thing Gregor and I were haggling over. What makes it tricky is that when a prover generates a proof, it re-runs the smart contract definition, so if the definition of the smart contract includes an empty state to start with, that's what will be used at proof time. If the definition of the contract doesn't reference offchain state at all to avoid this problem, then we can't use the data in provable methods.

2 plausible patterns are:

closure

function MyContract() {
  const offchainState = OffchainState.init();

  class MyContract {
    // ... reference offchainState here
    }
    
    return {
        state: offchainState,
        contract: MyContract
    }

Magic global state

// in offchain state definition
// manages a memoization map behind the scenes so that a new instance is only initialized the first time

  init(contractInstance: OffchainStateContract<Config>) {
    let key = 'COMPILE_TIME';
    let contractAddress = contractInstance.address;
    if (contractAddress.isConstant()) {
      key = contractAddress.toBase58();
    } else {
      Provable.asProver(() => {
        key = contractAddress.toBase58();
      });
    }
    let instance = memoizedInstances.get(key);
    if (instance === undefined) {
      instance = OffchainStateInstance();
      instance.setContractClass(
        contractInstance.constructor as OffchainStateContractClass<Config>
      );
      memoizedInstances.set(key, instance);
    }
    return instance;
}

// in contract file

class MyContract {
  // ... reference offchainState here
  
  offchainState = OffchainState.init(this);
}

Currently, option 2 is implemented, which I think is ok. All things equal, I don't like the magic global state, but I don't have an alternative that I like better.

Trivo25 · 2024-10-16T14:15:35Z

src/lib/mina/actions/offchain-contract.unit-test.ts

+assert((await contractA.getBalance(sender)).toBigInt() == 900n);
+assert((await contractB.getBalance(sender)).toBigInt() == 1300n);
+assert((await contractA.getBalance(receiver)).toBigInt() == 100n);
+assert((await contractB.getBalance(receiver)).toBigInt() == 200n);


imo it would be good to add a test where the state of both contracts actually diverge more aggressively (eg contract A doesn't settle for 10 actions, whereas contract B settles every action)

Ok, I will add this kind of test, rm the extra type I added, noted in Boray's comment, and add to the changelog. Anything else you want updated?

I need to get a bit familiar with the changes first, will do another pass

kadirchan · 2024-10-16T17:08:41Z

wen release 👀

boray · 2024-10-16T17:38:25Z

wen release 👀

It will be included in the next release after it's merged but you can use before that with npm i https://pkg.pr.new/o1-labs/o1js@<FIRST 7 DIGITS OF COMMIT HASH>. Commit hash should be the hash of merge commit in main branch.

45930 · 2024-10-16T19:24:41Z

src/lib/mina/actions/offchain-contract-tests/ExampleContract.ts

+    offchainState.emptyCommitments();
+
+  // o1js memoizes the offchain state by contract address so that this pattern works
+  offchainState: any = offchainState.init(this);


any is required here because the type of offchainState is actually

OffchainStateInstance<{ readonly accounts: { kind: "offchain-map"; keyType: typeof PublicKey; valueType: typeof UInt64; }; readonly totalSupply: { kind: "offchain-field"; type: typeof UInt64; }; }>

But that type is not exported, resulting in this error: Public property 'offchainState' of exported class has or is using name 'OffchainStateInstance' from external module "/Users/coby/Projects/o1js/src/lib/mina/actions/offchain-state" but cannot be named.

Seems fine to me to handle in a different PR eventually, as this API remains experimental.

as a rule of thumb: If we can do it now, might as well do it now :D

woah doesn't omitting this export completely break the type-safety of interacting with offchain state then?

Turns out to be a pretty simple fix

#1871

45930 · 2024-10-16T19:26:50Z

src/lib/mina/actions/offchain-contract-tests/multi-contract-instance.unit-test.ts

+import { ExampleContract } from './ExampleContract.js';
+import { settle, transfer } from './utils.js';
+
+const Local = await Mina.LocalBlockchain({ proofsEnabled: false });


proofs not enabled on the multi-contract test since it takes much longer. This also exercises both the proofs-enabled and proofs-disabled paths.

45930 · 2024-10-16T19:28:14Z

src/lib/mina/actions/offchain-contract-tests/single-contract-instance.unit-test.ts

+} from './ExampleContract.js';
+import { settle, transfer } from './utils.js';
+
+const Local = await Mina.LocalBlockchain({ proofsEnabled: true });


Proofs enabled on the shorter test to make sure the artifacts compile correctly and proofs can generate.

Trivo25 · 2024-10-17T13:43:06Z

src/lib/mina/actions/offchain-contract-tests/ExampleContract.ts

+    offchainState.emptyCommitments();
+
+  // o1js memoizes the offchain state by contract address so that this pattern works
+  offchainState: any = offchainState.init(this);


as a rule of thumb: If we can do it now, might as well do it now :D

Trivo25 · 2024-10-17T13:44:59Z

src/lib/mina/actions/offchain-contract-tests/ExampleContract.ts

@@ -0,0 +1,89 @@
+import {


I am okay with leaving it here for now but next time we add larger tests we should add them to the example folder and run them from there - makes it a bit easier to keep the internal code clean yet test it and provide a good example to the community

I'd be happy to follow convention, but it doesn't seem like there is one. This unit-test file was already here and I added onto it. It seemed like it was getting too big so I broke it out into different files. I'll follow up with you about how you see o1js being organized and we can work towards that.

I agree, definitely not your fault - just something I noticed with the larger test

mitschabaude · 2024-10-28T06:21:16Z

src/lib/mina/actions/offchain-contract.unit-test.ts

We should bring back the testLocal() test IMO, it was less verbose and had higher coverage of the implementation than the current proofsEnabled test

45930 added 7 commits September 24, 2024 16:54

first pass

c142de6

progress

89f774a

error in transfer function

45a0fca

sloppy version working

b106b5d

rm log lines, add multi-instance test

d0eeee2

better names and descriptions

a445be2

small cleanups

1feb7b4

45930 commented Oct 1, 2024

View reviewed changes

Merge branch 'main' into 2024-09-refactor-offchain-state

a6ad3a5

45930 marked this pull request as ready for review October 1, 2024 03:09

45930 requested review from mitschabaude, Trivo25 and boray October 1, 2024 03:09

demo the cleaner API

bf0e09e

mitschabaude reviewed Oct 2, 2024

View reviewed changes

boray reviewed Oct 7, 2024

View reviewed changes

Trivo25 reviewed Oct 16, 2024

View reviewed changes

Trivo25 approved these changes Oct 16, 2024

View reviewed changes

Merge branch 'main' into 2024-09-refactor-offchain-state

57642f2

45930 added 3 commits October 16, 2024 14:16

PR comments

78a23d7

undelete changelog!

1219a26

refactor test file into directory

ef3f7f8

45930 commented Oct 16, 2024

View reviewed changes

Trivo25 reviewed Oct 17, 2024

View reviewed changes

45930 merged commit f5a3f1d into main Oct 17, 2024
26 checks passed

45930 deleted the 2024-09-refactor-offchain-state branch October 17, 2024 13:56

This was referenced Oct 17, 2024

Include the type OffchainStateInstance in the internal API #1871

Merged

Release/v1.9.1 #1880

Merged

mitschabaude mentioned this pull request Oct 28, 2024

offchainState settle() throws error when remaining action exceed maxActionsPerProof #1887

Open

mitschabaude reviewed Oct 28, 2024

View reviewed changes

Refactor offchain state #1834

Refactor offchain state #1834

Conversation

45930 commented Sep 24, 2024 • edited Loading

Summary

Note

45930 Oct 1, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kadirchan commented Oct 1, 2024

Trivo25 commented Oct 1, 2024

45930 commented Oct 1, 2024

Trivo25 commented Oct 1, 2024

mitschabaude commented Oct 1, 2024

45930 commented Oct 1, 2024

45930 commented Oct 1, 2024

mitschabaude Oct 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

45930 commented Oct 7, 2024

dfstio commented Oct 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

45930 commented Oct 9, 2024

45930 commented Oct 9, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kadirchan commented Oct 16, 2024

boray commented Oct 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mitschabaude Oct 28, 2024 • edited Loading

Choose a reason for hiding this comment

45930 commented Sep 24, 2024 •

edited

Loading

45930 Oct 1, 2024 •

edited

Loading

mitschabaude Oct 2, 2024 •

edited

Loading

dfstio commented Oct 7, 2024 •

edited

Loading

boray commented Oct 16, 2024 •

edited

Loading

mitschabaude Oct 28, 2024 •

edited

Loading