@aztec/foundation 0.69.1 → 0.71.0

package/src/collection/array.ts

@@ -146,6 +146,11 @@ export function maxBy<T>(arr: T[], fn: (x: T) => number): T | undefined {
   return arr.reduce((max, x) => (fn(x) > fn(max) ? x : max), arr[0]);
 }
 
+/** Computes the sum of a numeric array. */
+export function sum(arr: number[]): number {
+  return arr.reduce((a, b) => a + b, 0);
+}
+
 /** Computes the median of a numeric array. Returns undefined if array is empty. */
 export function median(arr: number[]) {
   if (arr.length === 0) {
@@ -155,3 +160,29 @@ export function median(arr: number[]) {
   const mid = Math.floor(sorted.length / 2);
   return sorted.length % 2 !== 0 ? sorted[mid] : (sorted[mid - 1] + sorted[mid]) / 2;
 }
+
+/** Computes the mean of a numeric array. Returns undefined if the array is empty. */
+export function mean(values: number[]) {
+  if (values.length === 0) {
+    return undefined;
+  }
+  return values.reduce((a, b) => a + b, 0) / values.length;
+}
+
+/** Computes the variance of a numeric array. Returns undefined if there are less than 2 points. */
+export function variance(values: number[]) {
+  if (values.length < 2) {
+    return undefined;
+  }
+  const avg = mean(values)!;
+  const points = values.map(value => value * value + avg * avg - 2 * value * avg);
+  return sum(points) / (values.length - 1);
+}
+
+/** Computes the standard deviation of a numeric array. Returns undefined if there are less than 2 points. */
+export function stdDev(values: number[]) {
+  if (values.length < 2) {
+    return undefined;
+  }
+  return Math.sqrt(variance(values)!);
+}

package/src/config/env_var.ts

@@ -64,6 +64,7 @@ export type EnvVar =
   | 'NETWORK'
   | 'NO_PXE'
   | 'COIN_ISSUER_CONTRACT_ADDRESS'
+  | 'USE_GCLOUD_OBSERVABILITY'
   | 'OTEL_EXPORTER_OTLP_METRICS_ENDPOINT'
   | 'OTEL_EXPORTER_OTLP_TRACES_ENDPOINT'
   | 'OTEL_EXPORTER_OTLP_LOGS_ENDPOINT'
@@ -99,6 +100,7 @@ export type EnvVar =
   | 'P2P_TX_PROTOCOL'
   | 'P2P_UDP_ANNOUNCE_ADDR'
   | 'P2P_UDP_LISTEN_ADDR'
+  | 'P2P_ARCHIVED_TX_LIMIT'
   | 'PEER_ID_PRIVATE_KEY'
   | 'PROVER_BLOB_SINK_URL'
   | 'PROOF_VERIFIER_L1_START_BLOCK'
@@ -113,9 +115,9 @@ export type EnvVar =
   | 'PROVER_BROKER_JOB_TIMEOUT_MS'
   | 'PROVER_BROKER_POLL_INTERVAL_MS'
   | 'PROVER_BROKER_JOB_MAX_RETRIES'
-  | 'PROVER_BROKER_DATA_DIRECTORY'
   | 'PROVER_COORDINATION_NODE_URL'
   | 'PROVER_DISABLED'
+  | 'PROVER_FAILED_PROOF_STORE'
   | 'PROVER_ID'
   | 'PROVER_JOB_POLL_INTERVAL_MS'
   | 'PROVER_JOB_TIMEOUT_MS'
@@ -160,7 +162,6 @@ export type EnvVar =
   | 'TX_GOSSIP_VERSION'
   | 'TXE_PORT'
   | 'VALIDATOR_ATTESTATIONS_POLLING_INTERVAL_MS'
-  | 'VALIDATOR_ATTESTATIONS_WAIT_TIMEOUT_MS'
   | 'VALIDATOR_DISABLED'
   | 'VALIDATOR_PRIVATE_KEY'
   | 'VALIDATOR_REEXECUTE'
@@ -189,8 +190,10 @@ export type EnvVar =
   | 'L1_GAS_LIMIT_BUFFER_FIXED'
   | 'L1_GAS_PRICE_MIN'
   | 'L1_GAS_PRICE_MAX'
+  | 'L1_BLOB_FEE_PER_GAS_MAX'
   | 'L1_PRIORITY_FEE_BUMP_PERCENTAGE'
   | 'L1_PRIORITY_FEE_RETRY_BUMP_PERCENTAGE'
+  | 'L1_FIXED_PRIORITY_FEE_PER_GAS'
   | 'L1_TX_MONITOR_MAX_ATTEMPTS'
   | 'L1_TX_MONITOR_CHECK_INTERVAL_MS'
   | 'L1_TX_MONITOR_STALL_TIME_MS'
@@ -201,4 +204,5 @@ export type EnvVar =
   | 'FAUCET_INTERVAL_MS'
   | 'FAUCET_L1_ASSETS'
   | 'K8S_POD_NAME'
-  | 'K8S_POD_UID';
+  | 'K8S_POD_UID'
+  | 'K8S_NAMESPACE_NAME';

package/src/config/index.ts

@@ -1,6 +1,6 @@
 import { type EnvVar } from './env_var.js';
 
-export { EnvVar } from './env_var.js';
+export { type EnvVar } from './env_var.js';
 
 export interface ConfigMapping {
   env?: EnvVar;
@@ -102,7 +102,7 @@ export function optionalNumberConfigHelper(): Pick<ConfigMapping, 'parseEnv'> {
 export function booleanConfigHelper(
   defaultVal = false,
 ): Required<Pick<ConfigMapping, 'parseEnv' | 'defaultValue' | 'isBoolean'> & { parseVal: (val: string) => boolean }> {
-  const parse = (val: string | boolean) => (typeof val === 'boolean' ? val : ['1', 'true', 'TRUE'].includes(val));
+  const parse = (val: string | boolean) => (typeof val === 'boolean' ? val : parseBooleanEnv(val));
   return {
     parseEnv: parse,
     parseVal: parse,
@@ -111,6 +111,11 @@ export function booleanConfigHelper(
   };
 }
 
+/** Parses an env var as boolean. Returns true only if value is 1, true, or TRUE. */
+export function parseBooleanEnv(val: string | undefined): boolean {
+  return val !== undefined && ['1', 'true', 'TRUE'].includes(val);
+}
+
 /**
  * Safely parses a number from a string.
  * If the value is not a number or is not a safe integer, the default value is returned.

package/src/json-rpc/server/safe_json_rpc_server.ts

@@ -60,6 +60,12 @@ export class SafeJsonRpcServer {
           const message = err.issues.map(e => `${e.message} (${e.path.join('.')})`).join('. ') || 'Validation error';
           ctx.status = 400;
           ctx.body = { jsonrpc: '2.0', id: null, error: { code: -32701, message } };
+        } else if (this.http200OnError) {
+          ctx.body = {
+            jsonrpc: '2.0',
+            id: null,
+            error: { code: err.code || -32600, data: err.data, message: err.message },
+          };
         } else {
           ctx.status = 500;
           ctx.body = { jsonrpc: '2.0', id: null, error: { code: -32600, message: err.message ?? 'Internal error' } };
@@ -111,16 +117,8 @@ export class SafeJsonRpcServer {
         ctx.body = { jsonrpc, id, error: { code: -32601, message: `Method not found: ${method}` } };
       } else {
         ctx.status = 200;
-        try {
-          const result = await this.proxy.call(method, params);
-          ctx.body = { jsonrpc, id, result };
-        } catch (err: any) {
-          if (this.http200OnError) {
-            ctx.body = { jsonrpc, id, error: { code: err.code || -32600, data: err.data, message: err.message } };
-          } else {
-            throw err;
-          }
-        }
+        const result = await this.proxy.call(method, params);
+        ctx.body = { jsonrpc, id, result };
       }
     });
 
@@ -265,6 +263,12 @@ function makeAggregateHealthcheck(namedHandlers: NamespacedApiHandlers, log?: Lo
   };
 }
 
+type SafeJsonRpcServerOptions = {
+  http200OnError: boolean;
+  healthCheck?: StatusCheckFn;
+  log?: Logger;
+};
+
 /**
  * Creates a single SafeJsonRpcServer from multiple handlers.
  * @param servers - List of handlers to be combined.
@@ -272,9 +276,12 @@ function makeAggregateHealthcheck(namedHandlers: NamespacedApiHandlers, log?: Lo
  */
 export function createNamespacedSafeJsonRpcServer(
   handlers: NamespacedApiHandlers,
-  http200OnError = false,
-  log = createLogger('json-rpc:server'),
+  options: Omit<SafeJsonRpcServerOptions, 'healthcheck'> = {
+    http200OnError: false,
+    log: createLogger('json-rpc:server'),
+  },
 ): SafeJsonRpcServer {
+  const { http200OnError, log } = options;
   const proxy = new NamespacedSafeJsonProxy(handlers);
   const healthCheck = makeAggregateHealthcheck(handlers, log);
   return new SafeJsonRpcServer(proxy, http200OnError, healthCheck, log);
@@ -283,11 +290,11 @@ export function createNamespacedSafeJsonRpcServer(
 export function createSafeJsonRpcServer<T extends object = any>(
   handler: T,
   schema: ApiSchemaFor<T>,
-  http200OnError = false,
-  healthCheck?: StatusCheckFn,
+  options: SafeJsonRpcServerOptions = { http200OnError: false },
 ) {
+  const { http200OnError, log, healthCheck } = options;
   const proxy = new SafeJsonProxy(handler, schema);
-  return new SafeJsonRpcServer(proxy, http200OnError, healthCheck);
+  return new SafeJsonRpcServer(proxy, http200OnError, healthCheck, log);
 }
 
 /**

package/src/log/gcloud-logger-config.ts

@@ -0,0 +1,71 @@
+import { type pino } from 'pino';
+
+/* eslint-disable camelcase */
+
+const GOOGLE_CLOUD_TRACE_ID = 'logging.googleapis.com/trace';
+const GOOGLE_CLOUD_SPAN_ID = 'logging.googleapis.com/spanId';
+const GOOGLE_CLOUD_TRACE_SAMPLED = 'logging.googleapis.com/trace_sampled';
+
+/**
+ * Pino configuration for google cloud observability. Tweaks message and timestamp,
+ * adds trace context attributes, and injects severity level.
+ * Adapted from https://cloud.google.com/trace/docs/setup/nodejs-ot#config-structured-logging.
+ */
+export const GoogleCloudLoggerConfig = {
+  messageKey: 'message',
+  // Same as pino.stdTimeFunctions.isoTime but uses "timestamp" key instead of "time"
+  timestamp(): string {
+    return `,"timestamp":"${new Date(Date.now()).toISOString()}"`;
+  },
+  formatters: {
+    log(object: Record<string, unknown>): Record<string, unknown> {
+      // Add trace context attributes following Cloud Logging structured log format described
+      // in https://cloud.google.com/logging/docs/structured-logging#special-payload-fields
+      const { trace_id, span_id, trace_flags, ...rest } = object;
+
+      if (trace_id && span_id) {
+        return {
+          [GOOGLE_CLOUD_TRACE_ID]: trace_id,
+          [GOOGLE_CLOUD_SPAN_ID]: span_id,
+          [GOOGLE_CLOUD_TRACE_SAMPLED]: trace_flags ? trace_flags === '01' : undefined,
+          trace_flags, // Keep the original trace_flags for otel-pino-stream
+          ...rest,
+        };
+      }
+      return object;
+    },
+    level(label: string, level: number): object {
+      // Inspired by https://github.com/pinojs/pino/issues/726#issuecomment-605814879
+      // Severity labels https://cloud.google.com/logging/docs/reference/v2/rest/v2/LogEntry#LogSeverity
+      let severity: string;
+
+      switch (label as pino.Level | keyof typeof customLevels) {
+        case 'trace':
+        case 'debug':
+          severity = 'DEBUG';
+          break;
+        case 'verbose':
+        case 'info':
+          severity = 'INFO';
+          break;
+        case 'warn':
+          severity = 'WARNING';
+          break;
+        case 'error':
+          severity = 'ERROR';
+          break;
+        case 'fatal':
+          severity = 'CRITICAL';
+          break;
+        default:
+          severity = 'DEFAULT';
+          break;
+      }
+
+      return { severity, level };
+    },
+  },
+} satisfies pino.LoggerOptions;
+
+// Define custom logging levels for pino. Duplicate from pino-logger.ts.
+const customLevels = { verbose: 25 };

package/src/log/pino-logger.ts

@@ -5,12 +5,14 @@ import { type Writable } from 'stream';
 import { inspect } from 'util';
 
 import { compactArray } from '../collection/array.js';
+import { type EnvVar, parseBooleanEnv } from '../config/index.js';
+import { GoogleCloudLoggerConfig } from './gcloud-logger-config.js';
 import { getLogLevelFromFilters, parseEnv } from './log-filters.js';
 import { type LogLevel } from './log-levels.js';
 import { type LogData, type LogFn } from './log_fn.js';
 
 export function createLogger(module: string): Logger {
-  module = module.replace(/^aztec:/, '');
+  module = logNameHandlers.reduce((moduleName, handler) => handler(moduleName), module.replace(/^aztec:/, ''));
   const pinoLogger = logger.child({ module }, { level: getLogLevelFromFilters(logFilters, module) });
 
   // We check manually for isLevelEnabled to avoid calling processLogData unnecessarily.
@@ -56,6 +58,22 @@ function processLogData(data: LogData): LogData {
   return logDataHandlers.reduce((accum, handler) => handler(accum), data);
 }
 
+// Allow global hooks for tweaking module names.
+// Used in tests to add a uid to modules, so we can differentiate multiple nodes in the same process.
+type LogNameHandler = (module: string) => string;
+const logNameHandlers: LogNameHandler[] = [];
+
+export function addLogNameHandler(handler: LogNameHandler): void {
+  logNameHandlers.push(handler);
+}
+
+export function removeLogNameHandler(handler: LogNameHandler) {
+  const index = logNameHandlers.indexOf(handler);
+  if (index !== -1) {
+    logNameHandlers.splice(index, 1);
+  }
+}
+
 // Patch isLevelEnabled missing from pino/browser.
 function isLevelEnabled(logger: pino.Logger<'verbose', boolean>, level: LogLevel): boolean {
   return typeof logger.isLevelEnabled === 'function'
@@ -70,44 +88,14 @@ const [logLevel, logFilters] = parseEnv(process.env.LOG_LEVEL, defaultLogLevel);
 // Define custom logging levels for pino.
 const customLevels = { verbose: 25 };
 
-// inspired by https://github.com/pinojs/pino/issues/726#issuecomment-605814879
-const levelToSeverityFormatter = (label: string, level: number): object => {
-  // Severity labels https://cloud.google.com/logging/docs/reference/v2/rest/v2/LogEntry#LogSeverity
-  let severity: string;
-
-  switch (label as pino.Level | keyof typeof customLevels) {
-    case 'trace':
-    case 'debug':
-      severity = 'DEBUG';
-      break;
-    case 'verbose':
-    case 'info':
-      severity = 'INFO';
-      break;
-    case 'warn':
-      severity = 'WARNING';
-      break;
-    case 'error':
-      severity = 'ERROR';
-      break;
-    case 'fatal':
-      severity = 'CRITICAL';
-      break;
-    default:
-      severity = 'DEFAULT';
-      break;
-  }
-
-  return { severity, level };
-};
-
+// Global pino options, tweaked for google cloud if running there.
+const useGcloudObservability = parseBooleanEnv(process.env['USE_GCLOUD_OBSERVABILITY' satisfies EnvVar]);
 const pinoOpts: pino.LoggerOptions<keyof typeof customLevels> = {
   customLevels,
+  messageKey: 'msg',
   useOnlyCustomLevels: false,
   level: logLevel,
-  formatters: {
-    level: levelToSeverityFormatter,
-  },
+  ...(useGcloudObservability ? GoogleCloudLoggerConfig : {}),
 };
 
 export const levels = {
@@ -127,7 +115,7 @@ export const pinoPrettyOpts = {
   customLevels: 'fatal:60,error:50,warn:40,info:30,verbose:25,debug:20,trace:10',
   customColors: 'fatal:bgRed,error:red,warn:yellow,info:green,verbose:magenta,debug:blue,trace:gray',
   minimumLevel: 'trace' as const,
-  singleLine: !['1', 'true'].includes(process.env.LOG_MULTILINE ?? ''),
+  singleLine: !parseBooleanEnv(process.env['LOG_MULTILINE' satisfies EnvVar]),
 };
 
 const prettyTransport: pino.TransportTargetOptions = {
@@ -149,14 +137,14 @@ const stdioTransport: pino.TransportTargetOptions = {
 // would mean that all child loggers created before the telemetry-client is initialized would not have
 // this transport configured. Note that the target is defined as the export in the telemetry-client,
 // since pino will load this transport separately on a worker thread, to minimize disruption to the main loop.
-const otlpEndpoint = process.env.OTEL_EXPORTER_OTLP_LOGS_ENDPOINT;
+const otlpEndpoint = process.env['OTEL_EXPORTER_OTLP_LOGS_ENDPOINT' satisfies EnvVar];
+const otlpEnabled = !!otlpEndpoint && !useGcloudObservability;
 const otelOpts = { levels };
 const otelTransport: pino.TransportTargetOptions = {
   target: '@aztec/telemetry-client/otel-pino-stream',
   options: otelOpts,
   level: 'trace',
 };
-
 function makeLogger() {
   if (!isNode) {
     // We are on the browser.
@@ -168,10 +156,10 @@ function makeLogger() {
     return pino(pinoOpts, pino.destination(2));
   } else {
     // Regular nodejs with transports on worker thread, using pino-pretty for console logging if LOG_JSON
-    // is not set, and an optional OTLP transport if the OTLP endpoint is provided.
+    // is not set, and an optional OTLP transport if the OTLP endpoint is set.
     const targets: pino.TransportSingleOptions[] = compactArray([
-      ['1', 'true', 'TRUE'].includes(process.env.LOG_JSON ?? '') ? stdioTransport : prettyTransport,
-      otlpEndpoint ? otelTransport : undefined,
+      parseBooleanEnv(process.env.LOG_JSON) ? stdioTransport : prettyTransport,
+      otlpEnabled ? otelTransport : undefined,
     ]);
     return pino(pinoOpts, pino.transport({ targets, levels: levels.values }));
   }
@@ -186,7 +174,7 @@ logger.verbose(
     ...logFilters.reduce((accum, [module, level]) => ({ ...accum, [`log.${module}`]: level }), {}),
   },
   isNode
-    ? `Logger initialized with level ${logLevel}` + (otlpEndpoint ? ` with OTLP exporter to ${otlpEndpoint}` : '')
+    ? `Logger initialized with level ${logLevel}` + (otlpEnabled ? ` with OTLP exporter to ${otlpEndpoint}` : '')
     : `Browser console logger initialized with level ${logLevel}`,
 );
 

package/src/promise/running-promise.ts

@@ -23,6 +23,10 @@ export class RunningPromise {
    * Starts the running promise.
    */
   public start() {
+    if (this.running) {
+      this.logger.warn(`Attempted to start running promise that was already started`);
+      return;
+    }
     this.running = true;
 
     const poll = async () => {
@@ -54,6 +58,10 @@ export class RunningPromise {
    * and waits for the currently executing function to complete.
    */
   async stop(): Promise<void> {
+    if (!this.running) {
+      this.logger.warn(`Running promise was not started`);
+      return;
+    }
     this.running = false;
     this.interruptibleSleep.interrupt();
     await this.runningPromise;

package/src/queue/serial_queue.ts

@@ -6,6 +6,7 @@ import { FifoMemoryQueue } from './fifo_memory_queue.js';
 export class SerialQueue {
   private readonly queue = new FifoMemoryQueue<() => Promise<void>>();
   private runningPromise!: Promise<void>;
+  private started = false;
 
   /**
    * Initializes the execution of enqueued functions in the serial queue.
@@ -14,7 +15,11 @@ export class SerialQueue {
    * This method should be called once to start processing the queue.
    */
   public start() {
+    if (this.started) {
+      return;
+    }
     this.runningPromise = this.queue.process(fn => fn());
+    this.started = true;
   }
 
   /**

package/src/testing/files/index.ts

@@ -17,6 +17,11 @@ export function writeTestData(targetFileFromRepoRoot: string, contents: string |
   logger(`Wrote test data to ${targetFile}`);
 }
 
+export function readTestData(repoPath: string): Buffer {
+  const targetFile = getPathToFile(repoPath);
+  return readFileSync(targetFile);
+}
+
 /**
  * Looks for a variable assignment in the target file and updates the value, only if test data generation is enabled.
  * Note that a magic inline comment would be a cleaner approach, like `/* TEST-DATA-START *\/` and `/* TEST-DATA-END *\/`,
@@ -43,7 +48,7 @@ export function updateInlineTestData(targetFileFromRepoRoot: string, itemName: s
 /**
  * Updates the sample Prover.toml files in noir-projects/noir-protocol-circuits/crates/.
  * @remarks Requires AZTEC_GENERATE_TEST_DATA=1 to be set
- * To re-gen, run 'AZTEC_GENERATE_TEST_DATA=1 FAKE_PROOFS=1 yarn test:e2e-no-docker full.test '
+ * To re-gen, run 'AZTEC_GENERATE_TEST_DATA=1 FAKE_PROOFS=1 yarn test:e2e full.test '
  */
 export function updateProtocolCircuitSampleInputs(circuitName: string, value: string) {
   const logger = createConsoleLogger('aztec:testing:test_data');

package/src/trees/index.ts

@@ -1,4 +1,5 @@
-export * from './unbalanced_merkle_root.js';
+export * from './unbalanced_merkle_tree.js';
+export * from './unbalanced_tree_store.js';
 
 /**
  * A leaf of an indexed merkle tree.

package/src/trees/unbalanced_merkle_tree.ts

@@ -0,0 +1,103 @@
+import { padArrayEnd } from '@aztec/foundation/collection';
+import { sha256Trunc } from '@aztec/foundation/crypto';
+
+/**
+ * Computes the merkle root for an unbalanced tree.
+ *
+ * @dev Adapted from unbalanced_tree.ts.
+ * Calculates the tree upwards layer by layer until we reach the root.
+ * The L1 calculation instead computes the tree from right to left (slightly cheaper gas).
+ * TODO: A more thorough investigation of which method is cheaper, then use that method everywhere.
+ */
+export function computeUnbalancedMerkleRoot(leaves: Buffer[], emptyLeaf?: Buffer, hasher = sha256Trunc): Buffer {
+  // Pad leaves to 2
+  if (leaves.length < 2) {
+    if (emptyLeaf === undefined) {
+      throw new Error('Cannot compute a Merkle root with less than 2 leaves');
+    } else {
+      leaves = padArrayEnd(leaves, emptyLeaf, 2);
+    }
+  }
+
+  const depth = Math.ceil(Math.log2(leaves.length));
+  let [layerWidth, nodeToShift] =
+    leaves.length & 1 ? [leaves.length - 1, leaves[leaves.length - 1]] : [leaves.length, Buffer.alloc(0)];
+  // Allocate this layer's leaves and init the next layer up
+  let thisLayer = leaves.slice(0, layerWidth);
+  let nextLayer = [];
+  for (let i = 0; i < depth; i++) {
+    for (let j = 0; j < layerWidth; j += 2) {
+      // Store the hash of each pair one layer up
+      nextLayer[j / 2] = hasher(Buffer.concat([thisLayer[j], thisLayer[j + 1]]));
+    }
+    layerWidth /= 2;
+    if (layerWidth & 1) {
+      if (nodeToShift.length) {
+        // If the next layer has odd length, and we have a node that needs to be shifted up, add it here
+        nextLayer.push(nodeToShift);
+        layerWidth += 1;
+        nodeToShift = Buffer.alloc(0);
+      } else {
+        // If we don't have a node waiting to be shifted, store the next layer's final node to be shifted
+        layerWidth -= 1;
+        nodeToShift = nextLayer[layerWidth];
+      }
+    }
+    // reset the layers
+    thisLayer = nextLayer;
+    nextLayer = [];
+  }
+  // return the root
+  return thisLayer[0];
+}
+
+function getMaxBalancedTreeDepth(numLeaves: number) {
+  return Math.floor(Math.log2(numLeaves));
+}
+
+function getMaxUnbalancedTreeDepth(numLeaves: number) {
+  return Math.ceil(Math.log2(numLeaves));
+}
+
+function findPosition(
+  rootLevel: number,
+  leafLevel: number,
+  numLeaves: number,
+  indexOffset: number,
+  targetIndex: number,
+): { level: number; indexAtLevel: number } {
+  if (numLeaves <= 1) {
+    // Single leaf.
+    return { level: rootLevel, indexAtLevel: indexOffset };
+  }
+
+  // The largest balanced tree that can be created with the given number of leaves.
+  const maxBalancedTreeDepth = getMaxBalancedTreeDepth(numLeaves);
+  const numBalancedLeaves = 2 ** maxBalancedTreeDepth;
+  const numRemainingLeaves = numLeaves - numBalancedLeaves;
+
+  if (targetIndex < numBalancedLeaves) {
+    // Target is in the balanced tree.
+
+    // - If numRemainingLeaves is 0: this balanced tree is grown from the current root.
+    // - If numRemainingLeaves is not 0: the remaining leaves will form another tree, which will become the right child of the root.
+    //   And the balanced tree will be the left child of the root.
+    //   There will be an extra level between the root of the balanced tree and the current root.
+    const extraLevel = numRemainingLeaves ? 1 : 0;
+
+    return { level: rootLevel + maxBalancedTreeDepth + extraLevel, indexAtLevel: indexOffset + targetIndex };
+  } else {
+    // Target is in the right branch.
+    const rightBranchMaxLevel = getMaxUnbalancedTreeDepth(numRemainingLeaves);
+    const shiftedUp = leafLevel - rootLevel - rightBranchMaxLevel - 1;
+    const nextLeafLevel = leafLevel - shiftedUp;
+    const newIndexOffset = (indexOffset + numBalancedLeaves) >> shiftedUp;
+    const shiftedTargetIndex = targetIndex - numBalancedLeaves;
+    return findPosition(rootLevel + 1, nextLeafLevel, numRemainingLeaves, newIndexOffset, shiftedTargetIndex);
+  }
+}
+
+export function findLeafLevelAndIndex(numLeaves: number, leafIndex: number) {
+  const maxLevel = getMaxUnbalancedTreeDepth(numLeaves);
+  return findPosition(0, maxLevel, numLeaves, 0, leafIndex);
+}

package/src/trees/unbalanced_tree_store.ts

@@ -0,0 +1,102 @@
+import { findLeafLevelAndIndex } from './unbalanced_merkle_tree.js';
+
+export interface TreeNodeLocation {
+  level: number;
+  index: number;
+}
+
+interface TreeNode<T> {
+  value: T;
+  location: TreeNodeLocation;
+}
+
+export class UnbalancedTreeStore<T> {
+  #nodeMapping: Map<string, TreeNode<T>> = new Map();
+  readonly #numLeaves: number;
+
+  constructor(numLeaves: number) {
+    this.#numLeaves = numLeaves;
+  }
+
+  setLeaf(leafIndex: number, value: T): TreeNodeLocation {
+    if (leafIndex >= this.#numLeaves) {
+      throw new Error(`Expected at most ${this.#numLeaves} leaves. Received a leaf at index ${leafIndex}.`);
+    }
+
+    const { level, indexAtLevel } = findLeafLevelAndIndex(this.#numLeaves, leafIndex);
+    const location = {
+      level,
+      index: indexAtLevel,
+    };
+    this.#nodeMapping.set(this.#getKey(location), {
+      location,
+      value,
+    });
+    return location;
+  }
+
+  setNode({ level, index }: TreeNodeLocation, value: T) {
+    const location = {
+      level,
+      index,
+    };
+    this.#nodeMapping.set(this.#getKey(location), {
+      location,
+      value,
+    });
+  }
+
+  getParentLocation({ level, index }: TreeNodeLocation): TreeNodeLocation {
+    if (level === 0) {
+      throw new Error('Tree root does not have a parent.');
+    }
+
+    return { level: level - 1, index: Math.floor(index / 2) };
+  }
+
+  getSiblingLocation({ level, index }: TreeNodeLocation): TreeNodeLocation {
+    if (level === 0) {
+      throw new Error('Tree root does not have a sibling.');
+    }
+
+    return { level, index: index % 2 ? index - 1 : index + 1 };
+  }
+
+  getChildLocations({ level, index }: TreeNodeLocation): [TreeNodeLocation, TreeNodeLocation] {
+    const left = { level: level + 1, index: index * 2 };
+    const right = { level: level + 1, index: index * 2 + 1 };
+    return [left, right];
+  }
+
+  getLeaf(leafIndex: number) {
+    const { level, indexAtLevel } = findLeafLevelAndIndex(this.#numLeaves, leafIndex);
+    const location = {
+      level,
+      index: indexAtLevel,
+    };
+    return this.getNode(location);
+  }
+
+  getNode(location: TreeNodeLocation): T | undefined {
+    return this.#nodeMapping.get(this.#getKey(location))?.value;
+  }
+
+  getParent(location: TreeNodeLocation): T | undefined {
+    const parentLocation = this.getParentLocation(location);
+    return this.getNode(parentLocation);
+  }
+
+  getSibling(location: TreeNodeLocation): T | undefined {
+    const siblingLocation = this.getSiblingLocation(location);
+    return this.getNode(siblingLocation);
+  }
+
+  getChildren(location: TreeNodeLocation): [T | undefined, T | undefined] {
+    const [left, right] = this.getChildLocations(location);
+    return [this.getNode(left), this.getNode(right)];
+  }
+
+  #getKey(location: TreeNodeLocation) {
+    return `${location.level}-${location.index}`;
+  }
+}