build_setup_multi.py

#!/usr/bin/env python3
"""
build_setup_multi.py
Compile a Circom circuit (Allowed / VerifySign / Full) and run Groth16 setup.

Default PTAU policy:
  - Allowed     : local ceremony by default (small circuits)
  - VerifySign  : download prebuilt PTAU by default
  - Full        : download prebuilt PTAU by default

Override with:
  --ptau-source download
  --ptau-source local

Download source (default template):
  https://storage.googleapis.com/zkevm/ptau/powersOfTau28_hez_final_{K}.ptau
  (You can override with --ptau-url-template)
"""

import argparse
import math
import os
import re
import secrets
import shutil
import subprocess
from pathlib import Path
from urllib.request import urlopen, Request

class BuildError(RuntimeError): ...

def run(cmd: list[str], cwd: str | None = None, stdin_input: str | None = None) -> str:
    try:
        print(">>", " ".join(cmd))
        if stdin_input is not None:
            # For commands that need stdin input (like snarkjs contribute)
            proc = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                cwd=cwd
            )
            out, _ = proc.communicate(input=stdin_input)
            if proc.returncode != 0:
                raise subprocess.CalledProcessError(proc.returncode, cmd, out)
            print(out.strip())
            return out
        else:
            out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, text=True, cwd=cwd)
            print(out.strip())
            return out
    except subprocess.CalledProcessError as e:
        raise BuildError(f"Command failed: {' '.join(cmd)}\n{e.output}") from e

def ensure_tools():
    # circom: must succeed normally
    run(["circom", "--version"])

    # snarkjs: some builds exit non-zero on --version; tolerate that
    snarkjs_path = shutil.which("snarkjs")
    if not snarkjs_path:
        raise BuildError("Required tool 'snarkjs' not found on PATH.")

    try:
        out = run(["snarkjs", "--version"])
    except BuildError as e:
        # Accept if output looks like snarkjs banner/usage
        text = str(e)
        if "snarkjs@" in text or "Usage:" in text or "Full Command" in text:
            print("snarkjs present (non-zero exit on --version is OK).")
        else:
            raise

def parse_constraints(r1cs_info_output: str) -> int | None:
    m = re.search(r"# of Constraints:\s+(\d+)", r1cs_info_output)
    return int(m.group(1)) if m else None

def choose_k(constraints: int, user_k: int | None) -> int:
    if user_k:
        return user_k
    k = math.ceil(math.log2(max(1, constraints))) + 1  # minimal + 1 headroom
    return max(12, k)

def compile_circuit(circom_file: Path, outdir: Path, libroots: list[Path], force: bool):
    circ_name = circom_file.stem
    wasm_dir = outdir / f"{circ_name}_js"
    r1cs = outdir / f"{circ_name}.r1cs"
    sym  = outdir / f"{circ_name}.sym"
    wasm = wasm_dir / f"{circ_name}.wasm"

    if force:
        if wasm_dir.exists(): shutil.rmtree(wasm_dir, ignore_errors=True)
        for p in (r1cs, sym):
            if p.exists(): p.unlink()

    if r1cs.exists() and sym.exists() and wasm.exists():
        print(f"Compile skipped ({circ_name} artifacts already exist).")
        return r1cs, wasm, sym

    cmd = ["circom", str(circom_file), "--r1cs", "--wasm", "--sym", "-o", str(outdir)]
    # IMPORTANT: pass library ROOT folders (e.g., 'circomlib'), not their 'circuits/' subfolder
    for lib in libroots:
        cmd.extend(["-l", str(lib)])
    run(cmd)
    return r1cs, wasm, sym

# -------------------- PTAU handling --------------------

def _format_bytes(num: int) -> str:
    # Simple human-readable size helper
    units = ["B", "KB", "MB", "GB"]
    value = float(num)
    for unit in units:
        if value < 1024 or unit == units[-1]:
            return f"{value:.1f} {unit}"
        value /= 1024

def _download(url: str, dest: Path):
    print(f"Downloading PTAU: {url}")
    req = Request(url, headers={"User-Agent": "build-setup/1.0"})
    chunk_size = 1 << 20  # 1 MiB
    with urlopen(req) as resp, open(dest, "wb") as f:
        total = resp.headers.get("Content-Length")
        total_int = int(total) if total and total.isdigit() else None
        downloaded = 0
        last_percent = -5

        while True:
            chunk = resp.read(chunk_size)
            if not chunk:
                break
            f.write(chunk)
            downloaded += len(chunk)

            if total_int:
                percent = int(downloaded * 100 / total_int)
                if percent - last_percent >= 5:
                    print(f"  ... {percent:3d}% ({_format_bytes(downloaded)} / {_format_bytes(total_int)})")
                    last_percent = percent
            else:
                print(f"  ... downloaded {_format_bytes(downloaded)}")

    size = dest.stat().st_size
    print(f"Saved: {dest} ({_format_bytes(size)})")

def download_ptau(k: int, dest: Path, url_template: str):
    url = url_template.format(K=k)
    _download(url, dest)
    # Verify the downloaded PTAU file
    v = run(["snarkjs", "powersoftau", "verify", str(dest)])
    if "Powers Of tau file OK" not in v and "Powers of Tau Ok" not in v:
        raise BuildError("Downloaded PTAU did not verify with snarkjs.")
    print("PTAU verified")

def prepare_pot_local(k: int, force: bool, prefix: str | None = None, ptau_dir: Path | None = None) -> Path:
    ptau_prefix = prefix or f"pot{k}"
    if ptau_dir:
        ptau_dir.mkdir(parents=True, exist_ok=True)
        pot0   = ptau_dir / f"{ptau_prefix}_0000.ptau"
        pot1   = ptau_dir / f"{ptau_prefix}_0001.ptau"
        potfin = ptau_dir / f"{ptau_prefix}_final.ptau"
    else:
        pot0   = Path(f"{ptau_prefix}_0000.ptau")
        pot1   = Path(f"{ptau_prefix}_0001.ptau")
        potfin = Path(f"{ptau_prefix}_final.ptau")

    if force:
        for p in (pot0, pot1, potfin):
            if p.exists(): p.unlink()

    if not pot0.exists():
        run(["snarkjs", "powersoftau", "new", "bn128", str(k), str(pot0), "-v"])
    else:
        print(f"Reusing {pot0}")

    if not pot1.exists():
        # Provide entropy via stdin (non-interactive mode)
        # Generate random entropy string for the contribution
        entropy = secrets.token_hex(32) + "\n"
        run(["snarkjs", "powersoftau", "contribute", str(pot0), str(pot1), "--name=First contribution", "-v"], stdin_input=entropy)
    else:
        print(f"Reusing {pot1}")

    if not potfin.exists():
        run(["snarkjs", "powersoftau", "prepare", "phase2", str(pot1), str(potfin), "-v"])
    else:
        print(f"Reusing {potfin}")

    return potfin

def obtain_ptau(k: int, source: str, url_template: str, force: bool, prefix: str | None = None, ptau_dir: Path | None = None) -> Path:
    """
    source: 'download' or 'local'
    If 'download': fetch 'powersOfTau28_hez_final_{K}.ptau' and store as 'pot{K}_final.ptau'
    If 'local'   : run local ceremony to produce 'pot{K}_final.ptau'
    """
    ptau_prefix = prefix or f"pot{k}"
    if ptau_dir:
        ptau_dir.mkdir(parents=True, exist_ok=True)
        potfin = ptau_dir / f"{ptau_prefix}_final.ptau"
    else:
        potfin = Path(f"{ptau_prefix}_final.ptau")

    if force and potfin.exists():
        potfin.unlink()

    if source == "download":
        if not potfin.exists():
            download_ptau(k, potfin, url_template)
        else:
            print(f"Reusing {potfin} (download mode)")
        return potfin
    elif source == "local":
        return prepare_pot_local(k, force=force, prefix=ptau_prefix, ptau_dir=ptau_dir)
    else:
        raise BuildError(f"Unknown ptau source: {source}")

# -------------------- Groth16 setup --------------------

def groth16_setup(r1cs: Path, potfin: Path, circ_name: str, zkey_out: Path, vk_json: Path, force: bool, contrib_name: str, setup_dir: Path):
    zkey0 = setup_dir / f"{circ_name}_0000.zkey"

    if force:
        for p in (zkey0, zkey_out, vk_json):
            if p.exists(): p.unlink()

    if not zkey0.exists():
        run(["snarkjs", "groth16", "setup", str(r1cs), str(potfin), str(zkey0)])
    else:
        print(f"Reusing {zkey0}")

    if not zkey_out.exists():
        # Provide entropy via stdin (non-interactive mode)
        entropy = secrets.token_hex(32) + "\n"
        run(["snarkjs", "zkey", "contribute", str(zkey0), str(zkey_out), f"--name={contrib_name}", "-v"], stdin_input=entropy)
    else:
        print(f"Reusing {zkey_out}")

    if not vk_json.exists():
        run(["snarkjs", "zkey", "export", "verificationkey", str(zkey_out), str(vk_json)])
    else:
        print(f"Reusing {vk_json}")


# -------------------- Main --------------------

def main():
    parser = argparse.ArgumentParser(description="Compile + Groth16 setup for Allowed / VerifySign / Freshness / Full.")
    parser.add_argument("-c", "--circuit", choices=["Allowed", "VerifySign", "Freshness", "Full"], required=True)
    parser.add_argument("--circuits-dir", default="circuits", help="Directory containing your .circom files")
    parser.add_argument("-l", "--libroot", action="append", default=[],
                        help="Library ROOT to pass to circom via -l (repeatable). Example: -l circomlib -l circom-ecdsa-p256")
    parser.add_argument("--outdir", default="circuits", help="Output directory for r1cs/wasm/sym")
    parser.add_argument("--pot-k", type=int, default=None, help="Force K for Powers of Tau (2^K). If omitted, auto-choose from constraints.")
    parser.add_argument("--ptau-source", choices=["auto", "download", "local"], default="auto",
                        help="PTAU acquisition mode. 'auto' defaults to download for VerifySign/Full, local for Allowed.")
    parser.add_argument("--ptau-url-template",
                        default="https://storage.googleapis.com/zkevm/ptau/powersOfTau28_hez_final_{K}.ptau",
                        help="URL template for downloading PTAU (use {K} placeholder).")
    parser.add_argument("--ptau-prefix", default=None, help="Optional PTAU local filename prefix (default: pot{K})")
    parser.add_argument("--contrib-name", default="1st Contributor", help="Name for zkey contribution")
    parser.add_argument("--force", action="store_true", help="Rebuild / overwrite existing artifacts")
    args = parser.parse_args()

    ensure_tools()

    # Map circuit choice → file path
    circuits_dir = Path(args.circuits_dir)
    circ_file = circuits_dir / f"{args.circuit}.circom"
    if not circ_file.exists():
        raise BuildError(f"Circuit not found: {circ_file}")

    # Library roots (default per circuit if user didn’t pass any)
    default_libs = {
        "Allowed": ["circomlib"],
        "Freshness": ["circomlib"],
        "VerifySign": ["circomlib", "circom-ecdsa-p256"],
        "Full": ["circomlib", "circom-ecdsa-p256"],
    }
    if args.libroot:
        libroots = [Path(p) for p in args.libroot]
    else:
        libroots = [Path(p) for p in default_libs.get(args.circuit, ["circomlib"])]

    for lib in libroots:
        if not lib.exists():
            raise BuildError(f"Library root not found: {lib}  (tip: pass the ROOT folder, not its 'circuits/' subfolder)")

    # Create setup/<circuit_name>/ directory structure
    circ_name = circ_file.stem  # e.g., Allowed / VerifySign / Full
    setup_dir = Path("setup") / circ_name
    setup_dir.mkdir(parents=True, exist_ok=True)
    
    # PTAU files go in setup/ptau/ (shared across circuits)
    ptau_dir = Path("setup") / "ptau"
    ptau_dir.mkdir(parents=True, exist_ok=True)

    # Use setup/<circuit_name>/ as the output directory
    outdir = setup_dir
    print(f"Output directory: {outdir}")

    # 1) Compile
    r1cs, wasm, sym = compile_circuit(circ_file, outdir, libroots, force=args.force)

    # 2) Inspect constraints & decide K
    info_out = run(["snarkjs", "r1cs", "info", str(r1cs)])
    constraints = parse_constraints(info_out) or 0
    chosen_k = choose_k(constraints, args.pot_k)
    print(f"Chosen K = {chosen_k} (2^K = {1<<chosen_k}) for constraints ≈ {constraints}")

    # 3) PTAU source selection
    if args.ptau_source == "auto":
        ptau_mode = "download" if args.circuit in ("VerifySign", "Full") else "local"
    else:
        ptau_mode = args.ptau_source
    print(f"PTAU mode: {ptau_mode}")

    potfin = obtain_ptau(
        k=chosen_k,
        source=ptau_mode,
        url_template=args.ptau_url_template,
        force=args.force,
        prefix=args.ptau_prefix,
        ptau_dir=ptau_dir
    )

    # 4) Groth16 setup → zkey + vk (per circuit name)
    # Save zkey files in setup/<circuit_name>/
    zkey_out = setup_dir / f"{circ_name}_final.zkey"
    vk_json  = setup_dir / f"{circ_name}_verification_key.json"
    groth16_setup(r1cs, potfin, circ_name, zkey_out, vk_json, force=args.force, contrib_name=args.contrib_name, setup_dir=setup_dir)

    print("\nSetup complete.")
    print(f" Circuit    : {circ_name}")
    print(f" R1CS       : {r1cs}")
    print(f" WASM       : {wasm}")
    print(f" SYM        : {sym}")
    print(f" PTAU       : {potfin}")
    print(f" ZKey       : {zkey_out}")
    print(f" VK JSON    : {vk_json}")

if __name__ == "__main__":
    main()