MerkleTree.ts

// @ts-ignore
import { Fr } from '@aztec/bb.js';
// @ts-ignore -- no types

export async function pedersenLeftRight(
  barretenberg: any, 
  left: string, 
  right: string) {

    let leftBuffer = Fr.fromBufferReduce(Buffer.from(left.slice(2), 'hex'));
    let rightBuffer = Fr.fromBufferReduce(Buffer.from(right.slice(2), 'hex'));
    let hashRes = await barretenberg.pedersenPlookupCompressFields(leftBuffer, rightBuffer);
    return hashRes.toString()
}

export class MerkleTree  {
  zeroValue = "0xf35fcb490b7ea67c3ac26ed530fa5d8dfe8be344e7177ebb63fe02723fb6f725"; // sha256("Momiji")
  levels: number;
  defaultLeaves: string[]
  hashLeftRight: any;
  storage: Map<string, string>;
  zeros: string[];
  totalLeaves: number;
  barretenberg: any;

  constructor(
    levels, 
    barretenberg,
    defaultLeaves = [], 
    hashLeftRight = pedersenLeftRight) {
    this.levels = levels;
    this.hashLeftRight = hashLeftRight;
    this.storage = new Map();
    this.zeros = [];
    this.totalLeaves = 0;
    this.barretenberg = barretenberg;
    this.defaultLeaves = defaultLeaves
  }

  async init() {
    // build zeros depends on tree levels
    let currentZero = this.zeroValue;
    this.zeros.push(currentZero);
    for (let i = 0; i < this.levels; i++) {
      currentZero = await this.hashLeftRight(this.barretenberg, currentZero, currentZero);
      this.zeros.push(currentZero);
    }

    if (this.defaultLeaves.length > 0) {
      this.totalLeaves = this.defaultLeaves.length;

      // store leaves with key value pair
      let level = 0;
      this.defaultLeaves.forEach((leaf, index) => {
        this.storage.set(MerkleTree.indexToKey(level, index), leaf);
      });

      // build tree with initial leaves
      level++;
      let numberOfNodesInLevel = Math.ceil(this.totalLeaves / 2);
      for (level; level <= this.levels; level++) {
        for (let i = 0; i < numberOfNodesInLevel; i++) {
          const leftKey = MerkleTree.indexToKey(level - 1, 2 * i);
          const rightKey = MerkleTree.indexToKey(level - 1, 2 * i + 1);

          const left = this.storage.get(leftKey);
          const right = this.storage.get(rightKey) || this.zeros[level - 1];
          if (!left) throw new Error("leftKey not found");

          const node = await this.hashLeftRight(this.barretenberg, left, right);
          this.storage.set(MerkleTree.indexToKey(level, i), node);
        }
        numberOfNodesInLevel = Math.ceil(numberOfNodesInLevel / 2);
      }
    }
  }

  static indexToKey(level, index) {
    return `${level}-${index}`;
  }

  getIndex(leaf) {
    for (const [key, value] of this.storage) {
      if (value === leaf) {
        return Number(key.split("-")[1]);
      }
    }
    return -1;
  }

  root() {
    return this.storage.get(MerkleTree.indexToKey(this.levels, 0)) || this.zeros[this.levels];
  }

  proof(indexOfLeaf) {
    let pathElements: string[] = [];
    let pathIndices: number[] = [];

    const leaf = this.storage.get(MerkleTree.indexToKey(0, indexOfLeaf));
    if (!leaf) throw new Error("leaf not found");

    // store sibling into pathElements and target's indices into pathIndices
    const handleIndex = (level, currentIndex, siblingIndex) => {
      const siblingValue = this.storage.get(MerkleTree.indexToKey(level, siblingIndex)) || this.zeros[level];
      pathElements.push(siblingValue);
      pathIndices.push(currentIndex % 2);
    };

    this.traverse(indexOfLeaf, handleIndex);

    return {
      root: this.root(),
      pathElements,
      pathIndices,
      leaf: leaf,
    };
  }

  async insert(leaf) {
    const index = this.totalLeaves;
    await this.update(index, leaf, true);
    this.totalLeaves++;
  }

  async update(index, newLeaf, isInsert = false) {
    if (!isInsert && index >= this.totalLeaves) {
      throw Error("Use insert method for new elements.");
    } else if (isInsert && index < this.totalLeaves) {
      throw Error("Use update method for existing elements.");
    }

    let keyValueToStore: any[] = [];
    let currentElement = newLeaf;

    const handleIndex = async (level, currentIndex, siblingIndex) => {
      const siblingElement = this.storage.get(MerkleTree.indexToKey(level, siblingIndex)) || this.zeros[level];
  
      let left;
      let right;
      if (currentIndex % 2 === 0) {
        left = currentElement;
        right = siblingElement;
      } else {
        left = siblingElement;
        right = currentElement;
      }
  
      keyValueToStore.push({
        key: MerkleTree.indexToKey(level, currentIndex),
        value: currentElement,
      });
      currentElement = await this.hashLeftRight(this.barretenberg, left, right);
    };
    await this.traverse(index, handleIndex);

    // push root to the end
    keyValueToStore.push({
      key: MerkleTree.indexToKey(this.levels, 0),
      value: currentElement,
    });

    keyValueToStore.forEach(o => {
      this.storage.set(o.key, o.value);
    });
  }

  // traverse from leaf to root with handler for target node and sibling node
  async traverse(indexOfLeaf, handler) {
    let currentIndex = indexOfLeaf;
    for (let i = 0; i < this.levels; i++) {
      let siblingIndex;
      if (currentIndex % 2 === 0) {
        siblingIndex = currentIndex + 1;
      } else {
        siblingIndex = currentIndex - 1;
      }

      await handler(i, currentIndex, siblingIndex);
      currentIndex = Math.floor(currentIndex / 2);
    }
  }
}