Merge pull request #11 from Axis-Fi/merkle-improvements

Improvements to the merkle tree CLI

packages/oz-merkle-tree/README.md

@@ -9,7 +9,7 @@ bun install
 To run:
 
 ```bash
-bun run src/index.ts
+bun run src/index.ts file.csv
 ```
 
 This project was created using `bun init` in bun v1.1.8. [Bun](https://bun.sh) is a fast all-in-one JavaScript runtime.

packages/oz-merkle-tree/package.json

@@ -10,7 +10,8 @@
     "typescript": "^5.0.0"
   },
   "dependencies": {
-    "@openzeppelin/merkle-tree": "^1.0.6"
+    "@openzeppelin/merkle-tree": "^1.0.6",
+    "csv-parser": "^3.0.0"
   },
   "scripts": {
     "fmt:check": "prettier --check .",

packages/oz-merkle-tree/src/index.ts

@@ -1,41 +1,93 @@
+import fs from "fs";
+import csv from "csv-parser";
 import { StandardMerkleTree } from "@openzeppelin/merkle-tree";
 
-// TODOs
-// - Specify the types of the values
-// - Pass in the values from the CLI or file
+type Proof = {
+  root: string;
+  entries: {
+    value: string[];
+    proof: string;
+  }[];
+};
 
 const generateMerkleTree = (name: string, values: any[][], types: string[]) => {
   const tree = StandardMerkleTree.of(values, types);
 
-  console.log(name);
-
   console.log(`Merkle root: ${tree.root}`);
 
-  console.log("Proofs");
+  const proofs: Proof = {
+    root: tree.root,
+    entries: [],
+  };
 
   for (const [i, v] of tree.entries()) {
     const proof = tree.getProof(i);
-    console.log("Value:", v);
-    console.log("Proof:", proof);
+
+    proofs.entries.push({
+      value: v as string[],
+      proof: proof[0],
+    });
   }
 
-  Bun.write(`out/${name}.json`, JSON.stringify(tree.dump(), null, 2));
+  Bun.write(`out/${name}-tree.json`, JSON.stringify(tree.dump(), null, 2));
+  Bun.write(`out/${name}-proofs.json`, JSON.stringify(proofs, null, 2));
 
   console.log("");
+  console.log(`Merkle tree written to out/${name}-tree.json`);
+  console.log(`Proofs written to out/${name}-proofs.json`);
 };
 
-const addressValues = [
-  ["0x0000000000000000000000000000000000000004"],
-  ["0x0000000000000000000000000000000000000005"],
-];
-
-const allocatedAddressValues = [
-  ["0x0000000000000000000000000000000000000004", "5000000000000000000"],
-  ["0x0000000000000000000000000000000000000020", "0"],
-];
-
-generateMerkleTree("address", addressValues, ["address"]);
-generateMerkleTree("allocated-address", allocatedAddressValues, [
-  "address",
-  "uint256",
-]);
+const csvFilePath = process.argv[2];
+
+if (!csvFilePath) {
+  console.error(
+    "Please provide the path to the CSV file as a command-line argument.",
+  );
+  process.exit(1);
+}
+
+let csvHeaders: string[] = [];
+let headersValidated = false;
+const values: string[][] = [];
+
+// Get the filename from the path, without the file extension
+const fileName = csvFilePath.split("/").pop()?.split(".")[0];
+if (!fileName) {
+  console.error("Could not extract filename from the provided path.");
+  process.exit(1);
+}
+
+fs.createReadStream(csvFilePath)
+  .pipe(csv())
+  .on("headers", (headers) => {
+    if (
+      headers.length < 1 ||
+      headers.length > 2 ||
+      headers[0] !== "address" ||
+      (headers.length === 2 && headers[1] !== "amount")
+    ) {
+      console.error(
+        `CSV headers do not match the expected headers: "address" or "address,amount"`,
+      );
+      process.exit(1);
+    }
+    headersValidated = true;
+    csvHeaders = headers;
+  })
+  .on("data", (row) => {
+    if (!headersValidated) {
+      console.error("Cannot process data without header validation");
+      process.exit(1);
+    }
+
+    const rowValues =
+      csvHeaders.length === 1 ? [row.address] : [row.address, row.amount];
+    values.push(rowValues);
+  })
+  .on("end", () => {
+    console.log("CSV file successfully processed");
+    const types =
+      csvHeaders.length === 1 ? ["address"] : ["address", "uint256"];
+
+    generateMerkleTree(fileName, values, types);
+  });

packages/oz-merkle-tree/test.csv

@@ -0,0 +1,3 @@
+address,amount
+0x0000000000000000000000000000000000000004,5000000000000000000
+0x0000000000000000000000000000000000000020,0