/
forest-iterator.ts
143 lines (127 loc) · 5.15 KB
/
forest-iterator.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import {
AccountUpdate,
AccountUpdateForest,
AccountUpdateTreeBase,
TokenId,
} from '../account-update.js';
import { Field } from '../../provable/wrapped.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.createFromList(AccountUpdateForest);
class Layer extends Struct({
forest: AccountUpdateIterator.provable,
mayUseToken: AccountUpdate.MayUseToken.type,
}) {}
const ParentLayers = MerkleList.create<Layer>(Layer);
type MayUseToken = AccountUpdate['body']['mayUseToken'];
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 {
currentLayer: Layer;
unfinishedParentLayers: MerkleList<Layer>;
selfToken: Field;
constructor(
forest: MerkleListIterator<AccountUpdateTreeBase>,
mayUseToken: MayUseToken,
selfToken: Field
) {
this.currentLayer = { forest, mayUseToken };
this.unfinishedParentLayers = ParentLayers.empty();
this.selfToken = selfToken;
}
static create(forest: AccountUpdateForest, selfToken: Field) {
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?: string) {
assert(
this.currentLayer.forest.isAtStart(),
message ?? 'TokenAccountUpdateIterator not finished'
);
}
}