MerkleTree.ts 5.29 KB
Newer Older
John Doe's avatar
John Doe committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182
// @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);
    }
  }
}