Merge pull request #1834 from o1-labs/2024-09-refactor-offchain-state

Refactor offchain state

CHANGELOG.md

@@ -17,6 +17,10 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/o1js/compare/f15293a69...HEAD)
 
+### Fixes
+
+- Decouple offchain state instances from their definitions https://github.com/o1-labs/o1js/pull/1834
+
 ## [1.9.0](https://github.com/o1-labs/o1js/compare/450943...f15293a69) - 2024-10-15
 
 ### Added

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

@@ -0,0 +1,89 @@
+import {
+  SmartContract,
+  method,
+  state,
+  PublicKey,
+  UInt64,
+  Experimental,
+} from '../../../../index.js';
+
+export { offchainState, StateProof, ExampleContract };
+
+const { OffchainState } = Experimental;
+
+const offchainState = OffchainState(
+  {
+    accounts: OffchainState.Map(PublicKey, UInt64),
+    totalSupply: OffchainState.Field(UInt64),
+  },
+  { logTotalCapacity: 10, maxActionsPerProof: 5 }
+);
+
+class StateProof extends offchainState.Proof {}
+
+// example contract that interacts with offchain state
+class ExampleContract extends SmartContract {
+  @state(OffchainState.Commitments) offchainStateCommitments =
+    offchainState.emptyCommitments();
+
+  // o1js memoizes the offchain state by contract address so that this pattern works
+  offchainState: any = offchainState.init(this);
+
+  @method
+  async createAccount(address: PublicKey, amountToMint: UInt64) {
+    // setting `from` to `undefined` means that the account must not exist yet
+    this.offchainState.fields.accounts.update(address, {
+      from: undefined,
+      to: amountToMint,
+    });
+
+    // TODO using `update()` on the total supply means that this method
+    // can only be called once every settling cycle
+    let totalSupplyOption = await this.offchainState.fields.totalSupply.get();
+    let totalSupply = totalSupplyOption.orElse(0n);
+
+    this.offchainState.fields.totalSupply.update({
+      from: totalSupplyOption,
+      to: totalSupply.add(amountToMint),
+    });
+  }
+
+  @method
+  async transfer(from: PublicKey, to: PublicKey, amount: UInt64) {
+    let fromOption = await this.offchainState.fields.accounts.get(from);
+    let fromBalance = fromOption.assertSome('sender account exists');
+
+    let toOption = await this.offchainState.fields.accounts.get(to);
+    let toBalance = toOption.orElse(0n);
+
+    /**
+     * Update both accounts atomically.
+     *
+     * This is safe, because both updates will only be accepted if both previous balances are still correct.
+     */
+    this.offchainState.fields.accounts.update(from, {
+      from: fromOption,
+      to: fromBalance.sub(amount),
+    });
+
+    this.offchainState.fields.accounts.update(to, {
+      from: toOption,
+      to: toBalance.add(amount),
+    });
+  }
+
+  @method.returns(UInt64)
+  async getSupply() {
+    return (await this.offchainState.fields.totalSupply.get()).orElse(0n);
+  }
+
+  @method.returns(UInt64)
+  async getBalance(address: PublicKey) {
+    return (await this.offchainState.fields.accounts.get(address)).orElse(0n);
+  }
+
+  @method
+  async settle(proof: StateProof) {
+    await this.offchainState.settle(proof);
+  }
+}

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

@@ -0,0 +1,184 @@
+import {
+  SmartContract,
+  method,
+  state,
+  PublicKey,
+  UInt64,
+} from '../../../../index.js';
+import * as Mina from '../../mina.js';
+import assert from 'assert';
+import { ExampleContract } from './ExampleContract.js';
+import { settle, transfer } from './utils.js';
+
+const Local = await Mina.LocalBlockchain({ proofsEnabled: false });
+Mina.setActiveInstance(Local);
+
+const [
+  sender,
+  receiver1,
+  receiver2,
+  receiver3,
+  contractAccountA,
+  contractAccountB,
+] = Local.testAccounts;
+
+const contractA = new ExampleContract(contractAccountA);
+const contractB = new ExampleContract(contractAccountB);
+contractA.offchainState.setContractInstance(contractA);
+contractB.offchainState.setContractInstance(contractB);
+
+console.time('deploy contract');
+const deployTx = Mina.transaction(sender, async () => {
+  await contractA.deploy();
+  await contractB.deploy();
+});
+await deployTx.sign([sender.key, contractAccountA.key, contractAccountB.key]);
+await deployTx.prove();
+await deployTx.send().wait();
+console.timeEnd('deploy contract');
+
+console.time('create accounts');
+const accountCreationTx = Mina.transaction(sender, async () => {
+  await contractA.createAccount(sender, UInt64.from(1000));
+  await contractA.createAccount(receiver2, UInt64.from(1000));
+  await contractB.createAccount(sender, UInt64.from(1500));
+});
+await accountCreationTx.sign([sender.key]);
+await accountCreationTx.prove();
+await accountCreationTx.send().wait();
+console.timeEnd('create accounts');
+
+console.time('settle');
+await settle(contractA, sender);
+await settle(contractB, sender);
+console.timeEnd('settle');
+
+assert((await contractA.getSupply()).toBigInt() == 1000n);
+assert((await contractB.getSupply()).toBigInt() == 1500n);
+
+console.log('Initial balances:');
+console.log(
+  'Contract A, Sender: ',
+  (await contractA.offchainState.fields.accounts.get(sender)).value.toBigInt()
+);
+console.log(
+  'Contract B, Sender: ',
+  (await contractB.offchainState.fields.accounts.get(sender)).value.toBigInt()
+);
+assert((await contractA.getBalance(sender)).toBigInt() == 1000n);
+assert((await contractB.getBalance(sender)).toBigInt() == 1500n);
+
+console.time('transfer');
+await transfer(contractA, sender, receiver1, UInt64.from(100));
+await settle(contractA, sender);
+await transfer(contractA, sender, receiver2, UInt64.from(200));
+await settle(contractA, sender);
+await transfer(contractA, sender, receiver3, UInt64.from(300));
+await transfer(contractB, sender, receiver1, UInt64.from(200));
+console.timeEnd('transfer');
+
+console.time('settle');
+await settle(contractA, sender);
+await settle(contractB, sender);
+console.timeEnd('settle');
+
+console.log('After Settlement balances:');
+console.log(
+  'Contract A, Sender: ',
+  (await contractA.offchainState.fields.accounts.get(sender)).value.toBigInt()
+);
+console.log(
+  'Contract A, Receiver 1: ',
+  (
+    await contractA.offchainState.fields.accounts.get(receiver1)
+  ).value.toBigInt()
+);
+console.log(
+  'Contract A, Receiver 2: ',
+  (
+    await contractA.offchainState.fields.accounts.get(receiver2)
+  ).value.toBigInt()
+);
+console.log(
+  'Contract A, Receiver 3: ',
+  (
+    await contractA.offchainState.fields.accounts.get(receiver3)
+  ).value.toBigInt()
+);
+
+console.log(
+  'Contract B, Sender: ',
+  (await contractB.offchainState.fields.accounts.get(sender)).value.toBigInt()
+);
+console.log(
+  'Contract B, Receiver 1: ',
+  (
+    await contractB.offchainState.fields.accounts.get(receiver1)
+  ).value.toBigInt()
+);
+assert((await contractA.getBalance(sender)).toBigInt() == 400n);
+assert((await contractA.getBalance(receiver1)).toBigInt() == 100n);
+assert((await contractA.getBalance(receiver2)).toBigInt() == 200n);
+assert((await contractA.getBalance(receiver3)).toBigInt() == 300n);
+
+assert((await contractB.getBalance(sender)).toBigInt() == 1300n);
+assert((await contractB.getBalance(receiver1)).toBigInt() == 200n);
+
+console.time('advance contract A state but leave B unsettled');
+await transfer(contractA, sender, receiver1, UInt64.from(150)); // 250, 250, 200, 300
+await transfer(contractA, receiver2, receiver3, UInt64.from(20)); // 250, 250, 180, 320
+await settle(contractA, sender);
+
+await transfer(contractA, receiver1, receiver2, UInt64.from(50)); // 250, 200, 230, 320
+await transfer(contractA, receiver3, sender, UInt64.from(50)); // 300, 200, 230, 270
+await settle(contractA, sender);
+
+await transfer(contractB, sender, receiver1, UInt64.from(5));
+console.timeEnd('advance contract A state but leave B unsettled');
+
+assert((await contractA.getSupply()).toBigInt() == 1000n);
+assert((await contractB.getSupply()).toBigInt() == 1500n);
+
+console.log('Final balances:');
+console.log(
+  'Contract A, Sender: ',
+  (await contractA.offchainState.fields.accounts.get(sender)).value.toBigInt()
+);
+console.log(
+  'Contract A, Receiver 1: ',
+  (
+    await contractA.offchainState.fields.accounts.get(receiver1)
+  ).value.toBigInt()
+);
+console.log(
+  'Contract A, Receiver 2: ',
+  (
+    await contractA.offchainState.fields.accounts.get(receiver2)
+  ).value.toBigInt()
+);
+console.log(
+  'Contract A, Receiver 3: ',
+  (
+    await contractA.offchainState.fields.accounts.get(receiver3)
+  ).value.toBigInt()
+);
+
+console.log(
+  'Contract B, Sender: ',
+  (await contractB.offchainState.fields.accounts.get(sender)).value.toBigInt()
+);
+console.log(
+  'Contract B, Receiver: ',
+  (
+    await contractB.offchainState.fields.accounts.get(receiver1)
+  ).value.toBigInt()
+);
+
+assert((await contractA.getBalance(sender)).toBigInt() == 300n);
+assert((await contractA.getBalance(receiver1)).toBigInt() == 200n);
+assert((await contractA.getBalance(receiver2)).toBigInt() == 230n);
+assert((await contractA.getBalance(receiver3)).toBigInt() == 270n);
+
+// The 5 token transfer has not been settled
+assert((await contractB.getBalance(sender)).toBigInt() == 1300n);
+assert((await contractB.getBalance(receiver1)).toBigInt() == 200n);

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

@@ -0,0 +1,62 @@
+import { UInt64 } from '../../../../index.js';
+import * as Mina from '../../mina.js';
+import assert from 'assert';
+
+import {
+  ExampleContract,
+  offchainState as exampleOffchainState,
+} from './ExampleContract.js';
+import { settle, transfer } from './utils.js';
+
+const Local = await Mina.LocalBlockchain({ proofsEnabled: true });
+Mina.setActiveInstance(Local);
+
+const [sender, receiver, contractAccount] = Local.testAccounts;
+
+const contract = new ExampleContract(contractAccount);
+contract.offchainState.setContractInstance(contract);
+
+console.time('compile offchain state program');
+await exampleOffchainState.compile();
+console.timeEnd('compile offchain state program');
+
+console.time('compile contract');
+await ExampleContract.compile();
+console.timeEnd('compile contract');
+
+console.time('deploy contract');
+const deployTx = Mina.transaction(sender, async () => {
+  await contract.deploy();
+});
+await deployTx.sign([sender.key, contractAccount.key]);
+await deployTx.prove();
+await deployTx.send().wait();
+console.timeEnd('deploy contract');
+
+console.time('create accounts');
+const accountCreationTx = Mina.transaction(sender, async () => {
+  await contract.createAccount(sender, UInt64.from(1000));
+});
+await accountCreationTx.sign([sender.key]);
+await accountCreationTx.prove();
+await accountCreationTx.send().wait();
+console.timeEnd('create accounts');
+
+console.time('settle');
+await settle(contract, sender);
+console.timeEnd('settle');
+
+assert((await contract.getSupply()).toBigInt() == 1000n);
+assert((await contract.getBalance(sender)).toBigInt() == 1000n);
+
+console.time('transfer');
+await transfer(contract, sender, receiver, UInt64.from(100));
+console.timeEnd('transfer');
+
+console.time('settle');
+await settle(contract, sender);
+console.timeEnd('settle');
+
+assert((await contract.getSupply()).toBigInt() == 1000n);
+assert((await contract.getBalance(sender)).toBigInt() == 900n);
+assert((await contract.getBalance(receiver)).toBigInt() == 100n);

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

@@ -0,0 +1,28 @@
+import { PublicKey, UInt64 } from '../../../../index.js';
+import * as Mina from '../../mina.js';
+
+import { ExampleContract } from './ExampleContract.js';
+
+export { transfer, settle };
+
+async function transfer(
+  contract: ExampleContract,
+  sender: Mina.TestPublicKey,
+  receiver: PublicKey,
+  amount: UInt64
+) {
+  const tx = Mina.transaction(sender, async () => {
+    await contract.transfer(sender, receiver, amount);
+  });
+  tx.sign([sender.key]);
+  await tx.prove().send().wait();
+}
+
+async function settle(contract: ExampleContract, sender: Mina.TestPublicKey) {
+  const proof = await contract.offchainState.createSettlementProof();
+  const tx = Mina.transaction(sender, async () => {
+    await contract.settle(proof);
+  });
+  tx.sign([sender.key]);
+  await tx.prove().send().wait();
+}