From 873ec52f0197c4998144ddaa5366a3f8cbfdcf6f Mon Sep 17 00:00:00 2001
From: ioedeveloper <daviddisu8@gmail.com>
Date: Tue, 27 Aug 2024 12:25:34 +0100
Subject: [PATCH] Export witness as json

---
 .../circuit-compiler/src/app/actions/index.ts | 14 +++++++++--
 .../src/app/components/container.tsx          |  2 +-
 .../src/app/components/witness.tsx            | 25 +++++++++++++++----
 .../src/app/reducers/state.ts                 |  9 ++++++-
 .../src/app/services/circomPluginClient.ts    |  3 ++-
 apps/circuit-compiler/src/app/types/index.ts  |  2 ++
 .../src/app/tabs/locales/en/circuit.json      |  1 +
 7 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/apps/circuit-compiler/src/app/actions/index.ts b/apps/circuit-compiler/src/app/actions/index.ts
index 2f48faf2fb..67ac5813d4 100644
--- a/apps/circuit-compiler/src/app/actions/index.ts
+++ b/apps/circuit-compiler/src/app/actions/index.ts
@@ -19,12 +19,20 @@ export const compileCircuit = async (plugin: CircomPluginClient, appState: AppSt
   }
 }
 
-export const computeWitness = async (plugin: CircomPluginClient, status: string, witnessValues: Record<string, string>) => {
+export const computeWitness = async (plugin: CircomPluginClient, appState: AppState, dispatch: ICircuitAppContext['dispatch'], status: string, witnessValues: Record<string, string>) => {
   try {
     if (status !== "computing") {
       const input = JSON.stringify(witnessValues)
+      const witness = await plugin.computeWitness(input)
 
-      await plugin.computeWitness(input)
+      if (appState.exportWtnsJson) {
+        const wtns = await snarkjs.wtns.exportJson(witness)
+        const wtnsJson = wtns.map(wtn => wtn.toString())
+        const fileName = extractNameFromKey(appState.filePath)
+        const writePath = extractParentFromKey(appState.filePath) + `/.bin/${fileName.replace('.circom', '.wtn.json')}`
+
+        await plugin.call('fileManager', 'writeFile', writePath, JSON.stringify(wtnsJson, null, 2))
+      }
     } else {
       console.log('Existing witness computation in progress')
     }
@@ -110,6 +118,7 @@ export const generateProof = async (plugin: CircomPluginClient, appState: AppSta
       const { proof, publicSignals } = await snarkjs.groth16.prove(zkey_final, wtns, zkLogger(plugin, dispatch, 'SET_PROOF_FEEDBACK'))
       const verified = await snarkjs.groth16.verify(vKey, publicSignals, proof, zkLogger(plugin, dispatch, 'SET_PROOF_FEEDBACK'))
 
+      plugin.call('fileManager', 'writeFile', `${extractParentFromKey(appState.filePath)}/groth16/zk/build/proof.json`, JSON.stringify(proof, null, 2))
       plugin.call('terminal', 'log', { type: 'log', value: 'zk proof validity ' + verified })
       if (appState.exportVerifierCalldata) {
         const calldata = await snarkjs.groth16.exportSolidityCallData(proof, publicSignals)
@@ -120,6 +129,7 @@ export const generateProof = async (plugin: CircomPluginClient, appState: AppSta
       const { proof, publicSignals } = await snarkjs.plonk.prove(zkey_final, wtns, zkLogger(plugin, dispatch, 'SET_PROOF_FEEDBACK'))
       const verified = await snarkjs.plonk.verify(vKey, publicSignals, proof, zkLogger(plugin, dispatch, 'SET_PROOF_FEEDBACK'))
 
+      plugin.call('fileManager', 'writeFile', `${extractParentFromKey(appState.filePath)}/plonk/zk/build/proof.json`, JSON.stringify(proof, null, 2))
       plugin.call('terminal', 'log', { type: 'log', value: 'zk proof validity ' + verified })
       if (appState.exportVerifierCalldata) {
         const calldata = await snarkjs.plonk.exportSolidityCallData(proof, publicSignals)
diff --git a/apps/circuit-compiler/src/app/components/container.tsx b/apps/circuit-compiler/src/app/components/container.tsx
index 982395d11e..3a18465fa4 100644
--- a/apps/circuit-compiler/src/app/components/container.tsx
+++ b/apps/circuit-compiler/src/app/components/container.tsx
@@ -140,7 +140,7 @@ export function Container () {
             <RenderIf condition={circuitApp.appState.signalInputs.length > 0}>
               <Toggler title='circuit.computeWitness' dataId='witness_toggler' show={!!circuitApp.appState.setupExportStatus}>
                 <>
-                  <WitnessSection plugin={circuitApp.plugin} signalInputs={circuitApp.appState.signalInputs} status={circuitApp.appState.status} />
+                  <WitnessSection />
                   <RenderIf condition={circuitApp.appState.status !== 'computing'}>
                     <CompilerFeedback feedback={circuitApp.appState.computeFeedback} filePathToId={circuitApp.appState.filePathToId} openErrorLocation={handleOpenErrorLocation} hideWarnings={circuitApp.appState.hideWarnings} askGPT={askGPT} />
                   </RenderIf>
diff --git a/apps/circuit-compiler/src/app/components/witness.tsx b/apps/circuit-compiler/src/app/components/witness.tsx
index 0a14537cde..dd40b9b2d0 100644
--- a/apps/circuit-compiler/src/app/components/witness.tsx
+++ b/apps/circuit-compiler/src/app/components/witness.tsx
@@ -1,12 +1,14 @@
 import { RenderIf, RenderIfNot } from "@remix-ui/helper";
 import { FormattedMessage } from "react-intl";
-import { CompilerStatus } from "../types";
 import { computeWitness } from "../actions";
-import { useState } from "react";
-import type { CircomPluginClient } from "../services/circomPluginClient";
+import { useContext, useState } from "react";
 import * as remixLib from '@remix-project/remix-lib'
+import { CircuitAppContext } from "../contexts";
 
-export function WitnessSection ({ plugin, signalInputs, status }: {plugin: CircomPluginClient, signalInputs: string[], status: CompilerStatus}) {
+export function WitnessSection () {
+  const circuitApp = useContext(CircuitAppContext)
+  const { signalInputs, status, exportWtnsJson } = circuitApp.appState
+  const { plugin, dispatch, appState } = circuitApp
   const [witnessValues, setWitnessValues] = useState<Record<string, string>>({})
 
   const handleSignalInput = (e: any) => {
@@ -47,9 +49,22 @@ export function WitnessSection ({ plugin, signalInputs, status }: {plugin: Circo
               </div>
             ))
           }
+          <div className="custom-control custom-checkbox">
+            <input
+              className="custom-control-input"
+              type="checkbox"
+              title="Export Witness As JSON"
+              id="circuitExportWtnsJson"
+              onChange={() => { dispatch({ type: 'SET_EXPORT_WTNS_JSON', payload: !exportWtnsJson }) }}
+              checked={exportWtnsJson}
+            />
+            <label className="form-check-label custom-control-label pt-1" htmlFor="circuitExportWtnsJson">
+              <FormattedMessage id="circuit.exportWtnsJson" />
+            </label>
+          </div>
           <button
             className="btn btn-secondary btn-block d-block w-100 text-break mb-1 mt-1"
-            onClick={() => { computeWitness(plugin, status, witnessValues) }}
+            onClick={() => { computeWitness(plugin, appState, dispatch, status, witnessValues) }}
             disabled={(status === "compiling") || (status === "computing")}
             data-id="compute_witness_btn"
           >
diff --git a/apps/circuit-compiler/src/app/reducers/state.ts b/apps/circuit-compiler/src/app/reducers/state.ts
index 78707da38e..3fead45bfc 100644
--- a/apps/circuit-compiler/src/app/reducers/state.ts
+++ b/apps/circuit-compiler/src/app/reducers/state.ts
@@ -21,8 +21,9 @@ export const appInitialState: AppState = {
   ptauList: PTAU_LIST,
   ptauValue: "final_14.ptau",
   exportVerificationContract: true,
-  exportVerificationKey: true,
+  exportVerificationKey: false,
   exportVerifierCalldata: true,
+  exportWtnsJson: false,
   verificationKey: null,
   zKey: null
 }
@@ -138,6 +139,12 @@ export const appReducer = (state = appInitialState, action: Actions): AppState =
       setupExportStatus: action.payload
     }
 
+  case 'SET_EXPORT_WTNS_JSON':
+    return {
+      ...state,
+      exportWtnsJson: action.payload
+    }
+
   case 'SET_VERIFICATION_KEY':
     return {
       ...state,
diff --git a/apps/circuit-compiler/src/app/services/circomPluginClient.ts b/apps/circuit-compiler/src/app/services/circomPluginClient.ts
index f6817fe7a5..4998dbd0ae 100644
--- a/apps/circuit-compiler/src/app/services/circomPluginClient.ts
+++ b/apps/circuit-compiler/src/app/services/circomPluginClient.ts
@@ -237,7 +237,7 @@ export class CircomPluginClient extends PluginClient {
     }
   }
 
-  async computeWitness (input: string): Promise<void> {
+  async computeWitness (input: string): Promise<Uint8Array> {
     this.internalEvents.emit('circuit_computing_witness_start')
     this.emit('statusChanged', { key: 'loading', title: 'Computing...', type: 'info' })
     const wasmPath = this.lastCompiledCircuitPath
@@ -252,6 +252,7 @@ export class CircomPluginClient extends PluginClient {
     this._paq.push(['trackEvent', 'circuit-compiler', 'computeWitness', 'Witness computing successful'])
     this.internalEvents.emit('circuit_computing_witness_done')
     this.emit('statusChanged', { key: 'succeed', title: 'witness computed successfully', type: 'success' })
+    return witness
   }
 
   async resolveDependencies(filePath: string, fileContent: string, output?: Record<string, string>, depPath: string = '', blackPath: string[] = []): Promise<Record<string, string>> {
diff --git a/apps/circuit-compiler/src/app/types/index.ts b/apps/circuit-compiler/src/app/types/index.ts
index de36f80513..b1ad621b94 100644
--- a/apps/circuit-compiler/src/app/types/index.ts
+++ b/apps/circuit-compiler/src/app/types/index.ts
@@ -39,6 +39,7 @@ export interface ActionPayloadTypes {
   SET_EXPORT_VERIFICATION_CONTRACT: boolean,
   SET_EXPORT_VERIFICATION_KEY: boolean,
   SET_EXPORT_VERIFIER_CALLDATA: boolean,
+  SET_EXPORT_WTNS_JSON: boolean,
   SET_SETUP_EXPORT_STATUS: SetupExportStatus,
   SET_VERIFICATION_KEY: Record<string, any>,
   SET_ZKEY: any
@@ -71,6 +72,7 @@ export interface AppState {
   exportVerificationContract: boolean,
   exportVerificationKey: boolean,
   exportVerifierCalldata: boolean,
+  exportWtnsJson: boolean,
   verificationKey: Record<string, any>,
   zKey: Uint8Array
 }
diff --git a/apps/remix-ide/src/app/tabs/locales/en/circuit.json b/apps/remix-ide/src/app/tabs/locales/en/circuit.json
index 7d9fcc8428..4cf99ba63b 100644
--- a/apps/remix-ide/src/app/tabs/locales/en/circuit.json
+++ b/apps/remix-ide/src/app/tabs/locales/en/circuit.json
@@ -21,5 +21,6 @@
   "circuit.exportVerifierContract": "Export verifier contract",
   "circuit.exportVerificationKey": "Export verification key",
   "circuit.exportVerifierCalldata": "Export verifier calldata",
+  "circuit.exportWtnsJson": "Export witness as JSON",
   "circuit.runSetup": "Run setup"
 }