Skip to content

Commit a45933f

Browse files
rpanicrpanic
committed
Fixed new CompileRegistry to work
1 parent 92f8420 commit a45933f

File tree

15 files changed

+216
-233
lines changed

15 files changed

+216
-233
lines changed

‎packages/module/src/runtime/Runtime.ts‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
/* eslint-disable @typescript-eslint/no-unsafe-assignment,@typescript-eslint/no-unsafe-argument */
22
import { ZkProgram } from "o1js";
3-
import { DependencyContainer, injectable } from "tsyringe";
3+
import { container, DependencyContainer, injectable } from "tsyringe";
44
import {
55
StringKeyOf,
66
ModuleContainer,
@@ -16,6 +16,11 @@ import {
1616
MethodPublicOutput,
1717
StateServiceProvider,
1818
SimpleAsyncStateService,
19+
CompilableModule,
20+
CompileRegistry,
21+
RuntimeMethodExecutionContext,
22+
RuntimeTransaction,
23+
NetworkState,
1924
} from "@proto-kit/protocol";
2025

2126
import {
@@ -264,7 +269,7 @@ export class RuntimeZkProgrammable<
264269
@injectable()
265270
export class Runtime<Modules extends RuntimeModulesRecord>
266271
extends ModuleContainer<Modules>
267-
implements RuntimeEnvironment
272+
implements RuntimeEnvironment, CompilableModule
268273
{
269274
public static from<Modules extends RuntimeModulesRecord>(
270275
definition: RuntimeDefinition<Modules>
@@ -377,5 +382,16 @@ export class Runtime<Modules extends RuntimeModulesRecord>
377382
public get runtimeModuleNames() {
378383
return Object.keys(this.definition.modules);
379384
}
385+
386+
public async compile(registry: CompileRegistry) {
387+
return await registry.compileModule(async () => {
388+
const context = container.resolve(RuntimeMethodExecutionContext);
389+
context.setup({
390+
transaction: RuntimeTransaction.dummyTransaction(),
391+
networkState: NetworkState.empty(),
392+
});
393+
return await this.zkProgrammable.compile();
394+
});
395+
}
380396
}
381397
/* eslint-enable @typescript-eslint/no-unsafe-assignment,@typescript-eslint/no-unsafe-argument */

‎packages/protocol/src/compiling/AtomicCompileHelper.ts‎

Lines changed: 17 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,61 +1,49 @@
11
import {
22
AreProofsEnabled,
3+
CompileArtifact,
34
log,
45
MOCK_VERIFICATION_KEY,
56
} from "@proto-kit/common";
7+
import { SmartContract, VerificationKey } from "o1js";
68

7-
export type Artifact = string | object | undefined;
9+
export type ArtifactRecord = Record<string, CompileArtifact>;
810

9-
export type GenericCompileTarget<T extends Artifact> = {
10-
compile: () => Promise<T>;
11+
export type CompileTarget = {
12+
name: string;
13+
compile: () => Promise<CompileArtifact>;
1114
};
1215

1316
export class AtomicCompileHelper {
1417
public constructor(private readonly areProofsEnabled: AreProofsEnabled) {}
1518

1619
private compilationPromises: {
17-
[key: string]: Promise<Artifact | undefined>;
20+
[key: string]: Promise<CompileArtifact>;
1821
} = {};
1922

20-
// Generic params for zkProgrammable should be unknown, but verify makes those types invariant
21-
// public async zkProgrammable(zkProgrammable: ZkProgrammable<any, any>) {
22-
// await reduceSequential(
23-
// zkProgrammable.zkProgram,
24-
// async (acc, program) => {
25-
// const res = await this.program(program);
26-
// return {
27-
// ...acc,
28-
// [program.name]: res,
29-
// };
30-
// },
31-
// {}
32-
// );
33-
// }
34-
35-
public async program<ReturnArtifact extends Artifact>(
36-
name: string,
37-
contract: GenericCompileTarget<ReturnArtifact>,
23+
public async compileContract(
24+
contract: CompileTarget,
3825
overrideProofsEnabled?: boolean
39-
): Promise<ReturnArtifact> {
26+
): Promise<CompileArtifact> {
4027
let newPromise = false;
28+
const { name } = contract;
4129
if (this.compilationPromises[name] === undefined) {
4230
const proofsEnabled =
4331
overrideProofsEnabled ?? this.areProofsEnabled.areProofsEnabled;
44-
if (proofsEnabled) {
32+
// This wierd any is necessary otherwise the compiler optimized that check away
33+
if (proofsEnabled || !((contract as any) instanceof SmartContract)) {
4534
log.time(`Compiling ${name}`);
4635
this.compilationPromises[name] = contract.compile();
4736
newPromise = true;
4837
} else {
49-
// TODO Mock VK here is not at all generalized and safe
50-
// Better way would be to package the Smart contracts into a mock-compile as well
51-
this.compilationPromises[name] = Promise.resolve(MOCK_VERIFICATION_KEY);
38+
this.compilationPromises[name] = Promise.resolve({
39+
verificationKey: MOCK_VERIFICATION_KEY,
40+
});
5241
}
5342
}
5443
const result = await this.compilationPromises[name];
5544
if (newPromise) {
5645
log.timeEnd.info(`Compiling ${name}`);
5746
}
58-
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
59-
return result as ReturnArtifact;
47+
return result;
6048
}
6149
}

‎packages/protocol/src/compiling/CompilableModule.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import type { CompileRegistry } from "./CompileRegistry";
2-
import { Artifact } from "./AtomicCompileHelper";
2+
import { ArtifactRecord } from "./AtomicCompileHelper";
33

44
export interface CompilableModule {
5-
compile(registry: CompileRegistry): Promise<Artifact | void>;
5+
compile(registry: CompileRegistry): Promise<ArtifactRecord | void>;
66
}

‎packages/protocol/src/compiling/CompileRegistry.ts‎

Lines changed: 31 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,11 @@
11
import { inject, injectable, singleton } from "tsyringe";
22
import {
33
AreProofsEnabled,
4+
CompileArtifact,
45
mapSequential,
5-
MOCK_VERIFICATION_KEY,
66
} from "@proto-kit/common";
77

8-
import {
9-
Artifact,
10-
AtomicCompileHelper,
11-
GenericCompileTarget,
12-
} from "./AtomicCompileHelper";
8+
import { ArtifactRecord, AtomicCompileHelper } from "./AtomicCompileHelper";
139
import { CompilableModule } from "./CompilableModule";
1410

1511
/**
@@ -29,61 +25,60 @@ export class CompileRegistry {
2925

3026
public compile: AtomicCompileHelper;
3127

32-
private artifacts: Record<string, Artifact | "compiled-but-no-artifact"> = {};
28+
private artifacts: ArtifactRecord = {};
3329

34-
public async compileModule<ReturnArtifact extends Artifact>(
35-
// TODO Make name inferred by the module token
36-
name: string,
30+
public async compileModule(
3731
compile: (
38-
registry: CompileRegistry,
32+
compiler: AtomicCompileHelper,
3933
...args: unknown[]
40-
) => GenericCompileTarget<ReturnArtifact>,
34+
) => Promise<ArtifactRecord>,
4135
dependencies: Record<string, CompilableModule> = {}
42-
): Promise<ReturnArtifact | undefined> {
36+
): Promise<ArtifactRecord | undefined> {
4337
const collectedArtifacts = await mapSequential(
4438
Object.entries(dependencies),
4539
async ([depName, dep]) => {
4640
if (this.artifacts[depName] !== undefined) {
4741
return this.artifacts[depName];
4842
}
49-
const artifact =
50-
(await dep.compile(this)) ?? "compiled-but-no-artifact";
51-
this.artifacts[depName] = artifact;
43+
const artifact = await dep.compile(this);
44+
if (artifact !== undefined) {
45+
this.artifacts = {
46+
...this.artifacts,
47+
...artifact,
48+
};
49+
}
5250
return artifact;
5351
}
5452
);
5553

56-
const target = compile(this, ...collectedArtifacts);
57-
const artifact = await this.compile.program<ReturnArtifact>(name, target);
54+
const artifacts = await compile(this.compile, ...collectedArtifacts);
5855

59-
this.artifacts[name] = artifact ?? "compiled-but-no-artifact";
56+
this.artifacts = {
57+
...this.artifacts,
58+
...artifacts,
59+
};
6060

61-
return artifact;
61+
return artifacts;
6262
}
6363

64-
public getArtifact<ArtifactType extends Artifact>(name: string) {
64+
public getArtifact(name: string) {
6565
if (this.artifacts[name] === undefined) {
6666
throw new Error(
6767
`Artifact for ${name} not available, did you compile it via the CompileRegistry?`
6868
);
6969
}
70-
if (!this.areProofsEnabled.areProofsEnabled) {
71-
return MOCK_VERIFICATION_KEY;
72-
}
7370

74-
const artifact = this.artifacts[name];
75-
if (artifact === "compiled-but-no-artifact") {
76-
throw new Error(
77-
`Module ${name} didn't return the requested artifact even though proofs are enabled`
78-
);
79-
}
80-
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
81-
return artifact as ArtifactType;
71+
return this.artifacts[name];
72+
}
73+
74+
public addArtifactsRaw(artifacts: ArtifactRecord) {
75+
this.artifacts = {
76+
...this.artifacts,
77+
...artifacts,
78+
};
8279
}
8380

84-
public addArtifactsRaw(artifacts: Record<string, Artifact>) {
85-
Object.entries(artifacts).forEach(([key, value]) => {
86-
this.artifacts[key] = value;
87-
});
81+
public getAllArtifacts() {
82+
return this.artifacts;
8883
}
8984
}

‎packages/protocol/src/prover/block/BlockProver.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -929,8 +929,7 @@ export class BlockProver
929929
registry: CompileRegistry
930930
): Promise<Record<string, CompileArtifact> | undefined> {
931931
return await registry.compileModule(
932-
"BlockProver",
933-
() => this.zkProgrammable
932+
async () => await this.zkProgrammable.compile()
934933
);
935934
}
936935

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ import {
3030
import { StateTransitionWitnessProvider } from "./StateTransitionWitnessProvider";
3131
import { StateTransitionWitnessProviderReference } from "./StateTransitionWitnessProviderReference";
3232
import { CompilableModule } from "../../compiling/CompilableModule";
33-
import { Artifact } from "../../compiling/AtomicCompileHelper";
33+
import { ArtifactRecord } from "../../compiling/AtomicCompileHelper";
3434
import { CompileRegistry } from "../../compiling/CompileRegistry";
3535

3636
const errors = {
@@ -367,10 +367,9 @@ export class StateTransitionProver
367367
);
368368
}
369369

370-
public compile(registry: CompileRegistry): Promise<void | Artifact> {
370+
public compile(registry: CompileRegistry): Promise<void | ArtifactRecord> {
371371
return registry.compileModule(
372-
"StateTransitionProver",
373-
() => this.zkProgrammable
372+
async () => await this.zkProgrammable.compile()
374373
);
375374
}
376375

‎packages/protocol/src/settlement/ContractModule.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import {
77
import { SmartContract } from "o1js";
88
import { CompilableModule } from "../compiling/CompilableModule";
99
import type { CompileRegistry } from "../compiling/CompileRegistry";
10-
import { Artifact } from "../compiling/AtomicCompileHelper";
10+
import { ArtifactRecord } from "../compiling/AtomicCompileHelper";
1111

1212
export type SmartContractClassFromInterface<Type> = typeof SmartContract &
1313
TypedClass<Type>;
@@ -28,5 +28,5 @@ export abstract class ContractModule<ContractType, Config = NoConfig>
2828

2929
public abstract compile(
3030
registry: CompileRegistry
31-
): Promise<Artifact | undefined>;
31+
): Promise<ArtifactRecord | undefined>;
3232
}

‎packages/protocol/src/settlement/contracts/BridgeContractProtocolModule.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ export class BridgeContractProtocolModule extends ContractModule<
3636
}
3737

3838
public async compile(registry: CompileRegistry) {
39-
return await registry.compileModule("BridgeContract", () => BridgeContract);
39+
return await registry.compileModule(async (compiler) => ({
40+
BridgeContract: await BridgeContract.compile(),
41+
}));
4042
}
4143
}

‎packages/protocol/src/settlement/contracts/DispatchContractProtocolModule.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,10 @@ export class DispatchContractProtocolModule extends ContractModule<
4545
}
4646

4747
public async compile(registry: CompileRegistry) {
48-
return await registry.compileModule(
49-
"DispatchSmartContract",
50-
() => DispatchSmartContract
51-
);
48+
return await registry.compileModule(async (compiler) => ({
49+
DispatchSmartContract: await compiler.compileContract(
50+
DispatchSmartContract
51+
),
52+
}));
5253
}
5354
}

‎packages/protocol/src/settlement/contracts/SettlementContractProtocolModule.ts‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import {
88
} from "../ContractModule";
99
import { ProvableSettlementHook } from "../modularity/ProvableSettlementHook";
1010
import { CompileRegistry } from "../../compiling/CompileRegistry";
11-
import { Artifact } from "../../compiling/AtomicCompileHelper";
11+
import { ArtifactRecord } from "../../compiling/AtomicCompileHelper";
1212

1313
import { DispatchSmartContractBase } from "./DispatchSmartContract";
1414
import {
@@ -78,23 +78,22 @@ export class SettlementContractProtocolModule extends ContractModule<
7878

7979
public async compile(
8080
registry: CompileRegistry
81-
): Promise<Artifact | undefined> {
81+
): Promise<ArtifactRecord | undefined> {
8282
return await registry.compileModule(
83-
"SettlementContract",
84-
(
85-
registry2: CompileRegistry,
86-
bridgeVk: unknown,
87-
blockProverVk: unknown
88-
) => {
83+
async (compiler, bridgeVk: unknown, blockProverVk: unknown) => {
8984
SettlementSmartContractBase.args.BridgeContractVerificationKey =
9085
// TODO Infer type
9186
bridgeVk as VerificationKey;
9287

93-
return SettlementSmartContract;
88+
return {
89+
SettlementSmartContract: await compiler.compileContract(
90+
SettlementSmartContract
91+
),
92+
};
9493
},
9594
{
9695
BridgeContract: this.bridgeContractModule,
97-
BlockProver: this.blockProver.zkProgrammable,
96+
BlockProver: this.blockProver,
9897
}
9998
);
10099
}

0 commit comments

Comments
 (0)