o1js
Version:
TypeScript framework for zk-SNARKs and zkApps
92 lines • 5.08 kB
JavaScript
import { AccountUpdate, AccountUpdateForest, TokenId, } from '../account-update.js';
import { Provable } from '../../../provable/provable.js';
import { Struct } from '../../../provable/types/struct.js';
import { assert } from '../../../provable/gadgets/common.js';
import { MerkleListIterator, MerkleList } from '../../../provable/merkle-list.js';
export { TokenAccountUpdateIterator };
const AccountUpdateIterator = MerkleListIterator.create(AccountUpdateForest.prototype.innerProvable, AccountUpdateForest._nextHash, AccountUpdateForest.emptyHash);
class Layer extends Struct({
forest: AccountUpdateIterator,
mayUseToken: AccountUpdate.MayUseToken.type,
}) {
}
const ParentLayers = MerkleList.create(Layer);
const MayUseToken = AccountUpdate.MayUseToken;
/**
* Data structure to represent a forest of account updates that is being iterated over,
* in the context of a token manager contract.
*
* The iteration is done in a depth-first manner.
*
* ```ts
* let forest: AccountUpdateForest = ...;
* let tokenIterator = TokenAccountUpdateIterator.create(forest, tokenId);
*
* // process the first 5 account updates in the tree
* for (let i = 0; i < 5; i++) {
* let { accountUpdate, usesThisToken } = tokenIterator.next();
* // ... do something with the account update ...
* }
* ```
*
* **Important**: Since this is specifically used by token manager contracts to process their entire subtree
* of account updates, the iterator skips subtrees that don't inherit token permissions and can therefore definitely not use the token.
*
* So, the assumption is that the consumer of this iterator is only interested in account updates that use the token.
* We still can't avoid processing some account updates that don't use the token, therefore the iterator returns a boolean
* `usesThisToken` alongside each account update.
*/
class TokenAccountUpdateIterator {
constructor(forest, mayUseToken, selfToken) {
this.currentLayer = { forest, mayUseToken };
this.unfinishedParentLayers = ParentLayers.empty();
this.selfToken = selfToken;
}
static create(forest, selfToken) {
return new TokenAccountUpdateIterator(AccountUpdateIterator.startIteratingFromLast(forest), MayUseToken.ParentsOwnToken, selfToken);
}
/**
* Make a single step along a tree of account updates.
*
* This function is guaranteed to visit each account update in the tree that uses the token
* exactly once, when called repeatedly.
*
* The method makes a best effort to avoid visiting account updates that are not using the token,
* and in particular, to avoid returning dummy updates.
* However, neither can be ruled out. We're returning { update, usesThisToken: Bool } and let the
* caller handle the irrelevant case where `usesThisToken` is false.
*/
next() {
// get next account update from the current forest (might be a dummy)
let { accountUpdate, children } = this.currentLayer.forest.previous();
let childForest = AccountUpdateIterator.startIteratingFromLast(children);
let childLayer = {
forest: childForest,
mayUseToken: MayUseToken.InheritFromParent,
};
let update = accountUpdate.unhash();
let usesThisToken = update.tokenId.equals(this.selfToken);
// check if this account update / it's children can use the token
let canAccessThisToken = Provable.equal(MayUseToken.type, update.body.mayUseToken, this.currentLayer.mayUseToken);
let isSelf = TokenId.derive(update.publicKey, update.tokenId).equals(this.selfToken);
// if we don't have to check the children, ignore the forest by jumping to its end
let skipSubtree = canAccessThisToken.not().or(isSelf);
childForest.jumpToStartIf(skipSubtree);
// there are three cases for how to proceed:
// 1. if we have to process children, we step down and add the current layer to the stack of unfinished parent layers
// 2. if we don't have to process children, but are not finished with the current layer, we stay in the current layer
// (below, this is the case where the current layer is first pushed to and then popped from the stack of unfinished parent layers)
// 3. if both of the above are false, we step up to the next unfinished parent layer
let currentForest = this.currentLayer.forest;
let currentLayerFinished = currentForest.isAtStart();
let childLayerFinished = childForest.isAtStart();
this.unfinishedParentLayers.pushIf(currentLayerFinished.not(), this.currentLayer);
let currentOrParentLayer = this.unfinishedParentLayers.popIf(childLayerFinished);
this.currentLayer = Provable.if(childLayerFinished, Layer, currentOrParentLayer, childLayer);
return { accountUpdate: update, usesThisToken };
}
assertFinished(message) {
assert(this.currentLayer.forest.isAtStart(), message ?? 'TokenAccountUpdateIterator not finished');
}
}
//# sourceMappingURL=forest-iterator.js.map