From 3dda630f23dd9e82a247e4fdcf24c1a0c9c8a787 Mon Sep 17 00:00:00 2001 From: Las Safin Date: Tue, 18 Jun 2024 10:23:47 +0000 Subject: [PATCH] wip short-circuiting proof composition Doesn't compile, unfinished. The idea is to make an abstraction over "statements", such that they are either done by a recursive proof verification or by doing some minor logic in-circuit. For the action state extension statement, it would be short-circuiting if source and target are only N actions from each other (for e.g. N = 8), but more than that could use a recursive proof (and reusing the portion of the circuit to verify the last N ones). Then, we can define short-circuiting proof composition, currently called WrapTwo dumbly, such that if either of the two wrapped statements can be short-circuited, then we avoid an extra proof. This is very useful for reducing proofs to deal with the 2 recursive proof verifications per circuit limit, since you can stack WrapTwos on top of WrapTwos, and only proving when necessary, short-circuiting the rest. Closes #162 --- src/app/zeko/action_state_extension.ml | 140 ++++++++------------ src/app/zeko/action_state_extension.mli | 2 + src/app/zeko/dune | 2 +- src/app/zeko/wrap_two.ml | 165 ++++++++++++++++++++++++ src/app/zeko/wrapper.ml | 2 +- src/app/zeko/zeko_util.ml | 45 ++++++- 6 files changed, 265 insertions(+), 91 deletions(-) create mode 100644 src/app/zeko/wrap_two.ml diff --git a/src/app/zeko/action_state_extension.ml b/src/app/zeko/action_state_extension.ml index 0a81ebee7b..6b186136c2 100644 --- a/src/app/zeko/action_state_extension.ml +++ b/src/app/zeko/action_state_extension.ml @@ -15,65 +15,31 @@ module Stmt = struct type t = { source : F.t; target : F.t } [@@deriving show, snarky] end -module T = struct - type t = { source : F.t; target : F.t; proof : RefProof.t } - [@@deriving snarky] - - (* Do some checking here to avoid having two create unnecessary proofs, see issue #101 *) - let verify ?check ({ source; target; proof } : var) : - (Stmt.var, Nat.N2.n) Pickles.Inductive_rule.Previous_proof_statement.t = - let proof_must_verify = Boolean.not (Field.equal source target) in - let proof_must_verify = - match check with - | Some check -> - Boolean.(check &&& proof_must_verify) - | None -> - proof_must_verify - in - { public_input = { source; target } - ; proof_must_verify (* Don't check proof if source == target *) - ; proof - } - - let statement_var ({ source; target } : var) : Stmt.var = { source; target } - - let statement ({ source; target } : t) : Stmt.t = { source; target } -end - -include T - -(** Step `n` times *) -module MkStep (N : sig - val n : int -end) = -struct - module Actionss' = struct - module T = F - (* This should be treated as an instance of Actions, - however, Outside_hash_image.t is reserved, and is ignored *) - - let length = N.n - end - - module Actionss = SnarkList (Actionss') - +module Mk(Count : sig val n : int end) = struct module Witness = struct - type t = { actionss : Actionss.t; prev : T.t } [@@deriving snarky] + type t = Proof of Proof.t | List of F.t list end + module RefWitness = MkRef (Witness) + + type t = { source : F.t; target: F.t; witness : RefWitness.t } [@@deriving snarky] - include MkHandler (Witness) - - (* FIXME: slow, but probably fast enough for now *) - let construct_witness ~(actionss : Actions.t list) = - List.append - (List.map ~f:Actions.hash actionss) - (List.init - (N.n - List.length actionss) - ~f:(fun _ -> Outside_hash_image.t) ) + let dummy_proof = Pickles.Proof.dummy Nat.N2.n Nat.N2.n Nat.N1.n ~domain_log2:14 - let%snarkydef_ main Pickles.Inductive_rule.{ public_input = () } = - let Witness.{ actionss; prev = { source } as prev } = exists_witness () in - let f actions target' = + (* Do some checking here to avoid having two create unnecessary proofs, see issue #101 *) + (* Creates a constraint between source and target such that if there is some actionss in between + to link them, we return false (proof needn't be checked), but return true *) + let proof_must_verify_impl (witness : RefWitness.var) (t : Stmt.var) : Boolean.var = + let module FList' = struct + module T = F + let length = Count.n + end in + let module FList = SnarkList (FList') in + let actionss = exists FList.typ ~compute:(fun () -> + match As_prover.Ref.get witness with + | List actionss -> List.append actionss (List.init (Count.n - List.length actionss) ~f:(fun _ -> Outside_hash_image.t)) + | Proof _ -> List.init Count.n ~f:(fun _ -> Outside_hash_image.t) + ) in + let f actions target = (* Bad unsafe use, with mismatching data and hash, but it works *) let actions' = Data_as_hash.make_unsafe actions (As_prover.Ref.create (fun () -> [])) @@ -81,28 +47,41 @@ struct let open Field in if_ (equal actions @@ constant Outside_hash_image.t) - ~then_:target' - ~else_:(Actions.push_events_checked target' actions') + ~then_:target + ~else_:(Actions.push_events_checked target actions') in - let target' = List.fold_right ~init:source ~f actionss in - Pickles.Inductive_rule. - { previous_proof_statements = [ verify prev ] - ; public_output = Stmt.{ source; target = target' } - ; auxiliary_output = () - } + let target' = List.fold_right ~init:t.source ~f actionss in + Boolean.not (Field.equal target' t.target) - let rule tag : _ Pickles.Inductive_rule.t = - { identifier = "action state extension step" - ; prevs = [ tag ] - ; main - ; feature_flags = Pickles_types.Plonk_types.Features.none_bool + + (* Do some checking here to avoid having two create unnecessary proofs, see issue #101 *) + let verify ?check ({ source; target; witness } : var) : + (Stmt.var, Nat.N2.n) Pickles.Inductive_rule.Previous_proof_statement.t = + let proof_must_verify = proof_must_verify_impl witness {source; target} in + { public_input = ({ source; target } : Stmt.var) + ; proof_must_verify = (match check with + | Some check -> Boolean.(check &&& proof_must_verify) + | None -> proof_must_verify) + ; proof = As_prover.Ref.create (fun () -> + (match As_prover.Ref.get witness with + | Proof p -> p + | List _ -> dummy_proof + ) + ) } + + let statement_var ({ source; target } : var) : Stmt.var = { source; target } + + let statement ({ source; target } : t) : Stmt.t = { source; target } end +module P0 = Mk(struct let n = 0 end) +module P256 = Mk(struct let n = 256 end) + (** Merge two matching statements *) -module Merge = struct +module MkMerge (Left : module type of P256) (Right : module type of P256) = struct module Witness = struct - type t = { left : T.t; right : T.t } [@@deriving snarky] + type t = { left : Left.t; right : Right.t } [@@deriving snarky] end include MkHandler (Witness) @@ -111,7 +90,7 @@ module Merge = struct let Witness.{ left; right } = exists_witness () in Field.Assert.equal left.target right.source ; Pickles.Inductive_rule. - { previous_proof_statements = [ verify left; verify right ] + { previous_proof_statements = [ Left.verify left; Right.verify right ] ; public_output = Stmt.{ source = left.source; target = right.target } ; auxiliary_output = () } @@ -123,12 +102,9 @@ module Merge = struct ; feature_flags = Pickles_types.Plonk_types.Features.none_bool } end +module MergeLeft = MkMerge (P256) (P0) +module MergeRight = MkMerge (P0) (P256) -module N_2_8 = struct - let n = Int.pow 2 8 -end - -module Step = MkStep (N_2_8) open Async_kernel let compilation_result = @@ -144,13 +120,13 @@ let compilation_result = ~constraint_constants: (Genesis_constants.Constraint_constants.to_snark_keys_header constraint_constants ) - ~choices:(fun ~self -> [ Step.rule self; Merge.rule self ]) ) ) + ~choices:(fun ~self -> [ MergeLeft.rule self; MergeRight.rule self ]) ) ) let tag = lazy (match force compilation_result with tag, _, _, _ -> tag) -let step w = +let merge_left_ w = match force compilation_result with - | _, _, _, Pickles.Provers.[ step_; _ ] -> + | _, _, _, Pickles.Provers.[ merge_left_; _ ] -> time_async "Action_state_extension.step" (fun () -> step_ ~handler:(Step.handler w) () ) @@ -162,12 +138,6 @@ let merge (left : t) (right : t) : t Deferred.t = in ({ source = stmt.source; target = stmt.target; proof } : t) -let dummy_proof source : t = - { source - ; target = source - ; proof = Pickles.Proof.dummy Nat.N2.n Nat.N2.n Nat.N1.n ~domain_log2:14 - } - (* Head of list should be oldest, tail should be newest *) let prove ?(dummy = false) ~(source : F.t) (actionss : Actions.t list) : t Deferred.t = diff --git a/src/app/zeko/action_state_extension.mli b/src/app/zeko/action_state_extension.mli index bd92e46d12..674f33d8c0 100644 --- a/src/app/zeko/action_state_extension.mli +++ b/src/app/zeko/action_state_extension.mli @@ -11,6 +11,8 @@ val statement_var : var -> Stmt.var val statement : t -> Stmt.t +val proof_must_verify : Stmt.var -> Boolean.var + val verify : ?check:Boolean.var -> var diff --git a/src/app/zeko/dune b/src/app/zeko/dune index 970eeac8ec..67fc19a7ee 100644 --- a/src/app/zeko/dune +++ b/src/app/zeko/dune @@ -68,7 +68,7 @@ (backend bisect_ppx)) (preprocess (pps ppx_deriving.show ppx_deriving_snarky ppx_snarky ppx_mina ppx_version ppx_jane ppx_compare h_list.ppx)) - (modules zeko_util) + (modules zeko_util wrap_two) ) (library diff --git a/src/app/zeko/wrap_two.ml b/src/app/zeko/wrap_two.ml new file mode 100644 index 0000000000..e800c07ac5 --- /dev/null +++ b/src/app/zeko/wrap_two.ml @@ -0,0 +1,165 @@ +open Core_kernel +open Mina_base +open Snark_params.Tick.Run +(* Impure interface to snarky, FIXME: replace with pure one *) + +open Zkapp_basic +open Account_update +open Zeko_util + +module Opaque (X : CircuitType) = struct + type t = F.t array + type var = F.var array + let typ = + let Typ typ' = X.typ in + Typ.array ~length:typ'.size_in_field_elements F.typ +end + +module type Verifiable = sig + module Stmt : sig + type t + type var + val typ : (var, t) Typ.t + end + module Witness : sig + type t + val contains_proof : t -> bool + end + val verify : check:Boolean.var option -> Stmt.var -> Witness.t option V.t -> + ( Stmt.var + , Pickles_types.Nat.N2.n ) + Pickles.Inductive_rule.Previous_proof_statement.t + type tag_var + type tag_value + val tag : + unit -> ( + tag_var + , tag_value + , Pickles_types.Nat.N2.n + , Pickles_types.Nat.N2.n ) + Pickles.Tag.t +end + +let typ_size (Typ typ : ('var, 't) Typ.t) = typ.size_in_field_elements + +(** NB: X and Y must have same size in # field elements *) +module WrapTwo (X : Verifiable) (Y : Verifiable) = struct + let () = assert (Int.(typ_size X.Stmt.typ = typ_size Y.Stmt.typ)) + let stmt_size = typ_size X.Stmt.typ + + module Stmt = struct + type t = F.t array + type var = F.var array + let typ = Typ.array ~length:stmt_size F.typ + let combine : X.Stmt.t -> Y.Stmt.t -> t + = fun x y -> + let (Typ x_typ) = X.Stmt.typ in + let (Typ y_typ) = Y.Stmt.typ in + let x_fields, _ = x_typ.value_to_fields x in + let y_fields, _ = y_typ.value_to_fields y in + Array.zip_exn x_fields y_fields + |> Array.map ~f:(fun (x, y) -> Random_oracle.hash [|x ; y|]) + let combine_var : X.Stmt.var -> Y.Stmt.var -> var + = fun x y -> + let (Typ x_typ) = X.Stmt.typ in + let (Typ y_typ) = Y.Stmt.typ in + let x_fields, _ = x_typ.var_to_fields x in + let y_fields, _ = y_typ.var_to_fields y in + Array.zip_exn x_fields y_fields + |> Array.map ~f:(fun (x, y) -> Random_oracle.Checked.hash [|x ; y|]) + end + + module Witness = struct + (* The witness for the above statement/relation. + In Proof case, we have generated a proof that succinctly describes + the above statement. + In Short_circuit case, we skip generating a proof, + which is only possible if either X or Y don't need a proof. + If neither needn't a proof, we don't do need to verify any proof at all. + The core goal is to always (recursively) verify at most one proof. + *) + type t = Proof of Proof.t | Short_circuit of X.Witness.t * Y.Witness.t + end + module RefWitness = MkRef(Witness) + + type t = { x : X.Stmt.t; y : Y.Stmt.t; witness : RefWitness.t } + [@@deriving snarky] + + let contains_proof = function + | Witness.Proof _ -> true + | Witness.Short_circuit (x, y) -> assert (not @@ X.Witness.contains_proof x && Y.Witness.contains_proof y) ; X.Witness.contains_proof x || Y.Witness.contains_proof y + + let verify ~check ({ x; y; witness } : var) : + (Stmt.var, Pickles_types.Nat.N2.n) Pickles.Inductive_rule.Previous_proof_statement.t = + let public_input = exists Stmt.typ in + let proof_must_verify = exists Boolean.typ in + let verify_x_or_y = exists Boolean.typ in + let check_x = exists Boolean.typ in + let verify_x = X.verify x ~check:check_x (V.map witness ~f:(function None -> None | Some x -> )) in + Boolean.Assert.is_true @@ proof_must_verify ; + let proof = ref None in + { public_input = public_input + ; proof_must_verify (* Don't check proof if source == target *) + ; proof + } + + open Async_kernel + + let compilation_result = + lazy + (time "Action_state_extension.compile" (fun () -> + Pickles.compile () + ~override_wrap_domain:Pickles_base.Proofs_verified.N1 + ~cache:Cache_dir.cache ~public_input:(Output Stmt.typ) + ~auxiliary_typ:Typ.unit + ~branches:(module Nat.N2) + ~max_proofs_verified:(module Nat.N2) + ~name:"action state extension" + ~constraint_constants: + (Genesis_constants.Constraint_constants.to_snark_keys_header + constraint_constants ) + ~choices:(fun ~self -> [ Step.rule self; Merge.rule self ]) ) ) + + let tag = lazy (match force compilation_result with tag, _, _, _ -> tag) + + let step w = + match force compilation_result with + | _, _, _, Pickles.Provers.[ step_; _ ] -> + time_async "Action_state_extension.step" (fun () -> + step_ ~handler:(Step.handler w) () ) + + let merge (left : t) (right : t) : t Deferred.t = + match force compilation_result with + | _, _, _, Pickles.Provers.[ _; merge_ ] -> + let%map stmt, (), proof = + merge_ ~handler:(Merge.handler { left; right }) () + in + ({ source = stmt.source; target = stmt.target; proof } : t) + + let dummy_proof source : t = + { source + ; target = source + ; proof = Pickles.Proof.dummy Nat.N2.n Nat.N2.n Nat.N1.n ~domain_log2:14 + } + + (* Head of list should be oldest, tail should be newest *) + let prove ?(dummy = false) ~(source : F.t) (actionss : Actions.t list) : + t Deferred.t = + assert (List.length actionss <= N_2_8.n) (* FIXME: support, see FIXME below *) ; + if List.is_empty actionss then return (dummy_proof source) + else if dummy then + let target = + List.fold_right ~f:(Fun.flip Actions.push_events) ~init:source actionss + in + return { (dummy_proof source) with target } + else + let%map stmt, (), proof = + (* FIXME: split up when too big, see FIXME above *) + step + Step.Witness. + { actionss = Step.construct_witness ~actionss + ; prev = dummy_proof source + } + in + ({ source; target = stmt.target; proof } : t) +end diff --git a/src/app/zeko/wrapper.ml b/src/app/zeko/wrapper.ml index a91328f068..4c41624bd0 100644 --- a/src/app/zeko/wrapper.ml +++ b/src/app/zeko/wrapper.ml @@ -26,7 +26,7 @@ end module T = struct (** Akin to Transaction_snark.t *) - type t = { stmt : Stmt.t; proof : RefProof.t } [@@deriving snarky, yojson] + type t = { stmt : Stmt.t; proof : RefProof.t [@to_yojson Proof.to_yojson] [@of_yojson Proof.of_yojson] } [@@deriving snarky, yojson] let verify { stmt; proof } = Pickles.Inductive_rule.Previous_proof_statement. diff --git a/src/app/zeko/zeko_util.ml b/src/app/zeko/zeko_util.ml index 39a5e21ae7..42f3f88e8f 100644 --- a/src/app/zeko/zeko_util.ml +++ b/src/app/zeko/zeko_util.ml @@ -5,6 +5,7 @@ open Snark_params.Tick.Run open Zkapp_basic open Account_update + module V = Prover_value let attach_control_var : @@ -201,16 +202,22 @@ module F = struct fun fmt f -> Format.pp_print_string fmt @@ Field.Constant.to_string f end +module type CircuitType = sig + type t + type var + val typ : (var, t) Typ.t +end + (** To be used with deriving snarky, a reference to T with no in-circuit representation *) module MkRef (T : sig - type t [@@deriving yojson] + type t end) = struct - type t = T.t [@@deriving yojson] + type t = T.t - type var = T.t As_prover.Ref.t + type var = T.t V.t - let typ : (var, t) Typ.t = Typ.Internal.ref () + let typ : (var, t) Typ.t = V.typ () end (** A list of `length` `t`s *) @@ -310,3 +317,33 @@ let assert_var label expr = let assert_var_checked label expr = with_label label (fun () -> expr () |> run |> Boolean.Assert.is_true) + +(* TODO: I implemented this, but should we use this? +module V = struct + type 'a t = 'a option + let typ : ('a t, 'a) Typ.t = + Typ { + var_to_fields = (fun x -> ([||], x)) + ; var_of_fields = (fun (_, x) -> x) + ; value_to_fields = (fun x -> ([||], Some x)) + ; value_of_fields = (fun (_, x) -> Option.value_exn x) + ; size_in_field_elements = 0 + ; constraint_system_auxiliary = (fun () -> None) + ; check = (fun _ -> Snark_params.Tick.Checked.return ()) + } + let create (f : unit -> 'a) : 'a t = + if As_prover.in_prover_block () + then f () |> Some + else None + let value_exn : 'a t -> 'a = function + | Some x -> x + | None -> failwith "not proving" + let return : 'a -> 'a t = fun x -> create (fun () -> x) + let map : 'a t -> f:('a -> 'b) -> 'b t = fun x ~f = match x with + | Some x -> Some (f x) + | None -> None + let bind : 'a t -> f:('a -> 'b t) -> 'b t = fun x ~f = match x with + | Some x -> f x + | None -> None +end +*)