feat: Implement Halo2 Merkle Tree (#144)

* Add halo2 scaffold from https://github.com/axiom-crypto/halo2-scaffold/

* Delete unnecessary scaffolding.

* Add DualMux.

* Complete merkle tree.

* Remove more of the scaffold.

* Fix formatting.

circuits/halo2/.gitignore

@@ -0,0 +1,17 @@
+# Generated by Cargo
+# will have compiled files and executables
+/target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+
+# Added by cargo
+
+/target
+
+/params

circuits/halo2/Cargo.toml

@@ -0,0 +1,62 @@
+[package]
+name = "halo2-scaffold"
+version = "0.2.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+rand = "0.8"
+ark-std = { version = "0.3.0", features = ["print-trace"] }
+serde = { version = "1.0", default-features = false, features = ["derive"] }
+serde_json = "1.0"
+log = "0.4"
+env_logger = "0.10"
+clap = { version = "4.1", features = ["derive"] }
+clap-num = "1.0.2"
+
+# halo2
+halo2_proofs = { git = "https://github.com/privacy-scaling-explorations/halo2.git", tag = "v2023_02_02" }
+
+# Axiom's helper API with basic functions 
+halo2-base = { git = "https://github.com/axiom-crypto/halo2-lib", branch = "axiom-dev-0406" }
+# Axiom poseidon chip (adapted from Scroll)
+poseidon = { git = "https://github.com/axiom-crypto/halo2-lib", branch = "axiom-dev-0406" }
+# Axiom Evm wrapper 
+axiom-eth = { git = "https://github.com/axiom-crypto/axiom-eth.git", branch = "axiom-dev-0406", default-features = false, features = ["halo2-axiom", "aggregation", "evm", "clap"] }
+snark-verifier-sdk = { git = "https://github.com/axiom-crypto/snark-verifier.git", branch = "axiom-dev-0406", default-features = false, features = ["loader_halo2"] }
+
+
+[dev-dependencies]
+test-log = "0.2.11"
+
+[features]
+default = []
+
+# Dev / testing mode. We make opt-level = 3 to improve proving times (otherwise it is really slow)
+[profile.dev]
+opt-level = 3
+debug = 1 # change to 0 or 2 for more or less debug info
+overflow-checks = true # default
+incremental = true # default
+
+# Local "release" mode, more optimized than dev but faster to compile than release
+[profile.local]
+inherits = "dev"
+opt-level = 3
+# Set this to 1 or 2 to get more useful backtraces
+debug = 1
+debug-assertions = false
+panic = 'unwind'
+# better recompile times
+incremental = true
+lto = "thin"
+codegen-units = 16
+
+[profile.release]
+opt-level = 3
+debug = false
+debug-assertions = false
+lto = "fat"
+panic = "abort"
+incremental = false

circuits/halo2/README.md

@@ -0,0 +1,3 @@
+Built from base repository https://github.com/axiom-crypto/halo2-scaffold/
+
+To run the poseidon merkle tree run `cargo run --example merkle -- --name merkle -k 16 mock`

circuits/halo2/data/merkle.in

@@ -0,0 +1,5 @@
+{
+    "leaf": "7",
+    "path": ["2", "4", "8", "16", "32"],
+    "switch": ["1", "0", "1", "1", "0"]
+}

circuits/halo2/examples/merkle.rs

@@ -0,0 +1,76 @@
+use clap::Parser;
+use halo2_base::{
+    gates::GateChip, gates::GateInstructions, utils::ScalarField, AssignedValue, Context,
+    QuantumCell,
+};
+use halo2_scaffold::scaffold::{cmd::Cli, run};
+use poseidon::PoseidonChip;
+use serde::{Deserialize, Serialize};
+
+const T: usize = 3;
+const RATE: usize = 2;
+const R_F: usize = 8;
+const R_P: usize = 57;
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct CircuitInput {
+    pub leaf: String,
+    pub path: [String; 5],
+    pub switch: [String; 5],
+}
+
+// Verifies a merkle proof
+fn merkle_proof<F: ScalarField>(
+    ctx: &mut Context<F>,
+    inp: CircuitInput,
+    make_public: &mut Vec<AssignedValue<F>>,
+) {
+    // Load variables as private inputs
+    let leaf = ctx.load_witness(F::from_str_vartime(&inp.leaf).unwrap());
+    let path =
+        inp.path.map(|elem| ctx.load_witness(F::from_str_vartime(&elem).unwrap())).into_iter();
+    let switch =
+        inp.switch.map(|elem| ctx.load_witness(F::from_str_vartime(&elem).unwrap())).into_iter();
+    let gate = GateChip::<F>::default();
+
+    let mut next_hash = leaf;
+    for (s, p) in switch.zip(path) {
+        gate.assert_bit(ctx, s);
+        let [a, b] = dual_mux(&gate, ctx, next_hash, p, s);
+
+        let mut poseidon = PoseidonChip::<F, T, RATE>::new(ctx, R_F, R_P).unwrap();
+        poseidon.update(&[a, b]);
+        next_hash = poseidon.squeeze(ctx, &gate).unwrap();
+        println!("a: {:?}, b: {:?}, poseidon(a,b)): {:?}", a.value(), b.value(), next_hash.value());
+    }
+    make_public.push(next_hash); // make the root public
+    println!("Merkle root: {:?}", next_hash.value());
+
+    // TODO: test that the output is the same as circom
+}
+
+fn dual_mux<F: ScalarField>(
+    gate: &GateChip<F>,
+    ctx: &mut Context<F>,
+    a: impl Into<QuantumCell<F>>,
+    b: impl Into<QuantumCell<F>>,
+    switch: impl Into<QuantumCell<F>>,
+) -> [AssignedValue<F>; 2] {
+    let a = a.into();
+    let b = b.into();
+    let switch = switch.into();
+
+    let a_sub_b = gate.sub(ctx, a, b);
+    let b_sub_a = gate.sub(ctx, b, a);
+
+    let left = gate.mul_add(ctx, b_sub_a, switch, a); // left = (b-a)*s + a;
+    let right = gate.mul_add(ctx, a_sub_b, switch, b); // right = (a-b)*s + b;
+    [left, right]
+}
+
+fn main() {
+    env_logger::init();
+
+    let args = Cli::parse();
+    run(merkle_proof, args);
+}

circuits/halo2/rust-toolchain

@@ -0,0 +1 @@
+nightly-2022-10-28

circuits/halo2/rustfmt.toml

@@ -0,0 +1,2 @@
+max_width = 100
+use_small_heuristics = "Max" 

circuits/halo2/src/lib.rs

@@ -0,0 +1,4 @@
+#![feature(return_position_impl_trait_in_trait)]
+#![allow(incomplete_features)]
+
+pub mod scaffold;

circuits/halo2/src/scaffold/cmd.rs

@@ -0,0 +1,45 @@
+use clap::{Parser, Subcommand};
+use std::path::PathBuf;
+
+#[derive(Clone, Copy, Debug, Subcommand)]
+pub enum SnarkCmd {
+    /// Run the mock prover
+    Mock,
+    /// Generate new proving & verifying keys
+    Keygen,
+    /// Generate a new proof
+    Prove,
+    /// Verify a proof
+    Verify,
+}
+
+impl std::fmt::Display for SnarkCmd {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Mock => write!(f, "mock"),
+            Self::Keygen => write!(f, "keygen"),
+            Self::Prove => write!(f, "prove"),
+            Self::Verify => write!(f, "verify"),
+        }
+    }
+}
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+/// Command-line helper for various steps in ZK proving.
+pub struct Cli {
+    #[command(subcommand)]
+    pub command: SnarkCmd,
+    #[arg(short, long = "name")]
+    pub name: String,
+    #[arg(short = 'k', long = "degree")]
+    pub degree: u32,
+    #[arg(short, long = "input")]
+    pub input_path: Option<PathBuf>,
+    #[arg(long = "create-contract")]
+    pub create_contract: bool,
+    #[arg(short, long = "config-path")]
+    pub config_path: Option<PathBuf>,
+    #[arg(short, long = "data-path")]
+    pub data_path: Option<PathBuf>,
+}