From d9c4828ca18269f65563d048c8627d49323ff212 Mon Sep 17 00:00:00 2001 From: mrmr1993 Date: Sun, 11 Feb 2024 07:29:29 +0000 Subject: [PATCH] Add callback versions of constraint_system and generate_witness --- src/base/snark0.ml | 90 ++++++++++++++++++++++++++++++++++++++++++ src/base/snark_intf.ml | 19 +++++++++ 2 files changed, 109 insertions(+) diff --git a/src/base/snark0.ml b/src/base/snark0.ml index 04a296cf4..be2853840 100644 --- a/src/base/snark0.ml +++ b/src/base/snark0.ml @@ -1294,6 +1294,53 @@ module Run = struct let x = inject_wrapper x ~f:(fun x () -> mark_active ~f:x) in Perform.constraint_system ~run:as_stateful ~input_typ ~return_typ x ) + type ('input_var, 'return_var, 'result) manual_callbacks = + { run_circuit : 'a. ('input_var -> unit -> 'a) -> 'a + ; finish_computation : 'return_var -> 'result + } + + let constraint_system_manual ~input_typ ~return_typ = + let builder = + Run.Constraint_system_builder.build ~input_typ ~return_typ + in + (* FIXME: This behaves badly with exceptions. *) + let cached_state = ref None in + let cached_active_counters = ref None in + let run_circuit circuit = + (* Check the status. *) + if + Option.is_some !cached_state || Option.is_some !cached_active_counters + then failwith "Already generating constraint system" ; + (* Partial [finalize_is_running]. *) + cached_state := Some !state ; + builder.run_computation (fun input state' -> + (* Partial [as_stateful]. *) + state := state' ; + (* Partial [mark_active]. *) + let counters = !active_counters in + cached_active_counters := Some counters ; + active_counters := this_functor_id :: counters ; + (* Start the circuit. *) + circuit input () ) + in + let finish_computation return_var = + (* Check the status. *) + if + Option.is_none !cached_state || Option.is_none !cached_active_counters + then failwith "Constraint system not in a finalizable state" ; + (* Partial [mark_active]. *) + active_counters := Option.value_exn !cached_active_counters ; + (* Create an invalid state, to avoid re-runs. *) + cached_active_counters := None ; + (* Partial [as_stateful]. *) + let state' = !state in + let res = builder.finish_computation (state', return_var) in + (* Partial [finalize_is_running]. *) + state := Option.value_exn !cached_state ; + res + in + { run_circuit; finish_computation } + let generate_public_input t x : As_prover.Vector.t = finalize_is_running (fun () -> generate_public_input t x) @@ -1310,6 +1357,49 @@ module Run = struct Perform.generate_witness_conv ~run:as_stateful ~f ~input_typ ~return_typ x input ) + let generate_witness_manual ?handlers ~input_typ ~return_typ input = + let builder = + Run.Witness_builder.auxiliary_input ?handlers ~input_typ ~return_typ + input + in + (* FIXME: This behaves badly with exceptions. *) + let cached_state = ref None in + let cached_active_counters = ref None in + let run_circuit circuit = + (* Check the status. *) + if + Option.is_some !cached_state || Option.is_some !cached_active_counters + then failwith "Already generating constraint system" ; + (* Partial [finalize_is_running]. *) + cached_state := Some !state ; + builder.run_computation (fun input state' -> + (* Partial [as_stateful]. *) + state := state' ; + (* Partial [mark_active]. *) + let counters = !active_counters in + cached_active_counters := Some counters ; + active_counters := this_functor_id :: counters ; + (* Start the circuit. *) + circuit input () ) + in + let finish_computation return_var = + (* Check the status. *) + if + Option.is_none !cached_state || Option.is_none !cached_active_counters + then failwith "Constraint system not in a finalizable state" ; + (* Partial [mark_active]. *) + active_counters := Option.value_exn !cached_active_counters ; + (* Create an invalid state, to avoid re-runs. *) + cached_active_counters := None ; + (* Partial [as_stateful]. *) + let state' = !state in + let res = builder.finish_witness_generation (state', return_var) in + (* Partial [finalize_is_running]. *) + state := Option.value_exn !cached_state ; + res + in + { run_circuit; finish_computation } + let run_unchecked x = finalize_is_running (fun () -> Perform.run_unchecked ~run:as_stateful (fun () -> mark_active ~f:x) ) diff --git a/src/base/snark_intf.ml b/src/base/snark_intf.ml index 7db66f7fb..877296038 100644 --- a/src/base/snark_intf.ml +++ b/src/base/snark_intf.ml @@ -1380,6 +1380,25 @@ module type Run_basic = sig -> 'input_value -> Proof_inputs.t + type ('input_var, 'return_var, 'result) manual_callbacks = + { run_circuit : 'a. ('input_var -> unit -> 'a) -> 'a + ; finish_computation : 'return_var -> 'result + } + + (** Callback version of [constraint_system]. *) + val constraint_system_manual : + input_typ:('input_var, 'input_value) Typ.t + -> return_typ:('return_var, 'return_value) Typ.t + -> ('input_var, 'return_var, R1CS_constraint_system.t) manual_callbacks + + (** Callback version of [generate_witness]. *) + val generate_witness_manual : + ?handlers:(request -> response) list + -> input_typ:('input_var, 'input_value) Typ.t + -> return_typ:('return_var, 'return_value) Typ.t + -> 'input_value + -> ('input_var, 'return_var, Proof_inputs.t * 'return_value) manual_callbacks + (** Generate the public input vector for a given statement. *) val generate_public_input : ('input_var, 'input_value) Typ.t -> 'input_value -> Field.Constant.Vector.t