@cofhe/sdk 0.7.0 → 0.8.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (79) hide show
  1. package/CHANGELOG.md +37 -0
  2. package/acps/acp.ts +10 -1
  3. package/acps/index.ts +11 -0
  4. package/acps/onchain-utils.ts +188 -49
  5. package/acps/registry.ts +73 -0
  6. package/acps/test/onchainAccessStatus.test.ts +116 -0
  7. package/acps/types.ts +16 -0
  8. package/chains/chains/arbitrum.ts +14 -0
  9. package/chains/chains/ethereum.ts +14 -0
  10. package/chains/index.ts +5 -1
  11. package/chains/test/chains.test.ts +24 -3
  12. package/core/acps.ts +51 -143
  13. package/core/client.ts +13 -8
  14. package/core/clientTypes.ts +29 -4
  15. package/core/config.ts +25 -1
  16. package/core/consts.ts +14 -2
  17. package/core/decrypt/verifyDecryptResult.ts +4 -3
  18. package/core/encrypt/encryptInputsBuilder.ts +6 -0
  19. package/core/encrypt/zkPackProveVerify.ts +7 -1
  20. package/core/error.ts +1 -0
  21. package/core/fetchKeys.ts +2 -2
  22. package/core/index.ts +9 -1
  23. package/core/keyStore.ts +23 -10
  24. package/core/taskManager.ts +48 -0
  25. package/core/test/acps.test.ts +91 -1
  26. package/core/test/client.test.ts +41 -0
  27. package/core/test/config.test.ts +16 -0
  28. package/core/test/decrypt.test.ts +11 -4
  29. package/core/test/encryptInputsBuilder.test.ts +1 -1
  30. package/core/test/fetchKeys.test.ts +23 -2
  31. package/core/test/keyStore.test.ts +18 -6
  32. package/core/test/taskManager.test.ts +57 -0
  33. package/core/types.ts +8 -0
  34. package/dist/acps.cjs +314 -35
  35. package/dist/acps.d.cts +251 -4
  36. package/dist/acps.d.ts +251 -4
  37. package/dist/acps.js +2 -2
  38. package/dist/chains.cjs +26 -0
  39. package/dist/chains.d.cts +45 -1
  40. package/dist/chains.d.ts +45 -1
  41. package/dist/chains.js +1 -1
  42. package/dist/{chunk-43USWPEH.js → chunk-GGI7Y6AW.js} +59 -264
  43. package/dist/chunk-KSOV5B4B.js +50 -0
  44. package/dist/{chunk-ESMZCFJY.js → chunk-NMKMCAI3.js} +4 -2
  45. package/dist/{chunk-Q7CBWGQX.js → chunk-UJYOBHYA.js} +374 -37
  46. package/dist/{chunk-N6IDQRRU.js → chunk-UKQAJK4N.js} +25 -1
  47. package/dist/{clientTypes-BN3nbzYM.d.ts → clientTypes-5YPPR16m.d.ts} +77 -14
  48. package/dist/{clientTypes-CYZjFznO.d.cts → clientTypes-D3aIXmUp.d.cts} +77 -14
  49. package/dist/core.cjs +271 -136
  50. package/dist/core.d.cts +29 -8
  51. package/dist/core.d.ts +29 -8
  52. package/dist/core.js +4 -4
  53. package/dist/node.cjs +266 -136
  54. package/dist/node.d.cts +2 -2
  55. package/dist/node.d.ts +2 -2
  56. package/dist/node.js +5 -5
  57. package/dist/{acp-Wi6isVQI.d.cts → onchain-utils-CUAK8glx.d.cts} +31 -1
  58. package/dist/{acp-Wi6isVQI.d.ts → onchain-utils-CUAK8glx.d.ts} +31 -1
  59. package/dist/web.cjs +360 -150
  60. package/dist/web.d.cts +49 -3
  61. package/dist/web.d.ts +49 -3
  62. package/dist/web.js +51 -20
  63. package/dist/zkProve.worker.cjs +69 -15
  64. package/dist/zkProve.worker.js +23 -17
  65. package/node/storage.ts +2 -2
  66. package/node/test/inherited.test.ts +5 -1
  67. package/node/test/storage-memory-fallback.test.ts +47 -0
  68. package/package.json +5 -3
  69. package/web/index.ts +100 -26
  70. package/web/test/countingZkProveWorker.ts +31 -0
  71. package/web/test/inherited.web.test.ts +5 -1
  72. package/web/test/mainThreadPool.web.test.ts +174 -0
  73. package/web/test/threadPoolBenchHelper.ts +140 -0
  74. package/web/test/threadPoolOff.web.test.ts +17 -0
  75. package/web/test/threadPoolOn.web.test.ts +80 -0
  76. package/web/test/workerInitRace.web.test.ts +79 -0
  77. package/web/tfheThreadPool.ts +140 -0
  78. package/web/workerManager.ts +4 -2
  79. package/web/zkProve.worker.ts +42 -17
package/web/index.ts CHANGED
@@ -9,6 +9,8 @@ import {
9
9
  type ZkBuilderAndCrsGenerator,
10
10
  type FheKeyDeserializer,
11
11
  type EncryptableItem,
12
+ type TfheInitializer,
13
+ type TfheThreadsSetting,
12
14
  fheTypeToString,
13
15
  TFHE_RS_SAFE_SERIALIZATION_SIZE_LIMIT,
14
16
  } from '@/core';
@@ -25,21 +27,86 @@ import { getWorkerManager, terminateWorker, areWorkersAvailable } from './worker
25
27
  // reference `self` at module top — into the import graph during Next.js SSR.
26
28
  import type { TfheCompactPublicKey, ProvenCompactCiphertextList, CompactPkeCrs } from 'tfhe';
27
29
  import { hasDOM } from './const';
30
+ import { initTfheThreadPool, type TfheThreadPoolResult } from './tfheThreadPool.js';
28
31
 
29
- /**
30
- * Internal function to initialize TFHE for web
31
- * Called automatically on first encryption - users don't need to call this manually
32
- * @returns true if TFHE was initialized, false if already initialized
33
- */
34
32
  let tfheModule: typeof import('tfhe') | null = null;
35
33
  let tfheInitialized = false;
36
- async function initTfhe(): Promise<boolean> {
34
+ let tfheInitPromise: Promise<void> | null = null;
35
+ let threadPoolPromise: Promise<TfheThreadPoolResult> | null = null;
36
+ let threadPoolResult: TfheThreadPoolResult | null = null;
37
+
38
+ /**
39
+ * Outcome of the rayon thread pool setup on the main thread, or `null` if the
40
+ * main thread hasn't had to generate a proof yet. `enabled: false` carries a
41
+ * `reason`.
42
+ *
43
+ * Note this reflects the *main thread* only. When workers are enabled (the
44
+ * default) proving happens inside the zkProve worker, which runs its own pool,
45
+ * and the main thread never starts one — so `null` is the normal value there.
46
+ */
47
+ export function getTfheThreadPoolStatus(): TfheThreadPoolResult | null {
48
+ return threadPoolResult;
49
+ }
50
+
51
+ /**
52
+ * Internal TFHE initializer used on web.
53
+ * Called automatically on first encryption - users don't need to call this manually.
54
+ *
55
+ * The wasm instance is a module-level singleton, so initialization is memoised.
56
+ * It deliberately does NOT start the rayon thread pool: the main thread only
57
+ * needs one when it generates a proof itself, see `createStartMainThreadPool`.
58
+ *
59
+ * @returns true if it performed the initialization, false if TFHE was already initialized
60
+ */
61
+ const initTfhe: TfheInitializer = async () => {
37
62
  if (tfheInitialized) return false;
38
- tfheModule = await import('tfhe');
39
- await tfheModule.default();
40
- await tfheModule.init_panic_hook();
41
- tfheInitialized = true;
63
+
64
+ if (tfheInitPromise) {
65
+ // Another caller is already initializing; wait it out but don't claim
66
+ // credit for having done it.
67
+ await tfheInitPromise;
68
+ return false;
69
+ }
70
+
71
+ tfheInitPromise = (async () => {
72
+ const mod = await import('tfhe');
73
+ await mod.default();
74
+ await mod.init_panic_hook();
75
+ tfheModule = mod;
76
+ tfheInitialized = true;
77
+ })();
78
+
79
+ try {
80
+ await tfheInitPromise;
81
+ } catch (error) {
82
+ // Let the next call retry from scratch.
83
+ tfheInitPromise = null;
84
+ throw error;
85
+ }
86
+
42
87
  return true;
88
+ };
89
+
90
+ /**
91
+ * Starts the main thread's rayon pool right before the main thread generates a
92
+ * proof (workers disabled, unavailable, or the worker failed). With workers on
93
+ * the proof runs in the zkProve worker, which has its own pool, so a second one
94
+ * here would just be idle Web Workers plus another shared wasm memory.
95
+ *
96
+ * Must run before the FIRST main-thread proof: rayon locks in its global pool
97
+ * on first use, and `initThreadPool` fails from then on. Deserializing keys and
98
+ * packing inputs beforehand is fine. Memoised because `initThreadPool` may only
99
+ * run once per wasm instance — if an app creates several clients with different
100
+ * `tfheThreads` values, whichever proves on the main thread first wins.
101
+ *
102
+ * Best-effort: leaves tfhe single-threaded when the page isn't cross-origin
103
+ * isolated rather than failing the encryption.
104
+ */
105
+ function createStartMainThreadPool(tfheThreads: TfheThreadsSetting) {
106
+ return async (): Promise<void> => {
107
+ threadPoolPromise ??= initTfheThreadPool(requireTfhe(), tfheThreads);
108
+ threadPoolResult = await threadPoolPromise;
109
+ };
43
110
  }
44
111
 
45
112
  function requireTfhe(): typeof import('tfhe') {
@@ -101,22 +168,27 @@ const zkBuilderAndCrsGenerator: ZkBuilderAndCrsGenerator = (fhe: string, crs: st
101
168
  /**
102
169
  * Worker-enabled zkProve function
103
170
  * This submits proof generation to a Web Worker
171
+ *
172
+ * Bound to the client's `tfheThreads` setting, which the worker applies when it
173
+ * first initialises tfhe.
104
174
  */
105
- async function zkProveWithWorker(
106
- fheKeyHex: string,
107
- crsHex: string,
108
- items: EncryptableItem[],
109
- metadata: Uint8Array
110
- ): Promise<Uint8Array> {
111
- // Serialize items for worker (convert enum to string name)
112
- const serializedItems = items.map((item) => ({
113
- utype: fheTypeToString(item.utype),
114
- data: typeof item.data === 'bigint' ? item.data.toString() : item.data,
115
- }));
116
-
117
- // Submit to worker
118
- const workerManager = getWorkerManager();
119
- return await workerManager.submitProof(fheKeyHex, crsHex, serializedItems, metadata);
175
+ function createZkProveWithWorker(tfheThreads: TfheThreadsSetting) {
176
+ return async function zkProveWithWorker(
177
+ fheKeyHex: string,
178
+ crsHex: string,
179
+ items: EncryptableItem[],
180
+ metadata: Uint8Array
181
+ ): Promise<Uint8Array> {
182
+ // Serialize items for worker (convert enum to string name)
183
+ const serializedItems = items.map((item) => ({
184
+ utype: fheTypeToString(item.utype),
185
+ data: typeof item.data === 'bigint' ? item.data.toString() : item.data,
186
+ }));
187
+
188
+ // Submit to worker
189
+ const workerManager = getWorkerManager();
190
+ return await workerManager.submitProof(fheKeyHex, crsHex, serializedItems, metadata, tfheThreads);
191
+ };
120
192
  }
121
193
 
122
194
  /**
@@ -147,9 +219,10 @@ export function createCofheClient<TConfig extends CofheConfig>(config: TConfig):
147
219
  tfhePublicKeyDeserializer,
148
220
  compactPkeCrsDeserializer,
149
221
  initTfhe,
222
+ beforeMainThreadProve: createStartMainThreadPool(config.tfheThreads),
150
223
  // Always provide the worker function if available - config.useWorkers controls usage
151
224
  // areWorkersAvailable will return true if the Worker API is available and false in Node.js
152
- zkProveWorkerFn: areWorkersAvailable() ? zkProveWithWorker : undefined,
225
+ zkProveWorkerFn: areWorkersAvailable() ? createZkProveWithWorker(config.tfheThreads) : undefined,
153
226
  });
154
227
  }
155
228
 
@@ -182,6 +255,7 @@ export function createCofheClientWithCustomWorker(
182
255
  tfhePublicKeyDeserializer,
183
256
  compactPkeCrsDeserializer,
184
257
  initTfhe,
258
+ beforeMainThreadProve: createStartMainThreadPool(config.tfheThreads),
185
259
  zkProveWorkerFn: customZkProveWorkerFn,
186
260
  });
187
261
  }
@@ -0,0 +1,31 @@
1
+ /**
2
+ * Test-only wrapper around the real zkProve worker.
3
+ *
4
+ * tfhe's rayon threads are Web Workers that the zkProve worker creates itself,
5
+ * so the page never sees them. This wrapper counts them from the inside: it
6
+ * swaps this worker's global `Worker` for a counting subclass, then loads the
7
+ * real `zkProve.worker` into the same global scope. The zkProve protocol is
8
+ * untouched; one extra message type reports the count.
9
+ */
10
+
11
+ /* eslint-disable no-undef */
12
+
13
+ const scope = self as any;
14
+ const RealWorker = scope.Worker;
15
+ let rayonWorkersStarted = 0;
16
+
17
+ scope.Worker = class extends RealWorker {
18
+ constructor(url: string | URL, options?: WorkerOptions) {
19
+ super(url, options);
20
+ if (String(url).includes('workerHelpers')) rayonWorkersStarted++;
21
+ }
22
+ };
23
+
24
+ scope.addEventListener('message', (event: MessageEvent) => {
25
+ if (event.data?.type !== 'countRayonWorkers') return;
26
+ scope.postMessage({ id: event.data.id, type: 'rayonWorkerCount', count: rayonWorkersStarted });
27
+ });
28
+
29
+ // Loaded after the counting `Worker` is in place. The real worker installs its
30
+ // own `onmessage` handler and signals `ready` as usual.
31
+ import('../zkProve.worker');
@@ -17,6 +17,10 @@ import { createPublicClient, createWalletClient, http } from 'viem';
17
17
  import { privateKeyToAccount } from 'viem/accounts';
18
18
  import { arbitrumSepolia, baseSepolia, sepolia } from 'viem/chains';
19
19
 
20
+ // Live-chain suite: needs the funded TEST_PRIVATE_KEY (transacts against pre-deployed
21
+ // testnet fixtures). Without it — fork PRs, where secrets are unavailable — the suite
22
+ // skips; CI runs it separately with the key (test-live job).
23
+ const LIVE_TESTS = Boolean(TEST_PRIVATE_KEY);
20
24
  const DEFAULT_TEST_PRIVATE_KEY = '0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80';
21
25
  const BOB_PRIVATE_KEY = (TEST_PRIVATE_KEY || DEFAULT_TEST_PRIVATE_KEY) as `0x${string}`;
22
26
  const ALICE_PRIVATE_KEY = '0x59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d';
@@ -33,7 +37,7 @@ const VIEM_CHAINS: Record<number, Chain> = {
33
37
  const testViemChain = STAGING_TESTS ? stagingViemChain : arbitrumSepolia;
34
38
  const testCofheChain = STAGING_TESTS ? stagingCofhe : cofheArbSepolia;
35
39
 
36
- describe('@cofhe/web - Inherited Client Tests', () => {
40
+ describe.skipIf(!LIVE_TESTS)('@cofhe/web - Inherited Client Tests', () => {
37
41
  let cofheClient: CofheClient;
38
42
  let publicClient: PublicClient;
39
43
  let bobWalletClient: WalletClient;
@@ -0,0 +1,174 @@
1
+ import { STAGING_TESTS, stagingViemChain } from '../../core/test/stagingRedirect';
2
+ import { arbSepolia as cofheArbSepolia, stagingCofhe } from '@/chains';
3
+ import { Encryptable, fheTypeToString, type EncryptableItem } from '@/core';
4
+
5
+ import { describe, it, expect, beforeAll, afterAll } from 'vitest';
6
+ import type { PublicClient, WalletClient } from 'viem';
7
+ import { createPublicClient, createWalletClient, http } from 'viem';
8
+ import { privateKeyToAccount } from 'viem/accounts';
9
+ import { arbitrumSepolia as viemArbitrumSepolia } from 'viem/chains';
10
+ import {
11
+ createCofheClient,
12
+ createCofheConfig,
13
+ createCofheClientWithCustomWorker,
14
+ getTfheThreadPoolStatus,
15
+ } from '../index';
16
+ import { THREADPOOL_TESTS } from './threadPoolBenchHelper';
17
+
18
+ // Which thread needs a rayon pool?
19
+ //
20
+ // With workers on (the default) the ZK proof is generated inside the zkProve
21
+ // worker, which runs its own rayon pool. The main thread only deserializes keys
22
+ // and packs inputs, so a pool there is pure overhead: N idle Web Workers plus a
23
+ // second shared wasm memory. The main thread should start one only when it has
24
+ // to generate a proof itself.
25
+ //
26
+ // Rayon threads started by the main thread are `new Worker(...)` calls made by
27
+ // the page. The ones the zkProve worker starts are made inside that worker and
28
+ // never reach the page's `Worker` constructor. Counting page-level helper
29
+ // workers therefore measures the main-thread pool and nothing else.
30
+ //
31
+ // Own file on purpose: every vitest browser file gets a fresh tfhe wasm
32
+ // instance, and a rayon pool can only be started once per instance. The tests
33
+ // below share that instance and depend on their order.
34
+
35
+ const TEST_PRIVATE_KEY = '0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80';
36
+ const THREADS = 2;
37
+
38
+ const testViemChain = STAGING_TESTS ? stagingViemChain : viemArbitrumSepolia;
39
+ const testCofheChain = STAGING_TESTS ? stagingCofhe : cofheArbSepolia;
40
+
41
+ /**
42
+ * Drives the real zkProve worker. The SDK's worker manager points at the built
43
+ * `zkProve.worker.js`, which doesn't exist when tests run from source, so the
44
+ * default client always falls back to the main thread here. Loading the worker
45
+ * from its `.ts` source gives these tests a worker path that actually works.
46
+ */
47
+ function createRealWorkerProver(tfheThreads: number) {
48
+ const worker = new Worker(new URL('../zkProve.worker.ts', import.meta.url), { type: 'module' });
49
+
50
+ const ready = new Promise<void>((resolve, reject) => {
51
+ worker.addEventListener('message', function onReady(event: MessageEvent) {
52
+ if (event.data?.type !== 'ready') return;
53
+ worker.removeEventListener('message', onReady);
54
+ resolve();
55
+ });
56
+ worker.addEventListener('error', (event) => reject(new Error(`zkProve worker failed to load: ${event.message}`)));
57
+ });
58
+
59
+ let requests = 0;
60
+
61
+ const prove = async (
62
+ fheKeyHex: string,
63
+ crsHex: string,
64
+ items: EncryptableItem[],
65
+ metadata: Uint8Array
66
+ ): Promise<Uint8Array> => {
67
+ await ready;
68
+ const id = `main-thread-pool-test-${requests++}`;
69
+
70
+ return new Promise<Uint8Array>((resolve, reject) => {
71
+ worker.addEventListener('message', function onResponse(event: MessageEvent) {
72
+ if (event.data?.id !== id) return;
73
+ worker.removeEventListener('message', onResponse);
74
+ if (event.data.type === 'success') resolve(new Uint8Array(event.data.result));
75
+ else reject(new Error(event.data.error ?? 'zkProve worker error'));
76
+ });
77
+
78
+ worker.postMessage({
79
+ id,
80
+ type: 'zkProve',
81
+ fheKeyHex,
82
+ crsHex,
83
+ items: items.map((item) => ({
84
+ utype: fheTypeToString(item.utype),
85
+ data: typeof item.data === 'bigint' ? item.data.toString() : item.data,
86
+ })),
87
+ metadata: Array.from(metadata),
88
+ tfheThreads,
89
+ });
90
+ });
91
+ };
92
+
93
+ return { prove, terminate: () => worker.terminate() };
94
+ }
95
+
96
+ describe.skipIf(!THREADPOOL_TESTS)('@cofhe/sdk/web - main-thread rayon pool', () => {
97
+ const RealWorker = globalThis.Worker;
98
+ let rayonWorkersStartedByPage = 0;
99
+ let prover: ReturnType<typeof createRealWorkerProver>;
100
+ let publicClient: PublicClient;
101
+ let walletClient: WalletClient;
102
+ let consumingContract: `0x${string}`;
103
+ let proveContext: any;
104
+
105
+ beforeAll(async () => {
106
+ (globalThis as any).Worker = class extends RealWorker {
107
+ constructor(url: string | URL, options?: WorkerOptions) {
108
+ super(url, options);
109
+ if (String(url).includes('workerHelpers')) rayonWorkersStartedByPage++;
110
+ }
111
+ };
112
+
113
+ publicClient = createPublicClient({ chain: testViemChain, transport: http() });
114
+ const account = privateKeyToAccount(TEST_PRIVATE_KEY);
115
+ walletClient = createWalletClient({ chain: testViemChain, transport: http(), account });
116
+ consumingContract = account.address;
117
+
118
+ prover = createRealWorkerProver(THREADS);
119
+
120
+ const config = createCofheConfig({ supportedChains: [testCofheChain], tfheThreads: THREADS });
121
+ const client = createCofheClientWithCustomWorker(config, prover.prove);
122
+ await client.connect(publicClient, walletClient);
123
+
124
+ await client
125
+ .encryptInputs([Encryptable.uint32(7n)])
126
+ .setConsumingContract(consumingContract)
127
+ .onStep((step, context) => {
128
+ if (step === 'prove' && context?.isEnd) proveContext = context;
129
+ })
130
+ .execute();
131
+ }, 240000);
132
+
133
+ afterAll(() => {
134
+ (globalThis as any).Worker = RealWorker;
135
+ prover?.terminate();
136
+ });
137
+
138
+ it('generates the proof in the worker', () => {
139
+ expect(globalThis.crossOriginIsolated).toBe(true);
140
+ expect(proveContext?.usedWorker).toBe(true);
141
+ expect(proveContext?.workerFailedError).toBeUndefined();
142
+ });
143
+
144
+ it('does not start a rayon pool on the main thread when the worker proves', () => {
145
+ expect(rayonWorkersStartedByPage).toBe(0);
146
+ expect(getTfheThreadPoolStatus()).toBeNull();
147
+ });
148
+
149
+ // Runs after a worker-path encryption on the same wasm instance: the main
150
+ // thread has already deserialized keys and packed inputs by now, and the
151
+ // pool must still come up.
152
+ it('starts the pool once the main thread has to prove', async () => {
153
+ const config = createCofheConfig({
154
+ supportedChains: [testCofheChain],
155
+ tfheThreads: THREADS,
156
+ useWorkers: false,
157
+ });
158
+ const client = createCofheClient(config);
159
+ await client.connect(publicClient, walletClient);
160
+
161
+ let mainThreadProveContext: any;
162
+ await client
163
+ .encryptInputs([Encryptable.uint32(8n)])
164
+ .setConsumingContract(consumingContract)
165
+ .onStep((step, context) => {
166
+ if (step === 'prove' && context?.isEnd) mainThreadProveContext = context;
167
+ })
168
+ .execute();
169
+
170
+ expect(mainThreadProveContext?.usedWorker).toBe(false);
171
+ expect(rayonWorkersStartedByPage).toBe(THREADS);
172
+ expect(getTfheThreadPoolStatus()).toEqual({ enabled: true, threads: THREADS });
173
+ }, 240000);
174
+ });
@@ -0,0 +1,140 @@
1
+ /**
2
+ * Shared harness for the initThreadPool verification tests.
3
+ *
4
+ * Each vitest browser test file runs in its own iframe, so each gets a fresh
5
+ * tfhe wasm instance — which is what lets us compare single-threaded vs
6
+ * multi-threaded proving (`initThreadPool` may only run once per instance).
7
+ */
8
+
9
+ import { TFHE_RS_SAFE_SERIALIZATION_SIZE_LIMIT } from '../../core/consts';
10
+ import { STAGING_TESTS } from '../../core/test/stagingRedirect';
11
+ import { arbSepolia as cofheArbSepolia, stagingCofhe } from '@/chains';
12
+ import { initTfheThreadPool, type TfheThreadPoolResult } from '../tfheThreadPool.js';
13
+
14
+ // Same backend as the rest of the web suite: Arbitrum Sepolia, or staging when
15
+ // TEST_STAGING_ENABLED=true.
16
+ const COFHE_URL = (STAGING_TESTS ? stagingCofhe : cofheArbSepolia).coFheUrl;
17
+
18
+ // Injected by vitest.config.mts from TEST_THREADPOOL_ENABLED.
19
+ declare const __THREADPOOL_TESTS__: boolean | undefined;
20
+
21
+ /**
22
+ * Gate for the suites that generate real proofs on a rayon pool. They are
23
+ * CPU-heavy: inside the regular, fully parallel test run they starve other
24
+ * proving tests into timeouts on small CI runners, and skew their own timings.
25
+ * So they are skipped by default and run on their own, one file at a time:
26
+ *
27
+ * TEST_THREADPOOL_ENABLED=true pnpm test:threadpool
28
+ *
29
+ * which is what the "Test (thread pool)" CI job does.
30
+ */
31
+ export const THREADPOOL_TESTS = typeof __THREADPOOL_TESTS__ !== 'undefined' ? __THREADPOOL_TESTS__ : false;
32
+
33
+ function fromHexString(hexString: string): Uint8Array {
34
+ const cleanString = hexString.length % 2 === 1 ? `0${hexString}` : hexString;
35
+ const arr = cleanString.replace(/^0x/, '').match(/.{1,2}/g);
36
+ if (!arr) return new Uint8Array();
37
+ return new Uint8Array(arr.map((byte) => parseInt(byte, 16)));
38
+ }
39
+
40
+ /** Fetch the network's real FHE public key + CRS. Proving itself is fully local. */
41
+ export async function fetchNetworkKeys(securityZone = 0): Promise<{ fheKey: string; crs: string }> {
42
+ const post = async (path: string) => {
43
+ const res = await fetch(`${COFHE_URL}/${path}`, {
44
+ method: 'POST',
45
+ headers: { 'Content-Type': 'application/json' },
46
+ body: JSON.stringify({ securityZone }),
47
+ });
48
+ if (!res.ok) throw new Error(`${path} failed: ${res.status} ${res.statusText}`);
49
+ return res.json();
50
+ };
51
+
52
+ const [pk, crs] = await Promise.all([post('GetNetworkPublicKey'), post('GetCrs')]);
53
+ return { fheKey: (pk as { publicKey: string }).publicKey, crs: (crs as { crs: string }).crs };
54
+ }
55
+
56
+ export type BenchResult = {
57
+ threadPool: TfheThreadPoolResult;
58
+ /** Workers constructed during initThreadPool (the rayon threads) */
59
+ rayonWorkers: number;
60
+ /** Per-iteration proving durations in ms */
61
+ durations: number[];
62
+ medianMs: number;
63
+ };
64
+
65
+ /**
66
+ * Init tfhe (optionally with the rayon pool) and time `build_with_proof_packed`.
67
+ *
68
+ * @param threads rayon threads to request; <= 1 skips the pool entirely
69
+ * @param iterations how many proofs to time
70
+ */
71
+ export async function benchProve(threads: number, iterations = 5): Promise<BenchResult> {
72
+ const mod: any = await import('tfhe');
73
+ await mod.default();
74
+ await mod.init_panic_hook();
75
+
76
+ // Count the Workers initThreadPool spawns — direct evidence the rayon threads
77
+ // are real, separate from any timing signal.
78
+ const RealWorker = globalThis.Worker;
79
+ let rayonWorkers = 0;
80
+ (globalThis as any).Worker = class extends RealWorker {
81
+ constructor(url: string | URL, opts?: WorkerOptions) {
82
+ super(url as any, opts);
83
+ rayonWorkers++;
84
+ }
85
+ };
86
+
87
+ let threadPool: TfheThreadPoolResult;
88
+ try {
89
+ threadPool =
90
+ threads > 1
91
+ ? await initTfheThreadPool(mod, threads)
92
+ : { enabled: false, threads: 1, reason: 'skipped by benchmark' };
93
+ } finally {
94
+ (globalThis as any).Worker = RealWorker;
95
+ }
96
+
97
+ const { fheKey, crs } = await fetchNetworkKeys();
98
+ const pk = mod.TfheCompactPublicKey.safe_deserialize(fromHexString(fheKey), TFHE_RS_SAFE_SERIALIZATION_SIZE_LIMIT);
99
+ const zkCrs = mod.CompactPkeCrs.safe_deserialize(fromHexString(crs), TFHE_RS_SAFE_SERIALIZATION_SIZE_LIMIT);
100
+
101
+ // Dummy metadata: we only measure local proof generation, never verify it.
102
+ const metadata = new Uint8Array(32).fill(7);
103
+
104
+ const durations: number[] = [];
105
+ for (let i = 0; i < iterations; i++) {
106
+ const builder = mod.ProvenCompactCiphertextList.builder(pk);
107
+ builder.push_u128(BigInt(1000 + i));
108
+ builder.push_u64(BigInt(2000 + i));
109
+ builder.push_u32(3000 + i);
110
+
111
+ const started = performance.now();
112
+ builder.build_with_proof_packed(zkCrs, metadata, 1);
113
+ durations.push(performance.now() - started);
114
+ }
115
+
116
+ const sorted = [...durations].sort((a, b) => a - b);
117
+ const medianMs = sorted[Math.floor(sorted.length / 2)];
118
+
119
+ return { threadPool, rayonWorkers, durations, medianMs };
120
+ }
121
+
122
+ export function reportBench(label: string, result: BenchResult): void {
123
+ console.log(
124
+ `[THREADPOOL BENCH] ${label} ` +
125
+ JSON.stringify(
126
+ {
127
+ crossOriginIsolated: (globalThis as any).crossOriginIsolated,
128
+ hardwareConcurrency: navigator.hardwareConcurrency,
129
+ poolEnabled: result.threadPool.enabled,
130
+ poolThreads: result.threadPool.threads,
131
+ poolReason: result.threadPool.reason,
132
+ rayonWorkers: result.rayonWorkers,
133
+ medianMs: Math.round(result.medianMs),
134
+ durations: result.durations.map((d) => Math.round(d)),
135
+ },
136
+ null,
137
+ 2
138
+ )
139
+ );
140
+ }
@@ -0,0 +1,17 @@
1
+ import { describe, it, expect } from 'vitest';
2
+ import { THREADPOOL_TESTS, benchProve, reportBench } from './threadPoolBenchHelper.js';
3
+
4
+ // Baseline: same workload, no rayon pool. Kept in its own file because each
5
+ // vitest browser file gets a fresh tfhe wasm instance, and `initThreadPool`
6
+ // can only run once per instance — so single- vs multi-threaded can't be
7
+ // compared inside one file.
8
+ describe.skipIf(!THREADPOOL_TESTS)('tfhe proving — single-threaded baseline', () => {
9
+ it('proves without a rayon thread pool', async () => {
10
+ const result = await benchProve(1);
11
+ reportBench('single-threaded', result);
12
+
13
+ expect(result.threadPool.enabled).toBe(false);
14
+ expect(result.rayonWorkers).toBe(0);
15
+ expect(result.medianMs).toBeGreaterThan(0);
16
+ }, 300000);
17
+ });
@@ -0,0 +1,80 @@
1
+ import { describe, it, expect } from 'vitest';
2
+ import { MAX_AUTO_THREADS, initTfheThreadPool, resolveThreadCount } from '../tfheThreadPool.js';
3
+ import { THREADPOOL_TESTS, benchProve, reportBench } from './threadPoolBenchHelper.js';
4
+
5
+ describe('resolveThreadCount', () => {
6
+ it("derives from hardwareConcurrency for 'auto', capped", () => {
7
+ const expected = Math.min(navigator.hardwareConcurrency, MAX_AUTO_THREADS);
8
+
9
+ expect(resolveThreadCount('auto')).toBe(expected);
10
+ expect(resolveThreadCount()).toBe(expected); // 'auto' is the default
11
+ });
12
+
13
+ it('honours an explicit count, even above the auto cap', () => {
14
+ expect(resolveThreadCount(1)).toBe(1);
15
+ expect(resolveThreadCount(3)).toBe(3);
16
+ // An explicit request is trusted over hardwareConcurrency and the cap.
17
+ expect(resolveThreadCount(MAX_AUTO_THREADS + 4)).toBe(MAX_AUTO_THREADS + 4);
18
+ });
19
+
20
+ it('collapses false and nonsense values to single-threaded', () => {
21
+ expect(resolveThreadCount(false)).toBe(1);
22
+ expect(resolveThreadCount(0)).toBe(1);
23
+ expect(resolveThreadCount(-8)).toBe(1);
24
+ expect(resolveThreadCount(Number.NaN)).toBe(1);
25
+ });
26
+ });
27
+
28
+ describe('initTfheThreadPool', () => {
29
+ it('reports disabled without touching tfhe when config says false', async () => {
30
+ let called = false;
31
+ const result = await initTfheThreadPool(
32
+ {
33
+ initThreadPool: async () => {
34
+ called = true;
35
+ },
36
+ },
37
+ false
38
+ );
39
+
40
+ expect(result.enabled).toBe(false);
41
+ expect(result.threads).toBe(1);
42
+ expect(result.reason).toContain('config.tfheThreads');
43
+ expect(called).toBe(false);
44
+ });
45
+
46
+ it('degrades instead of throwing when initThreadPool fails', async () => {
47
+ const result = await initTfheThreadPool(
48
+ {
49
+ initThreadPool: async () => {
50
+ throw new Error('nested worker blocked');
51
+ },
52
+ },
53
+ 4
54
+ );
55
+
56
+ expect(result.enabled).toBe(false);
57
+ expect(result.threads).toBe(1);
58
+ expect(result.reason).toContain('nested worker blocked');
59
+ });
60
+ });
61
+
62
+ describe.skipIf(!THREADPOOL_TESTS)('tfhe proving — rayon thread pool', () => {
63
+ it('starts the pool and spawns real worker threads', async () => {
64
+ const expectedThreads = resolveThreadCount('auto');
65
+ const result = await benchProve(expectedThreads);
66
+ reportBench('multi-threaded', result);
67
+
68
+ // The page must be cross-origin isolated for this to work at all.
69
+ expect((globalThis as any).crossOriginIsolated).toBe(true);
70
+
71
+ expect(result.threadPool.enabled).toBe(true);
72
+ expect(result.threadPool.threads).toBe(expectedThreads);
73
+
74
+ // Direct evidence the threads are real: initThreadPool constructed one
75
+ // Worker per rayon thread.
76
+ expect(result.rayonWorkers).toBe(expectedThreads);
77
+
78
+ expect(result.medianMs).toBeGreaterThan(0);
79
+ }, 300000);
80
+ });