Building On Mina: A guide to writing smart contract implementing Merkle Trees using 01js

npm install -g zkapp-cli

zk project votingNft

cd votingNft

cd src
rm Add.test.ts Add.ts interact.ts

touch Voting.ts Voting.test.ts

//exit the directory
cd ..

import { Field, SmartContract, state, State, method, Struct, MerkleWitness, Poseidon, Provable } from "o1js";

export const treeHeight = 4

export class MerkleWitness4 extends MerkleWitness(treeHeight){}

export class Voting extends SmartContract{
    // declare state variables
    @state(Field) voters = State<Field>()
    @state(Field) winningNft = State<Field>()
    @state(Field) winningVotes = State<Field>()

    // root of the merkle tree
    @state(Field) treeRoot = State<Field>()

    init(){
      super.init()
      this.voters.set(Field(1))
    }

   // initialize the state on the contract
       @method async initState(initialRoot: Field){
     this.treeRoot.set(initialRoot)
   }
}

export class NFT extends Struct({nftName: Field, nftCreator: Field, nftVotes: Field}){

//method for the vote
vote(){
  this.nftVotes = this.nftVotes.add(1)
 }
}

export class Voting extends SmartContract{
  // declare state variables
  @state(Field) voters = State<Field>()
  @state(Field) winningNft = State<Field>()
  @state(Field) winningVotes = State<Field>()

  // root of the merkle tree
  @state(Field) treeRoot = State<Field>()
     init(){
       super.init()
       this.voters.set(Field(1))
     }

  // initialize the state of the contract
  @method async initState(initialRoot: Field){
    this.treeRoot.set(initialRoot)
   }

 //method to vote for the nft
 @method async voteForNft(nft: NFT, witness: MerkleWitness4){

  //Get the tree root and see if the nft is in it
  const treeRoot = this.treeRoot.getAndRequireEquals()

  // check to see whether the nft is in the merkle tree, the way we 
  // check to see if the nft is within the tree is by the witness       

  const nftRoot = witness.calculateRoot(Poseidon.hash(NFT.toFields(nft)))

  // check if the nft is in the tree, if it is not in the tree, the transaction will fail and the state will not be updated

   nftRoot.assertEquals(treeRoot)

  //if the nft is in the tree then the nft can be voted on
  nft.vote()

  //update the tree with the new nft hash
  const newNftRoot = witness.calculateRoot(Poseidon.hash(NFT.toFields(nft)))
  this.treeRoot.set(newNftRoot)

//we want to increment the number of voters after a vote has been cast      this.voters.set(this.voters.getAndRequireEquals().add(1))

//we are going to check if the nft has more votes than the current winning nft then we will update the winning nft

  const winningNft = this.winningNft.getAndRequireEquals()
  const winningVotes = this.winningVotes.getAndRequireEquals()

  const newWinningNft = Provable.if(winningVotes.lessThan(nft.nftVotes), nft.nftName, winningNft)
  const newWinningVotes = Provable.if(winningVotes.lessThan(nft.nftVotes), nft.nftName, winningVotes)

this.winningNft.set(newWinningNft)
this.winningVotes.set(newWinningVotes)
  }
}

import { AccountUpdate, CircuitString, Field, MerkleTree, Mina, Poseidon, PrivateKey, PublicKey } from 'o1js';
import { treeHeight, NFT, MerkleWitness4, Voting } from './Voting';

let proofsEnabled = false;

function createNFT(nftName: string, nftCreator: string) {
    const nft = new NFT({
        nftName: Poseidon.hash(CircuitString.toFields(CircuitString.fromString(nftName))),
        nftCreator: Poseidon.hash(CircuitString.toFields(CircuitString.fromString(nftCreator))),
        nftVotes: Field(0)
    });
    return nft;
}

function createTree(): MerkleTree {
    const tree = new MerkleTree(treeHeight);

    const nft0 = createNFT("nft0", "creator0");
    const nft1 = createNFT("nft1", "creator1");
    const nft2 = createNFT("nft2", "creator2");
    const nft3 = createNFT("nft3", "creator3");

    tree.setLeaf(0n, Poseidon.hash(NFT.toFields(nft0)));
    tree.setLeaf(1n, Poseidon.hash(NFT.toFields(nft1)));
    tree.setLeaf(2n, Poseidon.hash(NFT.toFields(nft2)));
    tree.setLeaf(3n, Poseidon.hash(NFT.toFields(nft3)));

    return tree;
}


describe('Voting', () => {
    let deployerAccount: Mina.TestPublicKey,
    deployerKey: PrivateKey,
    senderAccount: Mina.TestPublicKey,
    senderKey: PrivateKey,
    zkAppAddress: PublicKey,
    zkAppPrivateKey: PrivateKey,
    zkApp: Voting,
    tree: MerkleTree;

    beforeAll(async () => {
        if (proofsEnabled) await Voting.compile();
    });

    beforeEach(async () => {
        const Local = await Mina.LocalBlockchain({ proofsEnabled });
        Mina.setActiveInstance(Local);
        [deployerAccount, senderAccount] = Local.testAccounts;
        deployerKey = deployerAccount.key;
        senderKey = senderAccount.key;

        zkAppPrivateKey = PrivateKey.random();
        zkAppAddress = zkAppPrivateKey.toPublicKey();
        zkApp = new Voting(zkAppAddress);
        tree = createTree();
    });

async function localDeploy() {
        const txn = await Mina.transaction(deployerAccount, async () => {
            AccountUpdate.fundNewAccount(deployerAccount);
            await zkApp.deploy();
            await zkApp.initState(tree.getRoot());
        });
        await txn.prove();
        await txn.sign([deployerKey, zkAppPrivateKey]).send();
    }
})