Merge pull request #1874 from o1-labs/2024-10-bugfix-lb-performance

Build fix for lagrange basis performance

CHANGELOG.md

@@ -19,6 +19,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Fixes
 
+- Performance regression when compiling recursive circuits is fixed https://github.com/o1-labs/o1js/pull/1874
 - Decouple offchain state instances from their definitions https://github.com/o1-labs/o1js/pull/1834
 
 ## [1.9.0](https://github.com/o1-labs/o1js/compare/450943...f15293a69) - 2024-10-15

src/lib/proof-system/cache.ts

@@ -11,7 +11,13 @@ import { jsEnvironment } from '../../bindings/crypto/bindings/env.js';
 export { Cache, CacheHeader };
 
 // internal API
-export { readCache, writeCache, withVersion, cacheHeaderVersion };
+export {
+  readCache,
+  writeCache,
+  withVersion,
+  cacheHeaderVersion,
+  LAGRANGE_BASIS_PREFIX,
+};
 
 /**
  * Interface for storing and retrieving values, for caching.
@@ -90,6 +96,8 @@ type StepKeyHeader<Kind> = {
 type WrapKeyHeader<Kind> = { kind: Kind; programName: string; hash: string };
 type PlainHeader<Kind> = { kind: Kind };
 
+const LAGRANGE_BASIS_PREFIX = 'lagrange-basis' as const;
+
 /**
  * A header that is passed to the caching layer, to support rich caching strategies.
  *
@@ -101,7 +109,7 @@ type CacheHeader = (
   | WrapKeyHeader<'wrap-pk'>
   | WrapKeyHeader<'wrap-vk'>
   | PlainHeader<'srs'>
-  | PlainHeader<'lagrange-basis'>
+  | PlainHeader<typeof LAGRANGE_BASIS_PREFIX>
 ) &
   CommonHeader;
 

src/lib/proof-system/cached-lagrange-basis.unit-test.ts

@@ -0,0 +1,65 @@
+import { Cache, LAGRANGE_BASIS_PREFIX } from './cache.js';
+import { SelfProof, ZkProgram } from './zkprogram.js';
+import { Field } from '../provable/field.js';
+import { it, describe, after, before } from 'node:test';
+import { expect } from 'expect';
+import { promises as fs } from 'fs';
+
+const __cacheDirname = './.tmpcache';
+
+const exampleProgram = ZkProgram({
+  name: 'example',
+  publicOutput: Field,
+  methods: {
+    init: {
+      privateInputs: [],
+      async method() {
+        return new Field(0);
+      },
+    },
+    run: {
+      privateInputs: [SelfProof],
+      async method(p: SelfProof<undefined, Field>) {
+        return p.publicOutput.add(new Field(1));
+      },
+    },
+  },
+});
+
+describe('Compiling a program with a cache', () => {
+  const cache: Cache & { lagrangeBasisReadCount?: number } =
+    Cache.FileSystem(__cacheDirname);
+  const originalRead = cache.read;
+  cache.lagrangeBasisReadCount = 0;
+  cache.read = ({ persistentId, uniqueId, dataType }) => {
+    if (persistentId.startsWith(LAGRANGE_BASIS_PREFIX)) {
+      const readCount = cache.lagrangeBasisReadCount || 0;
+      cache.lagrangeBasisReadCount = readCount + 1;
+    }
+    return originalRead({ persistentId, uniqueId, dataType } as any);
+  };
+
+  before(async () => {
+    await fs.mkdir(__cacheDirname, { recursive: true });
+  });
+
+  after(async () => {
+    await fs.rm(__cacheDirname, { recursive: true });
+  });
+
+  /**
+   * This test is a regression test for https://github.com/o1-labs/o1js/issues/1869
+   * It ensures that the lagrange basis cache is accessed properly.  If the file system cache is not
+   * read during compile, that means that the lagrange basis was returned from WASM on the first attempt.
+   *
+   * This is not necessarily a problem.  If the WASM code is updated such that we expect the LB to be
+   * returned on the first try, and we explicitly skip the file system cache, then this test can be
+   * safely removed.  Otherwise, a failure here probably indicates a performance regression.
+   */
+  it('should attempt to read lagrange basis from the cache during compile', async () => {
+    cache.lagrangeBasisReadCount = 0;
+    await exampleProgram.compile({ cache });
+    expect(cache.lagrangeBasisReadCount).not.toBe(0);
+    cache.lagrangeBasisReadCount = 0;
+  });
+});