feat: merkle root tracker

src/index.ts

@@ -7,6 +7,7 @@ import {
   RLNInstance,
 } from "./rln.js";
 import { RLNContract } from "./rln_contract.js";
+import { MerkleRootTracker } from "./root_tracker.js";
 
 // reexport the create function, dynamically imported from rln.ts
 export async function create(): Promise<RLNInstance> {
@@ -24,6 +25,7 @@ export {
   ProofMetadata,
   RLNEncoder,
   RLNDecoder,
+  MerkleRootTracker,
   RLNContract,
   RLN_ABI,
   GOERLI_CONTRACT,

src/root_tracker.spec.ts

@@ -0,0 +1,56 @@
+import { assert, expect } from "chai";
+
+import { MerkleRootTracker } from "./root_tracker";
+
+describe("js-rln", () => {
+  it("should track merkle roots and backfill from block number", async function () {
+    const acceptableRootWindow = 3;
+
+    const tracker = new MerkleRootTracker(
+      acceptableRootWindow,
+      new Uint8Array([0, 0, 0, 0])
+    );
+    expect(tracker.roots()).to.have.length(1);
+    expect(tracker.buffer()).to.have.length(0);
+    expect(tracker.roots()[0]).to.deep.equal(new Uint8Array([0, 0, 0, 0]));
+
+    for (let i = 1; i <= 30; i++) {
+      tracker.pushRoot(i, new Uint8Array([0, 0, 0, i]));
+    }
+
+    expect(tracker.roots()).to.have.length(acceptableRootWindow);
+    expect(tracker.buffer()).to.have.length(20);
+    assert.sameDeepMembers(tracker.roots(), [
+      new Uint8Array([0, 0, 0, 30]),
+      new Uint8Array([0, 0, 0, 29]),
+      new Uint8Array([0, 0, 0, 28]),
+    ]);
+
+    // Buffer should keep track of 20 blocks previous to the current valid merkle root window
+    expect(tracker.buffer()[0]).to.be.eql(new Uint8Array([0, 0, 0, 8]));
+    expect(tracker.buffer()[19]).to.be.eql(new Uint8Array([0, 0, 0, 27]));
+
+    // Remove roots 29 and 30
+    tracker.backFill(29);
+    assert.sameDeepMembers(tracker.roots(), [
+      new Uint8Array([0, 0, 0, 28]),
+      new Uint8Array([0, 0, 0, 27]),
+      new Uint8Array([0, 0, 0, 26]),
+    ]);
+
+    expect(tracker.buffer()).to.have.length(18);
+    expect(tracker.buffer()[0]).to.be.eql(new Uint8Array([0, 0, 0, 8]));
+    expect(tracker.buffer()[17]).to.be.eql(new Uint8Array([0, 0, 0, 25]));
+
+    // Remove roots from block 15 onwards. These blocks exists within the buffer
+    tracker.backFill(15);
+    assert.sameDeepMembers(tracker.roots(), [
+      new Uint8Array([0, 0, 0, 14]),
+      new Uint8Array([0, 0, 0, 13]),
+      new Uint8Array([0, 0, 0, 12]),
+    ]);
+    expect(tracker.buffer()).to.have.length(4);
+    expect(tracker.buffer()[0]).to.be.eql(new Uint8Array([0, 0, 0, 8]));
+    expect(tracker.buffer()[3]).to.be.eql(new Uint8Array([0, 0, 0, 11]));
+  });
+});

src/root_tracker.ts

@@ -0,0 +1,88 @@
+class RootPerBlock {
+  constructor(public root: Uint8Array, public blockNumber: number) {}
+}
+
+const maxBufferSize = 20;
+
+export class MerkleRootTracker {
+  private validMerkleRoots: Array<RootPerBlock> = new Array<RootPerBlock>();
+  private merkleRootBuffer: Array<RootPerBlock> = new Array<RootPerBlock>();
+  constructor(
+    private acceptableRootWindowSize: number,
+    initialRoot: Uint8Array
+  ) {
+    this.pushRoot(0, initialRoot);
+  }
+
+  backFill(fromBlockNumber: number): void {
+    if (this.validMerkleRoots.length == 0) return;
+
+    let numBlocks = 0;
+    for (let i = this.validMerkleRoots.length - 1; i >= 0; i--) {
+      if (this.validMerkleRoots[i].blockNumber >= fromBlockNumber) {
+        numBlocks++;
+      }
+    }
+
+    if (numBlocks == 0) return;
+
+    const olderBlock = fromBlockNumber < this.validMerkleRoots[0].blockNumber;
+
+    // Remove last roots
+    let rootsToPop = numBlocks;
+    if (this.validMerkleRoots.length < rootsToPop) {
+      rootsToPop = this.validMerkleRoots.length;
+    }
+
+    this.validMerkleRoots = this.validMerkleRoots.slice(
+      0,
+      this.validMerkleRoots.length - rootsToPop
+    );
+
+    if (this.merkleRootBuffer.length == 0) return;
+
+    if (olderBlock) {
+      const idx = this.merkleRootBuffer.findIndex(
+        (x) => x.blockNumber == fromBlockNumber
+      );
+      if (idx > -1) {
+        this.merkleRootBuffer = this.merkleRootBuffer.slice(0, idx);
+      }
+    }
+
+    // Backfill the tree's acceptable roots
+    let rootsToRestore =
+      this.acceptableRootWindowSize - this.validMerkleRoots.length;
+    if (this.merkleRootBuffer.length < rootsToRestore) {
+      rootsToRestore = this.merkleRootBuffer.length;
+    }
+
+    for (let i = 0; i < rootsToRestore; i++) {
+      const x = this.merkleRootBuffer.pop();
+      if (x) this.validMerkleRoots.unshift(x);
+    }
+  }
+
+  pushRoot(blockNumber: number, root: Uint8Array): void {
+    this.validMerkleRoots.push(new RootPerBlock(root, blockNumber));
+
+    // Maintain valid merkle root window
+    if (this.validMerkleRoots.length > this.acceptableRootWindowSize) {
+      const x = this.validMerkleRoots.shift();
+      if (x) this.merkleRootBuffer.push(x);
+    }
+
+    // Maintain merkle root buffer
+    if (this.merkleRootBuffer.length > maxBufferSize) {
+      this.merkleRootBuffer.shift();
+    }
+  }
+
+  roots(): Array<Uint8Array> {
+    return this.validMerkleRoots.map((x) => x.root);
+  }
+
+  buffer(): Array<Uint8Array> {
+    return this.merkleRootBuffer.map((x) => x.root);
+  }
+}