feat: generate merkle tree

examples/CRISP/sdk/README.md

@@ -1,6 +1,6 @@
 # CRISP SDK
 
-This package is CRISP's TypeScript SDK for interacting with CRISP and the CRISP server. 
+This package is CRISP's TypeScript SDK for interacting with CRISP and the CRISP server.
 
 ## Installation
 
@@ -10,9 +10,9 @@ npm install @enclave/crisp-sdk
 
 ## Release
 
-The SDK is published on npmjs and does not follow the same versioning as the main Enclave packages. 
+The SDK is published on npmjs and does not follow the same versioning as the main Enclave packages.
 To release a new version, run the following command:
 
-```bash 
+```bash
 pnpm release
-``` 
+```

examples/CRISP/sdk/package.json

@@ -36,5 +36,10 @@
     "vitest": "^1.6.1",
     "typescript": "^5.0.0"
   },
-  "packageManager": "pnpm@10.7.1+sha512.2d92c86b7928dc8284f53494fb4201f983da65f0fb4f0d40baafa5cf628fa31dae3e5968f12466f17df7e97310e30f343a648baea1b9b350685dafafffdf5808"
+  "packageManager": "pnpm@10.7.1+sha512.2d92c86b7928dc8284f53494fb4201f983da65f0fb4f0d40baafa5cf628fa31dae3e5968f12466f17df7e97310e30f343a648baea1b9b350685dafafffdf5808",
+  "dependencies": {
+    "@zk-kit/lean-imt": "^2.2.4",
+    "poseidon-lite": "^0.3.0",
+    "viem": "2.30.6"
+  }
 }

examples/CRISP/sdk/src/index.ts

@@ -7,5 +7,6 @@
 export * from './token'
 export * from './state'
 export * from './constants'
+export * from './utils'
 
-export type { IRoundDetails, ITokenDetails } from './types'
+export type { IRoundDetails, IRoundDetailsResponse, ITokenDetails, IMerkleProof } from './types'

examples/CRISP/sdk/src/token.ts

@@ -6,12 +6,16 @@
 
 import { CRISP_SERVER_TOKEN_TREE_ENDPOINT } from './constants'
 
+import ERC20Votes from './artifacts/ERC20Votes.json'
+import { createPublicClient, http } from 'viem'
+import { localhost, sepolia } from 'viem/chains'
+
 /**
  * Get the merkle tree data from the CRISP server
  * @param serverUrl - The base URL of the CRISP server
  * @param e3Id - The e3Id of the round
  */
-export const getTreeData = async (serverUrl: string, e3Id: number) => {
+export const getTreeData = async (serverUrl: string, e3Id: number): Promise<bigint[]> => {
   const response = await fetch(`${serverUrl}/${CRISP_SERVER_TOKEN_TREE_ENDPOINT}`, {
     method: 'POST',
     headers: {
@@ -20,17 +24,50 @@ export const getTreeData = async (serverUrl: string, e3Id: number) => {
     body: JSON.stringify({ round_id: e3Id }),
   })
 
-  const hashes = await response.json()
+  const hashes = (await response.json()) as string[]
 
-  return hashes
+  // Convert hex strings to BigInts
+  return hashes.map((hash) => {
+    // Ensure the hash is treated as a hex string
+    if (!hash.startsWith('0x')) {
+      return BigInt('0x' + hash)
+    }
+    return BigInt(hash)
+  })
 }
 
-/**
- * Generate a Merkle proof for a given address to prove inclusion in the voters' list
- */
-export const generateMerkleProof = () => {}
-
 /**
  * Get the token balance at a specific block for a given address
+ * @param voterAddress - The address of the voter
+ * @param tokenAddress - The address of the token contract
+ * @param snapshotBlock - The block number at which to get the balance
+ * @param chainId - The chain ID of the network
+ * @returns The token balance as a bigint
  */
-export const getBalanceAt = () => {}
+export const getBalanceAt = async (voterAddress: string, tokenAddress: string, snapshotBlock: number, chainId: number): Promise<bigint> => {
+  let chain
+  switch (chainId) {
+    case 11155111:
+      chain = sepolia
+      break
+    case 31337:
+      chain = localhost
+      break
+    default:
+      throw new Error('Unsupported chainId')
+  }
+
+  const publicClient = createPublicClient({
+    transport: http(),
+    chain,
+  })
+
+  const balance = (await publicClient.readContract({
+    address: tokenAddress as `0x${string}`,
+    abi: ERC20Votes.abi,
+    functionName: 'getPastVotes',
+    args: [voterAddress as `0x${string}`, BigInt(snapshotBlock)],
+  })) as bigint
+
+  return balance
+}

examples/CRISP/sdk/src/types.ts

@@ -4,6 +4,8 @@
 // without even the implied warranty of MERCHANTABILITY
 // or FITNESS FOR A PARTICULAR PURPOSE.
 
+import type { LeanIMTMerkleProof } from '@zk-kit/lean-imt'
+
 /**
  * Interface representing the details of a specific round returned by the CRISP server
  */
@@ -50,3 +52,12 @@ export interface ITokenDetails {
   threshold: bigint
   snapshotBlock: bigint
 }
+
+/**
+ * Interface representing a Merkle proof
+ */
+export interface IMerkleProof {
+  leaf: bigint
+  index: number
+  proof: LeanIMTMerkleProof<bigint>
+}

examples/CRISP/sdk/src/utils.ts

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: LGPL-3.0-only
+//
+// This file is provided WITHOUT ANY WARRANTY;
+// without even the implied warranty of MERCHANTABILITY
+// or FITNESS FOR A PARTICULAR PURPOSE.
+
+import { poseidon2 } from 'poseidon-lite'
+import { LeanIMT } from '@zk-kit/lean-imt'
+
+import type { IMerkleProof } from './types'
+
+/**
+ * Hash a leaf node for the Merkle tree
+ * @param address The voter's address
+ * @param balance The voter's balance
+ * @returns The hashed leaf as a bigint
+ */
+export const hashLeaf = (address: string, balance: string): bigint => {
+  return poseidon2([address, balance])
+}
+
+/**
+ * Generate a new LeanIMT with the leaves provided
+ * @param leaves The leaves of the Merkle tree
+ * @returns the generated Merkle tree
+ */
+export const generateMerkleTree = (leaves: bigint[]): LeanIMT => {
+  return new LeanIMT((a, b) => poseidon2([a, b]), leaves)
+}
+
+/**
+ * Generate a Merkle proof for a given address to prove inclusion in the voters' list
+ * @param threshold The minimum balance required to be eligible
+ * @param balance The voter's balance
+ * @param address The voter's address
+ * @param leaves The leaves of the Merkle tree
+ */
+export const generateMerkleProof = (threshold: number, balance: number, address: string, leaves: bigint[]): IMerkleProof => {
+  if (balance < threshold) {
+    throw new Error('Balance is below the threshold')
+  }
+
+  const leaf = hashLeaf(address, balance.toString())
+
+  const index = leaves.findIndex((l) => l === leaf)
+
+  if (index === -1) {
+    throw new Error('Leaf not found in the tree')
+  }
+
+  const tree = generateMerkleTree(leaves)
+
+  const proof = tree.generateProof(index)
+
+  return {
+    leaf,
+    index,
+    proof,
+  }
+}

examples/CRISP/sdk/tests/utils.test.ts

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: LGPL-3.0-only
+//
+// This file is provided WITHOUT ANY WARRANTY;
+// without even the implied warranty of MERCHANTABILITY
+// or FITNESS FOR A PARTICULAR PURPOSE.
+
+import { expect, describe, it, beforeAll } from 'vitest'
+import { generateMerkleProof, generateMerkleTree, hashLeaf } from '../src/utils'
+import { getTreeData } from '../src'
+import { CRISP_SERVER_URL } from './constants'
+
+describe('Utils', () => {
+  let leaves: bigint[]
+
+  beforeAll(async () => {
+    leaves = await getTreeData(CRISP_SERVER_URL, 0)
+  })
+
+  describe('hashLeaf', () => {
+    it('should return a bigint hash of the two values', () => {
+      const leaf = hashLeaf('0x1234567890123456789012345678901234567890', '1000')
+      expect(typeof leaf).toBe('bigint')
+    })
+  })
+
+  describe('generateMerkleTree', () => {
+    it('should generate a merkle tree', () => {
+      const tree = generateMerkleTree(leaves)
+      expect(tree.root).toBeDefined()
+    })
+  })
+
+  describe('generateMerkleProof', () => {
+    const address = '0x1234567890123456789012345678901234567890'
+    const balance = 1000
+    it('should generate a merkle proof for a leaf', () => {
+      const proof = generateMerkleProof(0, balance, address, leaves)
+      expect(proof.leaf).toBe(hashLeaf(address, balance.toString()))
+    })
+    it('should throw if the leaf does not exist in the tree', () => {
+      expect(() => generateMerkleProof(0, balance, address, [])).toThrow('Leaf not found in the tree')
+      expect(() => generateMerkleProof(0, 999, address, leaves)).toThrow('Leaf not found in the tree')
+    })
+  })
+})

examples/CRISP/sdk/tsconfig.json

@@ -9,5 +9,5 @@
     "outDir": "dist",
     "declaration": true
   },
-  "include": ["src/**/*.ts"]
+  "include": ["src/**/*.ts", "src/**/*.json"]
 }

examples/CRISP/server/src/server/repo.rs

@@ -244,6 +244,7 @@ impl<S: DataStore> CrispE3Repository<S> {
 
     pub async fn set_token_holder_hashes(&mut self, hashes: Vec<String>) -> Result<()> {
         let key = self.crisp_key();
+
         self.store
             .modify(&key, |e3_obj: Option<E3Crisp>| {
                 e3_obj.map(|mut e| {
@@ -259,7 +260,6 @@ impl<S: DataStore> CrispE3Repository<S> {
 
     pub async fn get_token_holder_hashes(&self) -> Result<Vec<String>> {
         let e3_crisp = self.get_crisp().await?;
-
         Ok(e3_crisp.token_holder_hashes)
     }