o1-labs / o1js

TypeScript framework for zk-SNARKs and zkApps
https://docs.minaprotocol.com/en/zkapps/how-to-write-a-zkapp
Apache License 2.0
500 stars 110 forks source link

Batch Reducer #1676

Closed mitschabaude closed 2 months ago

mitschabaude commented 3 months ago

Introduces a new kind of reducer for the case that you only want to reduce actions in small batches. Avoids deadlock when more actions than your batch size are pending.

This targets zkapps which want to create account updates for each action.

Example application: A token airdrop, where users submit "claim" actions concurrently, and we process them in a reducer by paying out the claims and simultaneously removing the claims from our storage (to avoid double-claims).

The batch reducer unit test implements such an airdrop.

Batch reducer API

A batch reducer is declared by specifying an action type and a batch size:

let batchReducer = new BatchReducer({ actionType: Action, batchSize: 5 });

The action class is a generic type argument:

batchReducer satisfies Experimental.BatchReducer<typeof Action>;

There are essentially 3 methods for interacting with batchReducer, the first of which is familiar

// dispatch an action
batchReducer.dispatch(action: Action): void;

Reducing actions comes with a separate preparation step outside your contract:

// outside contract: prepare array of batches of actions
batchReducer.prepareBatches(): Promise<Array<{ batch: ActionBatch<Action>; proof: Proof }>>;

// inside contract: process one batch of actions
batchReducer.processBatch({ batch, proof }, callback: (action: Action, isDummy: Bool) => void): void;

processBatch() is what you would call in your "reducer" method. If we want to catch up with the current chain, we have to call that method, in a separate transaction, for each of the batches returned by prepareBatches().

Note: I went for a forEach() type of API instead of reduce(), because the API surface is simpler and usually allows you to easily implement reduce() like flows anyway, by mutating a variable in scope.

A batch reducer uses 2 onchain state fields (one more than a classical reducer):

class Contract extends SmartContract {
  @state(Field)
  actionState = State(BatchReducer.initialActionState);
  @state(Field)
  actionStack = State(BatchReducer.initialActionStack);
}

Batch reducer implementation

Under the hood, processBatch() creates a recursive proof of reversing the action order: Its output is a stack of actions where the top-most action is the next pending one, and the bottom-most action is the last one submitted to the chain.

The reason for this reversal is that popping the last element of a (Merkle) list can be implemented as a provable operation, which proves that the popped element was indeed part of the original list. So, once we created a reversed stack of actions, we can pop off the same stack for multiple transactions, without having to create the recursive proof again.

Reusing the same action stack for many transactions means that we have to store it on the contract. The actionStack stored on the contract is always guaranteed to consist only of pending actions in the correct (reversed) order; so, popping off this stack will never yield an invalid action. Once the stack runs out of actions, we have to call prepareBatch() again, to give us a fresh stack.

Note that I said that prepareBatch() creates a recursive proof of creating a new action stack. It being recursive means it can handle an arbitrary amount of pending actions, so we will always be able to catch up with the chain (note: stack proving is very efficient, it can do hundreds of action lists per proof). As an optimization, however, we don't actually prove the full stack offchain. We leave a final chunk to be proved inside processBatch() itself. If this final chunk is large enough to build the entire stack, then we don't even need a recursive proof. prepareBatch() will figure out all of that and return quickly if no proof is needed.

To understand the processBatch() logic, note that it handles two different cases:

  1. We are calling it with a freshly created stack of actions (to be completed within the method itself), which successfully connects us to the onchain action state. In this case, we overwrite the previous stack stored onchain (which typically will be exhausted or at least smaller than our new stack).
    • This is what happens in the first batch of the batches that prepareBatches() returns
  2. We are calling it with no new stack but instructing it to use the onchain stack
    • This happens on all subsequent batches that prepareBatches() returns

There is a lot of robustness in this design. We don't have to finish publishing transactions for all the batches before calling prepareBatches() again. We could always create a new stack before calling processBatch() -- worst case, the onchain stack still had some actions in it, so we wasted a small amount of recursive proving (note: this is only a waste if we even need a recursive proof). Also, we will never break anything by instructing processBatch() to use the onchain stack -- worst case, it's empty and no progress is made.

zkApp Testing framework

As an aside, this also introduces a little testing framework intended to make writing local blockchain tests easier and quicker. A "test" is modeled as a list of "instructions" that either:

Other additions

In the spirit of always adding the methods to o1js that we need ourselves, I added a whole bunch of new API in this PR. See the changelog. For example, conditional versions of many existing APIs:

These kind of APIs are especially useful for reducer-type logic where we can never do assertions and always have to handle invalid cases gracefully: by not doing anything. We need more of them!

mitschabaude commented 3 months ago

btw I had some explanatory comments here: https://github.com/o1-labs/o1js/pull/1676#pullrequestreview-2125551470

mitschabaude commented 2 months ago

What do you think about adding a variant that does this without recursion for smaller batches, directly in the smart contract method that calls reduce - essentially a low level variation of the recursion one with less abstractions?

The current implementation already does this! The smart contract method does the final "chunk" of the recursive work, which will be all of the work if there are few pending actions.

But we could have a variant which doesn't even need external input. Are you suggesting that? On the other hand, I don't see how this would improve the DX because you still need the most general method with recursion either way, as a backup

Trivo25 commented 2 months ago

But we could have a variant which doesn't even need external input. Are you suggesting that?

Yea, that's what I was referring to! A lightweight version that doesn't require any external inputs or proof/compilation/zkProgram dependencies.

mitschabaude commented 2 months ago

But we could have a variant which doesn't even need external input. Are you suggesting that?

Yea, that's what I was referring to! A lightweight version that doesn't require any external inputs or proof/compilation/zkProgram dependencies.

Right, good point about saving the compilation step as well! In that case, it will have to be on a separate smart contract method though. But I can imagine people making use of that.

I'm going to revisit the PR next week to see if I can break up the reducer logic into more modular steps. then it should be very easy to add variations of processBatch() and even let people write their own variations with more low-level methods

mitschabaude commented 2 months ago

well @Trivo25, on the OTHER hand.. I just remembered why I was convinced we would always need an external call 😅

people need an external call to figure out how many batches there are == how many transactions to create.

the lightweight method doesn't really have a way to tell you that, or even tell you that no transactions are currently needed.

I'm still in favor of adding more flexibility but I can't see when you wouldn't want to call an external method 🤔

mitschabaude commented 2 months ago

Bigger review incoming! But in the meantime, do we not intend to make the changes for testing public? Is this intended to be an internal only tool?

I think we should make them public. It would be nice to do some polishing on it though,

I don't want to put the burden of polishing the testing framework to be ready for release on this already huge PR

mitschabaude commented 2 months ago

@Trivo25 I started this draft PR which exposes more modularity and a way to do everything without external inputs (marked as "unsafe" though): https://github.com/o1-labs/o1js/pull/1733

Does this go in the direction of what you were thinking?