diff --git a/src/index.ts b/src/index.ts index d193cbd..d127c19 100644 --- a/src/index.ts +++ b/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 { @@ -24,6 +25,7 @@ export { ProofMetadata, RLNEncoder, RLNDecoder, + MerkleRootTracker, RLNContract, RLN_ABI, GOERLI_CONTRACT, diff --git a/src/root_tracker.spec.ts b/src/root_tracker.spec.ts new file mode 100644 index 0000000..d34e768 --- /dev/null +++ b/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])); + }); +}); diff --git a/src/root_tracker.ts b/src/root_tracker.ts new file mode 100644 index 0000000..d648d0c --- /dev/null +++ b/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 = new Array(); + private merkleRootBuffer: Array = new Array(); + 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 { + return this.validMerkleRoots.map((x) => x.root); + } + + buffer(): Array { + return this.merkleRootBuffer.map((x) => x.root); + } +}