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

@@ -15,94 +15,73 @@ 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 () -> []))
       in
       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 =

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

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

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

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.