Add Experimental Merkle Tree Implementation

src/examples/zkapps/merkle_tree/merkle_zkapp.ts

@@ -0,0 +1,132 @@
+/*
+Description: 
+
+This example describes how developers can use Merkle Trees as a basic off-chain storage tool.
+
+zkApps on Mina can only store a small amount of data on-chain, but many use cases require your application to at least reference big amounts of data.
+Merkle Trees give developers the power of storing large amounts of data off-chain, but proving its integrity to the on-chain smart contract!
+
+
+! Unfamiliar with Merkle Trees? No problem! Check out https://blog.ethereum.org/2015/11/15/merkling-in-ethereum/
+*/
+
+import {
+  SmartContract,
+  isReady,
+  shutdown,
+  Poseidon,
+  Field,
+  Experimental,
+  Permissions,
+  DeployArgs,
+  State,
+  state,
+  Circuit,
+  CircuitValue,
+  PublicKey,
+  UInt64,
+  prop,
+  Mina,
+  method,
+  UInt32,
+} from 'snarkyjs';
+
+await isReady;
+
+class Account extends CircuitValue {
+  @prop publicKey: PublicKey;
+  @prop points: UInt32;
+
+  constructor(publicKey: PublicKey, points: UInt32) {
+    super(publicKey, points);
+    this.publicKey = publicKey;
+    this.points = points;
+  }
+
+  hash(): Field {
+    return Poseidon.hash(this.toFields());
+  }
+}
+// we need the initiate tree root in order to tell the contract about our off-chain storage
+let initialCommitment: Field = Field.zero;
+/*
+  We want to write a smart contract that serves as a leaderboard,
+  but only has the commitment of the off-chain storage stored in an on-chain variable.
+  The accounts of all participants will be stored off-chain!
+  If a participant can guess the preimage of a hash, they will be granted one point :)
+*/
+
+class Leaderboard extends SmartContract {
+  // a commitment is a cryptographic primitive that allows us to commit to data, with the ability to "reveal" it later
+  @state(Field) commitment = State<Field>();
+
+  deploy(args: DeployArgs) {
+    super.deploy(args);
+    this.setPermissions({
+      ...Permissions.default(),
+      editState: Permissions.proofOrSignature(),
+    });
+    this.commitment.set(initialCommitment);
+  }
+
+  @method
+  guessPreimage(
+    guess: Field,
+    account: Account,
+    path: Experimental.MerkleWitness
+  ) {
+    // this is our hash! its the hash of the preimage "22", but keep it a secret!
+    let target =
+      Field(
+        17057234437185175411792943285768571642343179330449434169483610110583519635705
+      );
+    // if our guess preimage hashes to our target, we won a point!
+    Poseidon.hash([guess]).assertEquals(target);
+
+    // we fetch the on-chain commitment
+    let commitment = this.commitment.get();
+    this.commitment.assertEquals(commitment);
+
+    // we check that the account is within the committed Merkle Tree
+    //path.calculateRoot(account.hash()).assertEquals(commitment);
+
+    // we update the account and grant one point!
+    account.points = account.points.add(1);
+
+    // we calculate the new Merkle Root, based on the account changes
+    //let newCommitment = path.calculateRoot(account.hash());
+
+    //this.commitment.set(newCommitment);
+  }
+}
+
+// we initialize a new Merkle Tree with height 8
+const Tree = new Experimental.MerkleTree(8);
+
+let Local = Mina.LocalBlockchain();
+Mina.setActiveInstance(Local);
+
+// this map serves as our off-chain in-memory storage
+let Accounts: Map<string, Account> = new Map<string, Account>();
+
+let bob = new Account(Local.testAccounts[0].publicKey, UInt32.from(0));
+let alice = new Account(Local.testAccounts[1].publicKey, UInt32.from(0));
+let charlie = new Account(Local.testAccounts[2].publicKey, UInt32.from(0));
+let olivia = new Account(Local.testAccounts[3].publicKey, UInt32.from(0));
+
+Accounts.set('Bob', bob);
+Accounts.set('Alice', alice);
+Accounts.set('Charlie', charlie);
+Accounts.set('Olivia', olivia);
+
+// we now need "wrap" the Merkle tree around our off-chain storage
+
+Tree.setLeaf(0n, bob.hash());
+Tree.setLeaf(1n, alice.hash());
+Tree.setLeaf(2n, charlie.hash());
+Tree.setLeaf(3n, olivia.hash());
+
+// now that we got our accounts set up, we need the commitment to deploy our contract!
+initialCommitment = Tree.getRoot();
+
+console.log(Poseidon.hash([Field(22)]).toString());

src/index.ts

@@ -62,21 +62,34 @@ import {
 } from './lib/circuit_value';
 import { jsLayout, asFieldsAndAux } from './snarky/types';
 import { packToFields } from './lib/hash';
+import { MerkleTree, MerkleWitness as MerkleWitness_ } from './lib/merkle_tree';
 export { Experimental };
 
-/**
- * This module exposes APIs that are unstable, in the sense that the API surface is expected to change.
- * (Not unstable in the sense that they are less functional or tested than other parts.)
- */
-const Experimental = {
+const Experimental_ = {
   Reducer,
   createChildParty,
   memoizeWitness,
   // TODO: for testing, maybe remove later
   jsLayout,
   asFieldsAndAux,
   packToFields,
+  MerkleTree,
 };
+
+/**
+ * This module exposes APIs that are unstable, in the sense that the API surface is expected to change.
+ * (Not unstable in the sense that they are less functional or tested than other parts.)
+ */
 namespace Experimental {
+  export let Reducer = Experimental_.Reducer;
+  export let createChildParty = Experimental_.createChildParty;
+  export let memoizeWitness = Experimental_.memoizeWitness;
+  export let jsLayout = Experimental_.jsLayout;
+  export let asFieldsAndAux = Experimental_.asFieldsAndAux;
+  export let packToFields = Experimental_.packToFields;
+  export let MerkleTree = Experimental_.MerkleTree;
+
   export type AsFieldsAndAux<T, TJson> = AsFieldsAndAux_<T, TJson>;
+  export let MerkleWitness = MerkleWitness_;
+  export type MerkleWitness = MerkleWitness_;
 }

src/lib/merkle_tree.test.ts

@@ -0,0 +1,61 @@
+import {
+  isReady,
+  shutdown,
+  Poseidon,
+  Field,
+  Experimental,
+} from '../../dist/server';
+
+describe('Merkle Tree', () => {
+  beforeAll(async () => {
+    await isReady;
+  });
+  afterAll(async () => {
+    setTimeout(shutdown, 0);
+  });
+
+  it('root of empty tree of size 1', () => {
+    const tree = new Experimental.MerkleTree(1);
+    expect(tree.getRoot().toString()).toEqual(Field(0).toString());
+  });
+
+  it('root is correct', () => {
+    const tree = new Experimental.MerkleTree(2);
+    tree.setLeaf(0n, Field(1));
+    tree.setLeaf(1n, Field(2));
+    expect(tree.getRoot().toString()).toEqual(
+      Poseidon.hash([Field(1), Field(2)]).toString()
+    );
+  });
+
+  it('builds correct tree', () => {
+    const tree = new Experimental.MerkleTree(4);
+    tree.setLeaf(0n, Field(1));
+    tree.setLeaf(1n, Field(2));
+    tree.setLeaf(2n, Field(3));
+    expect(tree.validate(0n)).toBe(true);
+    expect(tree.validate(1n)).toBe(true);
+    expect(tree.validate(2n)).toBe(true);
+    expect(tree.validate(3n)).toBe(true);
+  });
+
+  it('tree of height 128', () => {
+    const tree = new Experimental.MerkleTree(128);
+
+    const index = 2n ** 64n;
+    expect(tree.validate(index)).toBe(true);
+
+    tree.setLeaf(index, Field(1));
+    expect(tree.validate(index)).toBe(true);
+  });
+
+  it('tree of height 256', () => {
+    const tree = new Experimental.MerkleTree(256);
+
+    const index = 2n ** 128n;
+    expect(tree.validate(index)).toBe(true);
+
+    tree.setLeaf(index, Field(1));
+    expect(tree.validate(index)).toBe(true);
+  });
+});

src/lib/merkle_tree.ts

@@ -0,0 +1,115 @@
+import { Circuit, CircuitValue, arrayProp } from './circuit_value';
+import { Poseidon } from './hash';
+import { Bool, Field } from './core';
+
+// external API
+export { Witness, MerkleTree, MerkleWitness };
+
+type Witness = { isLeft: boolean; sibling: Field }[];
+
+/**
+ * Levels are indexed from leafs (level 0) to root (level N - 1).
+ */
+class MerkleTree {
+  private nodes: Record<number, Record<string, Field>> = {};
+  private zeroes: Field[];
+
+  constructor(public readonly height: number) {
+    this.zeroes = [Field(0)];
+    for (let i = 1; i < height; i++) {
+      this.zeroes.push(Poseidon.hash([this.zeroes[i - 1], this.zeroes[i - 1]]));
+    }
+  }
+
+  getNode(level: number, index: bigint): Field {
+    return this.nodes[level]?.[index.toString()] ?? this.zeroes[level];
+  }
+
+  getRoot(): Field {
+    return this.getNode(this.height - 1, 0n);
+  }
+
+  private setNode(level: number, index: bigint, value: Field) {
+    (this.nodes[level] ??= {})[index.toString()] = value;
+  }
+
+  setLeaf(index: bigint, leaf: Field) {
+    this.setNode(0, index, leaf);
+    let currIndex = index;
+    for (let level = 1; level < this.height; level++) {
+      currIndex = currIndex / 2n;
+
+      const left = this.getNode(level - 1, currIndex * 2n);
+      const right = this.getNode(level - 1, currIndex * 2n + 1n);
+
+      this.setNode(level, currIndex, Poseidon.hash([left, right]));
+    }
+  }
+
+  getWitness(index: bigint): Witness {
+    const witness = [];
+    for (let level = 0; level < this.height - 1; level++) {
+      const isLeft = index % 2n === 0n;
+      witness.push({
+        isLeft,
+        sibling: this.getNode(level, isLeft ? index + 1n : index - 1n),
+      });
+      index = index / 2n;
+    }
+    return witness;
+  }
+
+  validate(index: bigint): boolean {
+    const path = this.getWitness(index);
+    let hash = this.getNode(0, index);
+    for (const node of path) {
+      hash = Poseidon.hash(
+        node.isLeft ? [hash, node.sibling] : [node.sibling, hash]
+      );
+    }
+
+    return hash.toString() === this.getRoot().toString();
+  }
+
+  fill(leaves: Field[]) {
+    leaves.forEach((value, index) => {
+      this.setLeaf(BigInt(index), value);
+    });
+  }
+}
+
+class MerkleWitness extends CircuitValue {
+  @arrayProp(Field, 255) path: Field[];
+  @arrayProp(Bool, 255) isLeft: Bool[];
+
+  constructor(witness: Witness) {
+    super(witness);
+
+    this.path = witness.map((item) => item.sibling);
+    this.isLeft = witness.map((item) => Bool(item.isLeft));
+  }
+
+  calculateRoot(leaf: Field): Field {
+    let hash = leaf;
+
+    for (let i = 1; i < 256; ++i) {
+      const left = Circuit.if(this.isLeft[i - 1], hash, this.path[i - 1]);
+      const right = Circuit.if(this.isLeft[i - 1], this.path[i - 1], hash);
+      hash = Poseidon.hash([left, right]);
+    }
+
+    return hash;
+  }
+
+  calculateIndex(): Field {
+    let powerOfTwo = Field(1);
+    let index = Field(0);
+
+    for (let i = 1; i < 256; ++i) {
+      index = Circuit.if(this.isLeft[i - 1], index, index.add(powerOfTwo));
+      powerOfTwo = powerOfTwo.mul(2);
+    }
+
+    return index;
+  }
+}