Skip to content

Commit 9dd8c82

Browse files
rpanicrpanic
committed
Changed STProver to work on appliedbatches
1 parent 90e64b8 commit 9dd8c82

File tree

6 files changed

+357
-253
lines changed

6 files changed

+357
-253
lines changed

‎packages/protocol/src/model/StateTransitionProvableBatch.ts‎

Lines changed: 122 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -1,67 +1,91 @@
1-
import { Bool, Provable, Struct } from "o1js";
1+
import { Bool, Field, Poseidon, Provable, Struct } from "o1js";
22
import {
3-
InMemoryMerkleTreeStorage,
43
range,
4+
InMemoryMerkleTreeStorage,
55
RollupMerkleTree,
6-
RollupMerkleTreeWitness,
6+
RollupMerkleTreeWitness
77
} from "@proto-kit/common";
88

99
import { constants } from "../Constants";
1010

1111
import { ProvableStateTransition } from "./StateTransition.js";
1212

1313
export class StateTransitionType {
14-
public static readonly normal = true;
15-
16-
public static readonly protocol = false;
14+
public static readonly nothing = 2;
15+
// The reason these are 0 and 1 is to efficiently check
16+
// x in [inside, closing] in-circuit via the boolean trick
17+
public static readonly closeAndApply = 1;
18+
public static readonly closeAndThrowAway = 0;
19+
}
1720

18-
public static isNormal(type: boolean) {
19-
return type === StateTransitionType.normal;
21+
export class ProvableStateTransitionType extends Struct({
22+
type: Field,
23+
}) {
24+
public static get nothing(): ProvableStateTransitionType {
25+
return this.from(StateTransitionType.nothing);
2026
}
2127

22-
public static isProtocol(type: boolean) {
23-
return type === StateTransitionType.protocol;
28+
public static get closeAndApply(): ProvableStateTransitionType {
29+
return this.from(StateTransitionType.closeAndApply);
2430
}
25-
}
2631

27-
export class ProvableStateTransitionType extends Struct({
28-
type: Bool,
29-
}) {
30-
public static get normal(): ProvableStateTransitionType {
31-
return new ProvableStateTransitionType({
32-
type: Bool(StateTransitionType.normal),
33-
});
32+
public static get closeAndThrowAway(): ProvableStateTransitionType {
33+
return this.from(StateTransitionType.closeAndThrowAway);
3434
}
3535

36-
public static get protocol(): ProvableStateTransitionType {
36+
private static from(constant: number) {
3737
return new ProvableStateTransitionType({
38-
type: Bool(StateTransitionType.protocol),
38+
type: Field(constant),
3939
});
4040
}
4141

42-
public isNormal(): Bool {
43-
return this.type;
42+
public isClosing() {
43+
const { type } = this;
44+
// check if base is 0 or 1
45+
// 0^2 == 0 && 1^2 == 1
46+
return type.mul(type).equals(type);
4447
}
48+
}
4549

46-
public isProtocol(): Bool {
47-
return this.type.not();
50+
export class AppliedStateTransitionBatch extends Struct({
51+
batchHash: Field,
52+
applied: Bool,
53+
}) {}
54+
55+
export class AppliedStateTransitionBatchState extends Struct({
56+
batchHash: Field,
57+
root: Field,
58+
}) {
59+
public hashOrZero(): Field {
60+
const hash = Poseidon.hash(AppliedStateTransitionBatchState.toFields(this));
61+
return Provable.if(this.batchHash.equals(0), Field(0), hash);
4862
}
4963
}
5064

65+
export class MerkleWitnessBatch extends Struct({
66+
witnesses: Provable.Array(
67+
RollupMerkleTreeWitness,
68+
constants.stateTransitionProverBatchSize
69+
),
70+
}) {}
71+
5172
/**
5273
* A Batch of StateTransitions to be consumed by the StateTransitionProver
5374
* to prove multiple STs at once
5475
*
55-
* transitionType:
56-
* true == normal ST, false == protocol ST
76+
* bases: Describes the state root on which the ST will be applied on
77+
* If it is zero, this means that this ST should connect with the previous one
78+
* If it is one, this means that the batch should be closed
5779
*/
5880
export class StateTransitionProvableBatch extends Struct({
5981
batch: Provable.Array(
6082
ProvableStateTransition,
6183
constants.stateTransitionProverBatchSize
6284
),
6385

64-
transitionTypes: Provable.Array(
86+
// bases: Provable.Array(Field, constants.stateTransitionProverBatchSize),
87+
88+
types: Provable.Array(
6589
ProvableStateTransitionType,
6690
constants.stateTransitionProverBatchSize
6791
),
@@ -71,75 +95,75 @@ export class StateTransitionProvableBatch extends Struct({
7195
constants.stateTransitionProverBatchSize
7296
),
7397
}) {
74-
public static fromMappings(
75-
transitions: {
76-
transition: ProvableStateTransition;
77-
type: ProvableStateTransitionType;
78-
}[],
79-
merkleWitnesses: RollupMerkleTreeWitness[]
80-
): StateTransitionProvableBatch {
81-
const batch = transitions.map((entry) => entry.transition);
82-
const transitionTypes = transitions.map((entry) => entry.type);
83-
const witnesses = merkleWitnesses.slice();
84-
// Check that order is correct
85-
let normalSTsStarted = false;
86-
transitionTypes.forEach((x) => {
87-
if (!normalSTsStarted && x.isNormal().toBoolean()) {
88-
normalSTsStarted = true;
89-
}
90-
if (normalSTsStarted && x.isProtocol().toBoolean()) {
91-
throw new Error("Order in initializing STBatch not correct");
92-
}
93-
});
94-
95-
while (batch.length < constants.stateTransitionProverBatchSize) {
96-
batch.push(ProvableStateTransition.dummy());
97-
transitionTypes.push(ProvableStateTransitionType.normal);
98-
witnesses.push(
99-
new RollupMerkleTree(new InMemoryMerkleTreeStorage()).getWitness(
100-
BigInt(0)
101-
)
102-
);
103-
}
104-
return new StateTransitionProvableBatch({
105-
batch,
106-
transitionTypes,
107-
merkleWitnesses: witnesses,
108-
});
109-
}
110-
111-
public static fromTransitions(
112-
transitions: ProvableStateTransition[],
113-
protocolTransitions: ProvableStateTransition[],
114-
merkleWitnesses: RollupMerkleTreeWitness[]
115-
): StateTransitionProvableBatch {
116-
const array = transitions.slice().concat(protocolTransitions);
117-
118-
const transitionTypes = range(0, transitions.length)
119-
.map(() => ProvableStateTransitionType.normal)
120-
.concat(
121-
range(0, protocolTransitions.length).map(
122-
() => ProvableStateTransitionType.protocol
123-
)
124-
);
125-
126-
while (array.length < constants.stateTransitionProverBatchSize) {
127-
array.push(ProvableStateTransition.dummy());
128-
transitionTypes.push(ProvableStateTransitionType.normal);
129-
}
130-
131-
return new StateTransitionProvableBatch({
132-
batch: array,
133-
transitionTypes,
134-
merkleWitnesses,
135-
});
136-
}
137-
138-
private constructor(object: {
139-
batch: ProvableStateTransition[];
140-
transitionTypes: ProvableStateTransitionType[];
141-
merkleWitnesses: RollupMerkleTreeWitness[];
142-
}) {
143-
super(object);
144-
}
98+
// public static fromMappings(
99+
// transitions: {
100+
// transition: ProvableStateTransition;
101+
// type: ProvableStateTransitionType;
102+
// }[],
103+
// merkleWitnesses: RollupMerkleTreeWitness[]
104+
// ): StateTransitionProvableBatch {
105+
// const batch = transitions.map((entry) => entry.transition);
106+
// const transitionTypes = transitions.map((entry) => entry.type);
107+
// const witnesses = merkleWitnesses.slice();
108+
// // Check that order is correct
109+
// let normalSTsStarted = false;
110+
// transitionTypes.forEach((x) => {
111+
// if (!normalSTsStarted && x.isNormal().toBoolean()) {
112+
// normalSTsStarted = true;
113+
// }
114+
// if (normalSTsStarted && x.isProtocol().toBoolean()) {
115+
// throw new Error("Order in initializing STBatch not correct");
116+
// }
117+
// });
118+
//
119+
// while (batch.length < constants.stateTransitionProverBatchSize) {
120+
// batch.push(ProvableStateTransition.dummy());
121+
// transitionTypes.push(ProvableStateTransitionType.normal);
122+
// witnesses.push(
123+
// new RollupMerkleTree(new InMemoryMerkleTreeStorage()).getWitness(
124+
// BigInt(0)
125+
// )
126+
// );
127+
// }
128+
// return new StateTransitionProvableBatch({
129+
// batch,
130+
// transitionTypes,
131+
// merkleWitnesses: witnesses,
132+
// });
133+
// }
134+
//
135+
// public static fromTransitions(
136+
// transitions: ProvableStateTransition[],
137+
// protocolTransitions: ProvableStateTransition[],
138+
// merkleWitnesses: RollupMerkleTreeWitness[]
139+
// ): StateTransitionProvableBatch {
140+
// const array = transitions.slice().concat(protocolTransitions);
141+
//
142+
// const transitionTypes = range(0, transitions.length)
143+
// .map(() => ProvableStateTransitionType.normal)
144+
// .concat(
145+
// range(0, protocolTransitions.length).map(
146+
// () => ProvableStateTransitionType.protocol
147+
// )
148+
// );
149+
//
150+
// while (array.length < constants.stateTransitionProverBatchSize) {
151+
// array.push(ProvableStateTransition.dummy());
152+
// transitionTypes.push(ProvableStateTransitionType.normal);
153+
// }
154+
//
155+
// return new StateTransitionProvableBatch({
156+
// batch: array,
157+
// transitionTypes,
158+
// merkleWitnesses,
159+
// });
160+
// }
161+
//
162+
// private constructor(object: {
163+
// batch: ProvableStateTransition[];
164+
// transitionTypes: ProvableStateTransitionType[];
165+
// merkleWitnesses: RollupMerkleTreeWitness[];
166+
// }) {
167+
// super(object);
168+
// }
145169
}

‎packages/protocol/src/prover/statetransition/StateTransitionProvable.ts‎

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,23 @@
11
import { Field, Proof, Struct } from "o1js";
22
import { WithZkProgrammable, CompilableModule } from "@proto-kit/common";
33

4-
import { StateTransitionProvableBatch } from "../../model/StateTransitionProvableBatch";
4+
import {
5+
AppliedStateTransitionBatch,
6+
AppliedStateTransitionBatchState,
7+
MerkleWitnessBatch,
8+
StateTransitionProvableBatch,
9+
} from "../../model/StateTransitionProvableBatch";
510

611
export class StateTransitionProverPublicInput extends Struct({
7-
stateTransitionsHash: Field,
8-
protocolTransitionsHash: Field,
9-
stateRoot: Field,
10-
protocolStateRoot: Field,
12+
batchesHash: Field,
13+
currentBatchStateHash: Field,
14+
root: Field,
1115
}) {}
1216

1317
export class StateTransitionProverPublicOutput extends Struct({
14-
stateTransitionsHash: Field,
15-
protocolTransitionsHash: Field,
16-
stateRoot: Field,
17-
protocolStateRoot: Field,
18+
batchesHash: Field,
19+
currentBatchStateHash: Field,
20+
root: Field,
1821
}) {}
1922

2023
export type StateTransitionProof = Proof<
@@ -30,7 +33,9 @@ export interface StateTransitionProvable
3033
CompilableModule {
3134
runBatch: (
3235
publicInput: StateTransitionProverPublicInput,
33-
batch: StateTransitionProvableBatch
36+
batch: StateTransitionProvableBatch,
37+
witnesses: MerkleWitnessBatch,
38+
currentAppliedBatch: AppliedStateTransitionBatchState
3439
) => Promise<StateTransitionProverPublicOutput>;
3540

3641
merge: (

0 commit comments

Comments
 (0)