Privacy dApps with ZK

import { useAccount, useWriteContract } from 'wagmi'
import { createPublicClient, http } from 'viem'
import { scrollSepolia } from 'viem/chains'
import { useState, useEffect } from 'react'
import { groth16 } from "snarkjs";

// Custom logic contract we just deployed
const CONTRACT_ADDRESS = import.meta.env.VITE_CIRCOM_CUSTOM_LOGIC_CONTRACT_ADDRESS as `0x${string}`

// In this tutorial we'll use Scroll Sepolia, but this code can also be used on Scroll Mainnet or any EVM chain that supports all precompiles needed for ZK proof
const publicClient = createPublicClient({
  chain: scrollSepolia,
  transport: http()
})

// Custom logic contract ABI to be able to send proofs on-chain
const verifierAbi = [
  {
    "inputs": [
      {
        "internalType": "uint256[2]",
        "name": "_pA",
        "type": "uint256[2]"
      },
      {
        "internalType": "uint256[2][2]",
        "name": "_pB",
        "type": "uint256[2][2]"
      },
      {
        "internalType": "uint256[2]",
        "name": "_pC",
        "type": "uint256[2]"
      },
      {
        "internalType": "uint256[1]",
        "name": "_pubSignals",
        "type": "uint256[1]"
      }
    ],
    "name": "sendProof",
    "outputs": [],
    "stateMutability": "nonpayable",
    "type": "function"
  },
  {
    "inputs": [],
    "name": "publicInput",
    "outputs": [
      {
        "internalType": "uint256",
        "name": "",
        "type": "uint256"
      }
    ],
    "stateMutability": "view",
    "type": "function"
  }
] as const

// Function responsible to generate private computation on the browser and then sending the computation proof on-chain
function useZKForm() {
  const [a, setA] = useState<string>("")
  const [b, setB] = useState<string>("")
  const [message, setMessage] = useState<string>("Connect your wallet")
  const [publicInput, setPublicInput] = useState<number | null>(null)
  const { writeContract, isPending: isSendingProof } = useWriteContract()

  // Reads the contract state, in this case the multiplication result that was previously submmited
  const fetchPublicInput = async () => {
    try {
      const result = await publicClient.readContract({
        address: CONTRACT_ADDRESS,
        abi: verifierAbi,
        functionName: "publicInput",
      })
      setPublicInput(Number(result))
    } catch (err) {
      console.error("Error fetching public input:", err)
    }
  }

  // Generates the ZK proof based on private parameters, and then sends the proofs to the smart contract on Scroll
  const sendProof = async () => {
    setMessage("Generating proof...")

    // Let's start by generating a proof by passing private inputs and the ZK artifacts generated by the circom compiler
    // Notice the ZK artifacts should be shared publicly in your website
    const { proof, publicSignals } = await groth16.fullProve(
      { a: Number(a), b: Number(b) },
      "./zk_artifacts/myCircuit.wasm",
      "./zk_artifacts/myCircuit_final.zkey"
    )

    // Now let's verify the proof locally, this is optional
    setMessage("Verifying off‑chain...")
    const vkey = await fetch("/verification_key.json").then((r) =>
      r.json()
    )
    const valid = await groth16.verify(vkey, publicSignals, proof)
    if (!valid) {
      setMessage("Proof verification failed")
      return
    }

    // We now translate the proof in a format compatible with what our smart contract expects
    const pA = proof.pi_a.slice(0, 2)  // Take first two elements
    const pB = proof.pi_b.slice(0, 2).map((row: string[]) => row.slice(0, 2))
    const pC = proof.pi_c.slice(0, 2)

    // Once the proof is ready we send it on-chain
    setMessage("Proof generated please confirm transaction.")
    writeContract({
      address: CONTRACT_ADDRESS,
      abi: verifierAbi,
      functionName: "sendProof",
      args: [pA, pB, pC, publicSignals],
    }, {
      onSuccess: (txHash) => {
        setMessage("Executing...")
        publicClient.waitForTransactionReceipt({ hash: txHash })
          .then(() => {
            setMessage("Success!")
            fetchPublicInput()
          })
          .catch((err) => {
            console.error("ERROR! Transaction reverted:", err)
            setMessage("Transaction failed")
          })
      },
      onError: (err) => {
        console.error("ERROR! Transaction reverted:", err)
        setMessage("Transaction failed")
      }
    })
  }

  return {
    a,
    setA,
    b,
    setB,
    message,
    publicInput,
    isSendingProof,
    sendProof,
    fetchPublicInput
  }
}

// React UI that handles contract reads, and private inputs for the ZK proof generation
export default function ZKForm() {
  const account = useAccount()
  const {
    a, setA,
    b, setB,
    message,
    publicInput,
    isSendingProof,
    sendProof,
    fetchPublicInput
  } = useZKForm()

  useEffect(() => {
    if (account.status === 'connected') {
      fetchPublicInput()
    }
  }, [account.status, fetchPublicInput])

  return (
    <div className="p-4 max-w-md mx-auto">
      <h2 className="text-xl font-semibold mb-2">ZK Multiplication</h2>
      <p className="mb-4 text-sm text-gray-600">{message}</p>
      <input
        className="border p-2 mb-2 w-full"
        type="number"
        placeholder="a"
        value={a}
        onChange={(e) => setA(e.target.value)}
      />
      <input
        className="border p-2 mb-4 w-full"
        type="number"
        placeholder="b"
        value={b}
        onChange={(e) => setB(e.target.value)}
      />
      {account.status === 'connected' && (
        <button
          onClick={sendProof}
          disabled={isSendingProof}
          className="bg-blue-600 text-white px-4 py-2 rounded disabled:opacity-50"
        >
          {isSendingProof ? 'Sending...' : 'Send Proof'}
        </button>
      )}
      {publicInput !== null && (
        <p className="mt-4 text-sm text-gray-700">
          Last stored result: {publicInput}
        </p>
      )}
    </div>
  )
}