Got all recursive tests and configurable tests working

src/lib/pickles/cache.ml

@@ -186,8 +186,9 @@ module Wrap = struct
               header path ) )
 
   let read_or_generate
-      ~(custom_gate_type : (* TODO: JES: Delete this *)
-      Backend.Tick.Field.t Kimchi_types.polish_token array option )
+      ~(custom_gate_type :
+         (* TODO: JES: Delete this *)
+         Backend.Tick.Field.t Kimchi_types.polish_token array option )
       ~prev_challenges cache k_p k_v typ return_typ main =
     let module Vk = Verification_key in
     let open Impls.Wrap in
@@ -203,6 +204,7 @@ module Wrap = struct
              (Keypair.create ~pk ~vk:(Backend.Tock.Keypair.vk pk), d)
          | Error _e ->
              let r =
+               (* JES: TODO: Try removing custom_gate_type here *)
                Common.time "wrapkeygen" (fun () ->
                    constraint_system ~input_typ:typ ~return_typ main
                    |> Keypair.generate ~custom_gate_type:None ~prev_challenges )

src/lib/pickles/cache.mli

@@ -21,7 +21,8 @@ module Step : sig
   end
 
   val read_or_generate :
-       custom_gate_type:Impls.Step.Field.Constant.t Kimchi_types.polish_token array option
+       custom_gate_type:
+         Impls.Step.Field.Constant.t Kimchi_types.polish_token array option
     -> prev_challenges:int
     -> Key_cache.Spec.t list
     -> Key.Proving.t lazy_t
@@ -61,7 +62,8 @@ module Wrap : sig
   end
 
   val read_or_generate :
-       custom_gate_type:Impls.Step.Field.Constant.t Kimchi_types.polish_token array option
+       custom_gate_type:
+         Impls.Step.Field.Constant.t Kimchi_types.polish_token array option
     -> prev_challenges:Core_kernel.Int.t
     -> Key_cache.Spec.t list
     -> Key.Proving.t Core_kernel.Lazy.t

src/lib/pickles/compile.ml

@@ -420,12 +420,12 @@ struct
     in
     let full_signature = { Full_signature.padded; maxes = (module Maxes) } in
     Timer.clock __LOC__ ;
-    (* Note: we have the limitation that the custom_gate_type for all choices must be the same *)
     let feature_flags, custom_gate_type =
       let rec go :
           type a b c d.
              (a, b, c, d) H4.T(IR).t
-          -> Opt.Flag.t Plonk_types.Features.Full.t * (Backend.Tick.Field.t Kimchi_types.polish_token array option) =
+          -> Opt.Flag.t Plonk_types.Features.Full.t
+             * Backend.Tick.Field.t Kimchi_types.polish_token array option =
        fun rules ->
         match rules with
         | [] ->
@@ -440,8 +440,12 @@ struct
             , r.custom_gate_type )
         | r :: rules ->
             let feature_flags, custom_gate_type = go rules in
-            (* JES: TODO Assert below *)
-            (* assert (custom_gate_type = r.custom_gate_type) ; *)
+            (* Note: For now we only support one choice when custom gates are defined *)
+            if
+              Option.is_some custom_gate_type
+              || Option.is_some r.custom_gate_type
+            then
+              failwith "Only a single choice is allowed with configurable gates" ;
             ( Plonk_types.Features.Full.map2
                 (Plonk_types.Features.to_full ~or_:( || ) r.feature_flags)
                 feature_flags ~f:(fun enabled flag ->
@@ -682,8 +686,11 @@ struct
       in
       let r =
         Common.time "wrap read or generate " (fun () ->
-            Cache.Wrap.read_or_generate (* Due to Wrap_hack *)
-              ~custom_gate_type:None (* We don't override custom gates on the wrap proof *)
+            Cache.Wrap.read_or_generate
+            (* Due to Wrap_hack *)
+            (* JES: TODO: Try removing custom_gate_type here *)
+              ~custom_gate_type:None
+                (* We don't override custom gates on the wrap proof *)
                 (* Note: custom_gate_type needed for case when gate is overriden (true) and witness is valid *)
               ~prev_challenges:2 cache disk_key_prover disk_key_verifier typ
               (Snarky_backendless.Typ.unit ())
@@ -1136,7 +1143,8 @@ let compile_with_wrap_main_override_promise :
   in
   let custom_gate_type =
     let compiled = Types_map.lookup_compiled self.id in
-    compiled.custom_gate_type in
+    compiled.custom_gate_type
+  in
   let module P = struct
     type statement = value
 

src/lib/pickles/compile.mli

@@ -125,7 +125,8 @@ module Side_loaded : sig
        name:string
     -> max_proofs_verified:(module Nat.Add.Intf with type n = 'n1)
     -> feature_flags:Opt.Flag.t Plonk_types.Features.t
-    -> custom_gate_type:Backend.Tick.Field.t Kimchi_types.polish_token array option
+    -> custom_gate_type:
+         Backend.Tick.Field.t Kimchi_types.polish_token array option
     -> typ:('var, 'value) Impls.Step.Typ.t
     -> ('var, 'value, 'n1, Verification_key.Max_branches.n) Tag.t
 

src/lib/pickles/impls.mli

@@ -21,7 +21,8 @@ module Step : sig
     val create : pk:Proving_key.t -> vk:Verification_key.t -> t
 
     val generate :
-         custom_gate_type:Impl.Field.Constant.t Kimchi_types.polish_token array option
+         custom_gate_type:
+           Impl.Field.Constant.t Kimchi_types.polish_token array option
       -> prev_challenges:int
       -> Kimchi_pasta_constraint_system.Vesta_constraint_system.t
       -> t
@@ -110,7 +111,8 @@ module Wrap : sig
     val create : pk:Proving_key.t -> vk:Verification_key.t -> t
 
     val generate :
-         custom_gate_type:Impl.Field.Constant.t Kimchi_types.polish_token array option
+         custom_gate_type:
+           Impl.Field.Constant.t Kimchi_types.polish_token array option
       -> prev_challenges:int
       -> Kimchi_pasta_constraint_system.Pallas_constraint_system.t
       -> t

src/lib/pickles/inductive_rule.ml

@@ -116,7 +116,8 @@ type ( 'prev_vars
          'a_var main_input
       -> ('prev_vars, 'widths, 'ret_var, 'auxiliary_var) main_return
   ; feature_flags : bool Pickles_types.Plonk_types.Features.t
-  ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+  ; custom_gate_type :
+      Backend.Tick.Field.t Kimchi_types.polish_token array option
   }
 
 module T

src/lib/pickles/inductive_rule.mli

@@ -114,7 +114,8 @@ type ( 'prev_vars
          'a_var main_input
       -> ('prev_vars, 'widths, 'ret_var, 'auxiliary_var) main_return
   ; feature_flags : bool Pickles_types.Plonk_types.Features.t
-  ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+  ; custom_gate_type :
+      Backend.Tick.Field.t Kimchi_types.polish_token array option
   }
 
 module T

src/lib/pickles/pickles_intf.mli

@@ -235,7 +235,8 @@ module type S = sig
              'a_var main_input
           -> ('prev_vars, 'widths, 'ret_var, 'auxiliary_var) main_return
       ; feature_flags : bool Pickles_types.Plonk_types.Features.t
-      ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+      ; custom_gate_type :
+          Backend.Tick.Field.t Kimchi_types.polish_token array option
       }
   end
 
@@ -244,7 +245,7 @@ module type S = sig
 
   val verify_promise :
        ?chunking_data:chunking_data
-    -> ?custom_gate_type: Backend.Tick.Field.t Kimchi_types.polish_token array
+    -> ?custom_gate_type:Backend.Tick.Field.t Kimchi_types.polish_token array
     -> (module Nat.Intf with type n = 'n)
     -> (module Statement_value_intf with type t = 'a)
     -> Verification_key.t
@@ -341,7 +342,8 @@ module type S = sig
          name:string
       -> max_proofs_verified:(module Nat.Add.Intf with type n = 'n1)
       -> feature_flags:Opt.Flag.t Plonk_types.Features.t
-      -> custom_gate_type:Backend.Tick.Field.t Kimchi_types.polish_token array option
+      -> custom_gate_type:
+           Backend.Tick.Field.t Kimchi_types.polish_token array option
       -> typ:('var, 'value) Impls.Step.Typ.t
       -> ('var, 'value, 'n1, Verification_key.Max_branches.n) Tag.t
 

src/lib/pickles/step_branch_data.ml

@@ -53,7 +53,8 @@ type ( 'a_var
               and type return_value = 'ret_value
               and type auxiliary_value = 'auxiliary_value )
       ; feature_flags : bool Plonk_types.Features.t
-      ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+      ; custom_gate_type :
+          Backend.Tick.Field.t Kimchi_types.polish_token array option
       }
       -> ( 'a_var
          , 'a_value

src/lib/pickles/step_branch_data.mli

@@ -56,7 +56,8 @@ type ( 'a_var
               and type return_value = 'ret_value
               and type statement = 'a_value )
       ; feature_flags : bool Plonk_types.Features.t
-      ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+      ; custom_gate_type :
+          Backend.Tick.Field.t Kimchi_types.polish_token array option
       }
       -> ( 'a_var
          , 'a_value

src/lib/pickles/step_verifier.ml

@@ -855,8 +855,7 @@ struct
   (* TODO: This needs to handle the fact of variable length evaluations.
      Meaning it needs opt sponge. *)
   let finalize_other_proof (type b branches)
-      (module Proofs_verified : Nat.Add.Intf with type n = b)
-      ?custom_gate_type
+      (module Proofs_verified : Nat.Add.Intf with type n = b) ~custom_gate_type
       ~(step_domains :
          [ `Known of (Domains.t, branches) Vector.t | `Side_loaded ] ) ~zk_rows
       ~(* TODO: Add "actual proofs verified" so that proofs don't

src/lib/pickles/step_verifier.mli

@@ -38,7 +38,8 @@ val assert_n_bits :
 
 val finalize_other_proof :
      (module Pickles_types.Nat.Add.Intf with type n = 'b)
-  -> ?custom_gate_type:Backend.Tick.Field.t Kimchi_types.polish_token array option
+  -> custom_gate_type:
+       Backend.Tick.Field.t Kimchi_types.polish_token array option
   -> step_domains:
        [ `Known of (Import.Domains.t, 'branches) Pickles_types.Vector.t
        | `Side_loaded ]

src/lib/pickles/test/customizable_gate/test_customizable_gate.ml

@@ -4,9 +4,9 @@ open Pickles.Impls.Step
 
 let perform_step_tests = true
 
-let perform_recursive_tests = false
+let perform_recursive_tests = true
 
-let perform_step_choices_test = false
+let perform_step_choices_test = true
 
 let () = Pickles.Backend.Tick.Keypair.set_urs_info []
 
@@ -212,8 +212,7 @@ let () =
     test ~step_only:true ~custom_gate_type:None ~valid_witness:true ;
     (* Customised as Conditional gate; valid witness *)
     (* Note: Requires Cache.Wrap.read_or_generate to have custom_gate_type passed to it *)
-    test ~step_only:true ~custom_gate_type:conditional_gate (* JES: TODO: custom gate def *)
-      ~valid_witness:true ;
+    test ~step_only:true ~custom_gate_type:conditional_gate ~valid_witness:true ;
 
     (* Customised as ForeignFieldAdd gate; invalid witness *)
     let test_failed =
@@ -230,8 +229,7 @@ let () =
     let test_failed =
       try
         let _cs =
-          test ~step_only:true
-            ~custom_gate_type:conditional_gate (* JES: TODO: custom gate def *)
+          test ~step_only:true ~custom_gate_type:conditional_gate
             ~valid_witness:false
         in
         false
@@ -246,9 +244,7 @@ let () =
     test ~step_only:false ~custom_gate_type:None ~valid_witness:true ;
 
     (* Customised as Conditional gate; valid witness *)
-    test ~step_only:false
-      ~custom_gate_type:None (* JES: TODO: custom gate def *)
-      ~valid_witness:true ;
+    test ~step_only:false ~custom_gate_type:conditional_gate ~valid_witness:true ;
 
     (* Customised as ForeignFieldAdd gate; invalid witness *)
     let test_failed =
@@ -265,8 +261,7 @@ let () =
     let test_failed =
       try
         let _cs =
-          test ~step_only:false
-            ~custom_gate_type:None (* JES: TODO: custom gate def *)
+          test ~step_only:false ~custom_gate_type:conditional_gate
             ~valid_witness:false
         in
         false
@@ -298,9 +293,7 @@ let () =
                   ; prevs = []
                   ; main =
                       (fun _ ->
-                        create_customisable_circuit
-                          ~custom_gate_type:None
-                            (* JES: TODO: custom gate def *)
+                        create_customisable_circuit ~custom_gate_type:None
                           ~valid_witness:true ;
 
                         { previous_proof_statements = []
@@ -310,14 +303,14 @@ let () =
                   ; feature_flags =
                       Pickles_types.Plonk_types.Features.
                         { none_bool with foreign_field_add = true }
-                  ; custom_gate_type = None (* JES: TODO: custom gate def *)
+                  ; custom_gate_type = None
                   }
                 ; { identifier = "customizable gate (choice 2)"
                   ; prevs = []
                   ; main =
                       (fun _ ->
-                        create_customisable_circuit ~custom_gate_type:None
-                          ~valid_witness:true ;
+                        create_customisable_circuit
+                          ~custom_gate_type:conditional_gate ~valid_witness:true ;
 
                         { previous_proof_statements = []
                         ; public_output = ()
@@ -326,7 +319,7 @@ let () =
                   ; feature_flags =
                       Pickles_types.Plonk_types.Features.
                         { none_bool with foreign_field_add = true }
-                  ; custom_gate_type = None
+                  ; custom_gate_type = conditional_gate
                   }
                 ] )
               ()

src/lib/pickles/test/customizable_optional_gate/customizable_optional_gate.ml

@@ -134,14 +134,44 @@ let register_feature_test (name, specific_feature_flags, custom_gate_type) =
     specific_feature_flags Plonk_types.Features.none_bool custom_gate_type
     custom_gate_type
 
+(* User-supplied conditional gate in RPN
+ *     w(0) = w(1) * w(3) + (1 - w(3)) * w(2)
+ *)
+let conditional_gate =
+  Some
+    Kimchi_types.
+      [| Cell { col = Index ForeignFieldAdd; row = Curr }
+       ; Cell { col = Witness 3; row = Curr }
+       ; Dup
+       ; Mul
+       ; Cell { col = Witness 3; row = Curr }
+       ; Sub
+       ; Alpha
+       ; Pow 1l
+       ; Cell { col = Witness 0; row = Curr }
+       ; Cell { col = Witness 3; row = Curr }
+       ; Cell { col = Witness 1; row = Curr }
+       ; Mul
+       ; Literal (Impl.Field.Constant.of_int 1)
+       ; Cell { col = Witness 3; row = Curr }
+       ; Sub
+       ; Cell { col = Witness 2; row = Curr }
+       ; Mul
+       ; Add
+       ; Sub
+       ; Mul
+       ; Add
+       ; Mul
+      |]
+
 let () =
   let configurations =
     [ ( "customizable gate (ffadd)"
       , Plonk_types.Features.{ none_bool with foreign_field_add = true }
       , None )
     ; ( "customizable gate (conditional)"
       , Plonk_types.Features.{ none_bool with foreign_field_add = true }
-      , None (* JES: TODO: set to gate definition *) )
+      , conditional_gate )
     ]
   in
   List.iter ~f:register_feature_test configurations ;

src/lib/pickles/types_map.ml

@@ -19,7 +19,8 @@ module Basic = struct
     ; wrap_key : Tick.Inner_curve.Affine.t array Plonk_verification_key_evals.t
     ; wrap_vk : Impls.Wrap.Verification_key.t
     ; feature_flags : Opt.Flag.t Plonk_types.Features.Full.t
-    ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+    ; custom_gate_type :
+        Backend.Tick.Field.t Kimchi_types.polish_token array option
     ; num_chunks : int
     ; zk_rows : int
     }
@@ -42,7 +43,8 @@ module Side_loaded = struct
       { max_proofs_verified : (module Nat.Add.Intf with type n = 'n1)
       ; public_input : ('var, 'value) Impls.Step.Typ.t
       ; feature_flags : Opt.Flag.t Plonk_types.Features.Full.t
-      ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+      ; custom_gate_type :
+          Backend.Tick.Field.t Kimchi_types.polish_token array option
       ; branches : 'n2 Nat.t
       ; num_chunks : int
       ; zk_rows : int
@@ -102,7 +104,8 @@ module Compiled = struct
     ; wrap_domains : Domains.t
     ; step_domains : (Domains.t, 'branches) Vector.t
     ; feature_flags : Opt.Flag.t Plonk_types.Features.Full.t
-    ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+    ; custom_gate_type :
+        Backend.Tick.Field.t Kimchi_types.polish_token array option
     ; num_chunks : int
     ; zk_rows : int
     }
@@ -123,7 +126,8 @@ module Compiled = struct
     ; wrap_domains : Domains.t
     ; step_domains : (Domains.t, 'branches) Vector.t
     ; feature_flags : Opt.Flag.t Plonk_types.Features.Full.t
-    ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+    ; custom_gate_type :
+        Backend.Tick.Field.t Kimchi_types.polish_token array option
     ; num_chunks : int
     ; zk_rows : int
     }
@@ -173,7 +177,8 @@ module For_step = struct
           Impls.Step.field Pickles_base.Proofs_verified.One_hot.Checked.t ]
     ; step_domains : [ `Known of (Domains.t, 'branches) Vector.t | `Side_loaded ]
     ; feature_flags : Opt.Flag.t Plonk_types.Features.Full.t
-    ; custom_gate_type : Backend.Tick.Field.t Kimchi_types.polish_token array option
+    ; custom_gate_type :
+        Backend.Tick.Field.t Kimchi_types.polish_token array option
     ; num_chunks : int
     ; zk_rows : int
     }