Merge pull request #2460 from o1-labs/florian/properly-fix-cache

Store all cache artefacts

Author: Trivo25

CHANGELOG.md

@@ -32,6 +32,9 @@ This project adheres to
 - Fixed a verification key regression that was caused by incorrectly enabling a
   proof system feature that wasn't needed.
   https://github.com/o1-labs/o1js/pull/2449
+- Fixed an edge case where not all the artefacts needed for the cache were
+  stored properly, resulting in them being re-computed after loading the cache.
+  https://github.com/o1-labs/o1js/pull/2460
 
 ### Deprecated
 

src/bindings/crypto/bindings/srs.ts

@@ -1,5 +1,5 @@
 import type { Wasm, RustConversion } from '../bindings.js';
-import type { WasmFpSrs, WasmFqSrs } from '../../compiled/node_bindings/plonk_wasm.cjs';
+import { type WasmFpSrs, type WasmFqSrs } from '../../compiled/node_bindings/plonk_wasm.cjs';
 import { PolyComm } from './kimchi-types.js';
 import {
   type CacheHeader,
@@ -24,6 +24,8 @@ function empty(): SrsStore {
 
 const srsStore = { fp: empty(), fq: empty() };
 
+const CacheReadRegister = new Map<string, boolean>();
+
 let cache: Cache | undefined;
 
 function setSrsCache(c: Cache) {
@@ -143,16 +145,15 @@ function srsPerField(f: 'fp' | 'fq', wasm: Wasm, conversion: RustConversion) {
         } else {
           // try to read lagrange basis from cache / recompute and write if not found
           let header = cacheHeaderLagrange(f, domainSize);
-
-          let didRead = readCache(cache, header, (bytes) => {
-            let comms: PolyCommJson[] = JSON.parse(new TextDecoder().decode(bytes));
-            let mlComms = polyCommsFromJSON(comms);
-            let wasmComms = conversion[f].polyCommsToRust(mlComms);
-
-            setLagrangeBasis(srs, domainSize, wasmComms);
-            return true;
-          });
-
+          let didRead = readCacheLazy(
+            cache,
+            header,
+            conversion,
+            f,
+            srs,
+            domainSize,
+            setLagrangeBasis
+          );
           if (didRead !== true) {
             // not in cache
             if (cache.canWrite) {
@@ -167,13 +168,39 @@ function srsPerField(f: 'fp' | 'fq', wasm: Wasm, conversion: RustConversion) {
               lagrangeCommitment(srs, domainSize, i);
             }
           }
-
           // here, basis is definitely stored on the srs
           let c = maybeLagrangeCommitment(srs, domainSize, i);
           assert(c !== undefined, 'commitment exists after setting');
           commitment = c;
         }
       }
+
+      // edge case for when we have a writeable cache and the basis was already stored on the srs
+      // but we didn't store it in the cache seperately yet
+      if (commitment && cache && cache.canWrite) {
+        let header = cacheHeaderLagrange(f, domainSize);
+        let didRead = readCacheLazy(
+          cache,
+          header,
+          conversion,
+          f,
+          srs,
+          domainSize,
+          setLagrangeBasis
+        );
+        // only proceed for entries we haven't written to the cache yet
+        if (didRead !== true) {
+          // same code as above - write the lagrange basis to the cache if it wasn't there already
+          // currently we re-generate the basis via `getLagrangeBasis` - we could derive this from the
+          // already existing `commitment` instead, but this is simpler and the performance impact is negligible
+          let wasmComms = getLagrangeBasis(srs, domainSize);
+          let mlComms = conversion[f].polyCommsFromRust(wasmComms);
+          let comms = polyCommsToJSON(mlComms);
+          let bytes = new TextEncoder().encode(JSON.stringify(comms));
+
+          writeCache(cache, header, bytes);
+        }
+      }
       return conversion[f].polyCommFromRust(commitment);
     },
 
@@ -222,3 +249,24 @@ function polyCommsFromJSON(json: PolyCommJson[]): MlArray<PolyComm> {
     return [0, MlArray.mapTo(shifted, OrInfinity.fromJSON)];
   });
 }
+
+function readCacheLazy(
+  cache: Cache,
+  header: CacheHeader,
+  conversion: RustConversion,
+  f: 'fp' | 'fq',
+  srs: WasmSrs,
+  domainSize: number,
+  setLagrangeBasis: (srs: WasmSrs, domainSize: number, comms: Uint32Array) => void
+) {
+  if (CacheReadRegister.get(header.uniqueId) === true) return true;
+  return readCache(cache, header, (bytes) => {
+    let comms: PolyCommJson[] = JSON.parse(new TextDecoder().decode(bytes));
+    let mlComms = polyCommsFromJSON(comms);
+    let wasmComms = conversion[f].polyCommsToRust(mlComms);
+
+    setLagrangeBasis(srs, domainSize, wasmComms);
+    CacheReadRegister.set(header.uniqueId, true);
+    return true;
+  });
+}