diff --git a/codex/codex.nim b/codex/codex.nim index b8905205..2fb1b286 100644 --- a/codex/codex.nim +++ b/codex/codex.nim @@ -289,8 +289,9 @@ proc new*( store = NetworkStore.new(engine, repoStore) prover = if config.prover: - let backend = - config.initializeBackend().expect("Unable to create prover backend.") + let backend = config.initializeBackend(taskpool = taskpool).expect( + "Unable to create prover backend." + ) some Prover.new(store, backend, config.numProofSamples) else: none Prover diff --git a/codex/slots/proofs/backendfactory.nim b/codex/slots/proofs/backendfactory.nim index 7aba27d8..34706191 100644 --- a/codex/slots/proofs/backendfactory.nim +++ b/codex/slots/proofs/backendfactory.nim @@ -2,6 +2,7 @@ import os import strutils import pkg/chronos import pkg/chronicles +import pkg/taskpools import pkg/questionable import pkg/confutils/defs import pkg/stew/io2 @@ -11,7 +12,9 @@ import ../../conf import ./backends import ./backendutils -proc initializeFromConfig(config: CodexConf, utils: BackendUtils): ?!AnyBackend = +proc initializeFromConfig( + config: CodexConf, utils: BackendUtils, taskpool: Taskpool +): ?!AnyBackend = if not fileAccessible($config.circomR1cs, {AccessFlags.Read}) or not endsWith($config.circomR1cs, ".r1cs"): return failure("Circom R1CS file not accessible") @@ -27,7 +30,7 @@ proc initializeFromConfig(config: CodexConf, utils: BackendUtils): ?!AnyBackend trace "Initialized prover backend from cli config" success( utils.initializeCircomBackend( - $config.circomR1cs, $config.circomWasm, $config.circomZkey + $config.circomR1cs, $config.circomWasm, $config.circomZkey, taskpool ) ) @@ -41,14 +44,14 @@ proc zkeyFilePath(config: CodexConf): string = config.circuitDir / "proof_main.zkey" proc initializeFromCircuitDirFiles( - config: CodexConf, utils: BackendUtils + config: CodexConf, utils: BackendUtils, taskpool: Taskpool ): ?!AnyBackend {.gcsafe.} = if fileExists(config.r1csFilePath) and fileExists(config.wasmFilePath) and fileExists(config.zkeyFilePath): trace "Initialized prover backend from local files" return success( utils.initializeCircomBackend( - config.r1csFilePath, config.wasmFilePath, config.zkeyFilePath + config.r1csFilePath, config.wasmFilePath, config.zkeyFilePath, taskpool ) ) @@ -68,11 +71,11 @@ proc suggestDownloadTool(config: CodexConf) = instructions proc initializeBackend*( - config: CodexConf, utils: BackendUtils = BackendUtils() + config: CodexConf, utils: BackendUtils = BackendUtils(), taskpool: Taskpool ): ?!AnyBackend = - without backend =? initializeFromConfig(config, utils), cliErr: + without backend =? initializeFromConfig(config, utils, taskpool), cliErr: info "Could not initialize prover backend from CLI options...", msg = cliErr.msg - without backend =? initializeFromCircuitDirFiles(config, utils), localErr: + without backend =? initializeFromCircuitDirFiles(config, utils, taskpool), localErr: info "Could not initialize prover backend from circuit dir files...", msg = localErr.msg suggestDownloadTool(config) diff --git a/codex/slots/proofs/backends/circomcompat.nim b/codex/slots/proofs/backends/circomcompat.nim index 1d2e3e19..dc4cc440 100644 --- a/codex/slots/proofs/backends/circomcompat.nim +++ b/codex/slots/proofs/backends/circomcompat.nim @@ -9,9 +9,11 @@ {.push raises: [].} -import std/sugar +import std/[sugar, atomics] import pkg/chronos +import pkg/taskpools +import pkg/chronos/threadsync import pkg/questionable/results import pkg/circomcompat @@ -22,6 +24,7 @@ import ../../../contracts import ./converters export circomcompat, converters +export taskpools type CircomCompat* = object @@ -35,9 +38,24 @@ type zkeyPath: string # path to the zkey file backendCfg: ptr CircomBn254Cfg vkp*: ptr CircomKey + taskpool: Taskpool NormalizedProofInputs*[H] {.borrow: `.`.} = distinct ProofInputs[H] + ProveTask = object + circom: ptr CircomCompat + ctx: ptr CircomCompatCtx + proof: ptr Proof + success: Atomic[bool] + signal: ThreadSignalPtr + + VerifyTask = object + proof: ptr CircomProof + vkp: ptr CircomKey + inputs: ptr CircomInputs + success: ptr Atomic[bool] + signal: ThreadSignalPtr + func normalizeInput*[H]( self: CircomCompat, input: ProofInputs[H] ): NormalizedProofInputs[H] = @@ -79,7 +97,28 @@ proc release*(self: CircomCompat) = if not isNil(self.vkp): self.vkp.unsafeAddr.release_key() -proc prove[H](self: CircomCompat, input: NormalizedProofInputs[H]): ?!CircomProof = +proc circomProveTask(task: ptr ProveTask) {.gcsafe.} = + defer: + discard task[].signal.fireSync() + + var proofPtr: ptr Proof = nil + try: + if ( + let res = task.circom.backendCfg.prove_circuit(task.ctx, proofPtr.addr) + res != ERR_OK + ) or proofPtr == nil: + task.success.store(false) + return + + copyProof(task.proof, proofPtr[]) + task.success.store(true) + finally: + if proofPtr != nil: + proofPtr.addr.release_proof() + +proc asyncProve*[H]( + self: CircomCompat, input: NormalizedProofInputs[H], proof: ptr Proof +): Future[?!void] {.async.} = doAssert input.samples.len == self.numSamples, "Number of samples does not match" doAssert input.slotProof.len <= self.datasetDepth, @@ -143,13 +182,11 @@ proc prove[H](self: CircomCompat, input: NormalizedProofInputs[H]): ?!CircomProo for s in input.samples: var - merklePaths = s.merklePaths.mapIt(it.toBytes) + merklePaths = s.merklePaths.mapIt(@(it.toBytes)).concat data = s.cellData.mapIt(@(it.toBytes)).concat if ctx.push_input_u256_array( - "merklePaths".cstring, - merklePaths[0].addr, - uint (merklePaths[0].len * merklePaths.len), + "merklePaths".cstring, merklePaths[0].addr, uint (merklePaths.len) ) != ERR_OK: return failure("Failed to push merkle paths") @@ -157,44 +194,127 @@ proc prove[H](self: CircomCompat, input: NormalizedProofInputs[H]): ?!CircomProo ERR_OK: return failure("Failed to push cell data") - var proofPtr: ptr Proof = nil + without threadPtr =? ThreadSignalPtr.new(): + return failure("Unable to create thread signal") - let proof = + defer: + threadPtr.close().expect("closing once works") + + var task = ProveTask(circom: addr self, ctx: ctx, proof: proof, signal: threadPtr) + + let taskPtr = addr task + + doAssert task.circom.taskpool.numThreads > 1, + "Must have at least one separate thread or signal will never be fired" + task.circom.taskpool.spawn circomProveTask(taskPtr) + let threadFut = threadPtr.wait() + + try: + await threadFut.join() + except CatchableError as exc: try: - if (let res = self.backendCfg.prove_circuit(ctx, proofPtr.addr); res != ERR_OK) or - proofPtr == nil: - return failure("Failed to prove - err code: " & $res) - - proofPtr[] + await threadFut + except AsyncError as asyncExc: + return failure(asyncExc.msg) finally: - if proofPtr != nil: - proofPtr.addr.release_proof() + if exc of CancelledError: + raise (ref CancelledError) exc + else: + return failure(exc.msg) - success proof + if not task.success.load(): + return failure("Failed to prove circuit") -proc prove*[H](self: CircomCompat, input: ProofInputs[H]): ?!CircomProof = - self.prove(self.normalizeInput(input)) + success() + +proc prove*[H]( + self: CircomCompat, input: ProofInputs[H] +): Future[?!CircomProof] {.async, raises: [CancelledError].} = + var proof = newProof() + defer: + destroyProof(proof) + + try: + if error =? (await self.asyncProve(self.normalizeInput(input), proof)).errorOption: + return failure(error) + return success(deepCopy(proof)[]) + except CancelledError as exc: + raise exc + +proc circomVerifyTask(task: ptr VerifyTask) {.gcsafe.} = + defer: + task[].inputs[].releaseCircomInputs() + discard task[].signal.fireSync() + let res = verify_circuit(task[].proof, task[].inputs, task[].vkp) + if res == ERR_OK: + task[].success[].store(true) + elif res == ERR_FAILED_TO_VERIFY_PROOF: + task[].success[].store(false) + else: + task[].success[].store(false) + error "Failed to verify proof", errorCode = res + +proc asyncVerify*[H]( + self: CircomCompat, + proof: CircomProof, + inputs: ProofInputs[H], + success: ptr Atomic[bool], +): Future[?!void] {.async.} = + var proofPtr = unsafeAddr proof + var inputs = inputs.toCircomInputs() + + without threadPtr =? ThreadSignalPtr.new(): + return failure("Unable to create thread signal") + + defer: + threadPtr.close().expect("closing once works") + + var task = VerifyTask( + proof: proofPtr, + vkp: self.vkp, + inputs: addr inputs, + success: success, + signal: threadPtr, + ) + + let taskPtr = addr task + + doAssert self.taskpool.numThreads > 1, + "Must have at least one separate thread or signal will never be fired" + + self.taskpool.spawn circomVerifyTask(taskPtr) + + let threadFut = threadPtr.wait() + + try: + await threadFut.join() + except CatchableError as exc: + try: + await threadFut + except AsyncError as asyncExc: + return failure(asyncExc.msg) + finally: + if exc of CancelledError: + raise (ref CancelledError) exc + else: + return failure(exc.msg) + + success() proc verify*[H]( self: CircomCompat, proof: CircomProof, inputs: ProofInputs[H] -): ?!bool = +): Future[?!bool] {.async, raises: [CancelledError].} = ## Verify a proof using a ctx ## - - var - proofPtr = unsafeAddr proof - inputs = inputs.toCircomInputs() - + var res = newVerifyResult() + defer: + destroyVerifyResult(res) try: - let res = verify_circuit(proofPtr, inputs.addr, self.vkp) - if res == ERR_OK: - success true - elif res == ERR_FAILED_TO_VERIFY_PROOF: - success false - else: - failure("Failed to verify proof - err code: " & $res) - finally: - inputs.releaseCircomInputs() + if error =? (await self.asyncVerify(proof, inputs, res)).errorOption: + return failure(error) + return success(res[].load()) + except CancelledError as exc: + raise exc proc init*( _: type CircomCompat, @@ -206,6 +326,7 @@ proc init*( blkDepth = DefaultBlockDepth, cellElms = DefaultCellElms, numSamples = DefaultSamplesNum, + taskpool: Taskpool, ): CircomCompat = ## Create a new ctx ## @@ -237,4 +358,5 @@ proc init*( numSamples: numSamples, backendCfg: cfg, vkp: vkpPtr, + taskpool: taskpool, ) diff --git a/codex/slots/proofs/backends/converters.nim b/codex/slots/proofs/backends/converters.nim index ee771477..b0c197f8 100644 --- a/codex/slots/proofs/backends/converters.nim +++ b/codex/slots/proofs/backends/converters.nim @@ -10,6 +10,7 @@ {.push raises: [].} import pkg/circomcompat +import std/atomics import ../../../contracts import ../../types @@ -22,6 +23,9 @@ type CircomProof* = Proof CircomKey* = VerifyingKey CircomInputs* = Inputs + VerifyResult* = Atomic[bool] + +export VerifyResult proc toCircomInputs*(inputs: ProofInputs[Poseidon2Hash]): CircomInputs = var @@ -52,3 +56,32 @@ func toG2*(g: CircomG2): G2Point = func toGroth16Proof*(proof: CircomProof): Groth16Proof = Groth16Proof(a: proof.a.toG1, b: proof.b.toG2, c: proof.c.toG1) + +proc newProof*(): ptr Proof = + result = cast[ptr Proof](allocShared0(sizeof(Proof))) + +proc newVerifyResult*(): ptr VerifyResult = + result = cast[ptr VerifyResult](allocShared0(sizeof(VerifyResult))) + +proc destroyVerifyResult*(result: ptr VerifyResult) = + if result != nil: + deallocShared(result) + +proc destroyProof*(proof: ptr Proof) = + if proof != nil: + deallocShared(proof) + +proc copyInto*(dest: var G1, src: G1) = + copyMem(addr dest.x[0], addr src.x[0], 32) + copyMem(addr dest.y[0], addr src.y[0], 32) + +proc copyInto*(dest: var G2, src: G2) = + for i in 0 .. 1: + copyMem(addr dest.x[i][0], addr src.x[i][0], 32) + copyMem(addr dest.y[i][0], addr src.y[i][0], 32) + +proc copyProof*(dest: ptr Proof, src: Proof) = + if not isNil(dest): + copyInto(dest.a, src.a) + copyInto(dest.b, src.b) + copyInto(dest.c, src.c) diff --git a/codex/slots/proofs/backendutils.nim b/codex/slots/proofs/backendutils.nim index 0e334ace..b6995896 100644 --- a/codex/slots/proofs/backendutils.nim +++ b/codex/slots/proofs/backendutils.nim @@ -1,8 +1,13 @@ import ./backends +import pkg/taskpools type BackendUtils* = ref object of RootObj method initializeCircomBackend*( - self: BackendUtils, r1csFile: string, wasmFile: string, zKeyFile: string + self: BackendUtils, + r1csFile: string, + wasmFile: string, + zKeyFile: string, + taskpool: Taskpool, ): AnyBackend {.base, gcsafe.} = - CircomCompat.init(r1csFile, wasmFile, zKeyFile) + CircomCompat.init(r1csFile, wasmFile, zKeyFile, taskpool = taskpool) diff --git a/codex/slots/proofs/prover.nim b/codex/slots/proofs/prover.nim index 36fc0a05..ee75bbfe 100644 --- a/codex/slots/proofs/prover.nim +++ b/codex/slots/proofs/prover.nim @@ -72,7 +72,7 @@ proc prove*( return failure(err) # prove slot - without proof =? self.backend.prove(proofInput), err: + without proof =? await self.backend.prove(proofInput), err: error "Unable to prove slot", err = err.msg return failure(err) @@ -83,7 +83,11 @@ proc verify*( ): Future[?!bool] {.async.} = ## Prove a statement using backend. ## Returns a future that resolves to a proof. - self.backend.verify(proof, inputs) + without res =? (await self.backend.verify(proof, inputs)), err: + error "Unable to verify proof", err = err.msg + return failure(err) + + return success(res) proc new*( _: type Prover, store: BlockStore, backend: AnyBackend, nSamples: int