Clean up code related to customisable pickles

src/lib/pickles/cache.ml

@@ -185,11 +185,7 @@ module Wrap = struct
                 (Kimchi_bindings.Protocol.Index.Fq.write (Some true) t.index)
               header path ) )
 
-  let read_or_generate
-      ~(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 read_or_generate ~prev_challenges cache k_p k_v typ return_typ main =
     let module Vk = Verification_key in
     let open Impls.Wrap in
     let s_p = storable in
@@ -204,10 +200,9 @@ 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 )
+                   |> Keypair.generate ~prev_challenges )
              in
              ignore
                ( Key_cache.Sync.write cache s_p k (Keypair.pk r)

src/lib/pickles/cache.mli

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

src/lib/pickles/compile.ml

@@ -686,12 +686,7 @@ struct
       in
       let r =
         Common.time "wrap read or generate " (fun () ->
-            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 *)
+            Cache.Wrap.read_or_generate (* Due to Wrap_hack *)
               ~prev_challenges:2 cache disk_key_prover disk_key_verifier typ
               (Snarky_backendless.Typ.unit ())
               main )

src/lib/pickles/impls.ml

@@ -165,13 +165,9 @@ module Wrap = struct
 
     let create = Fields.create
 
-    let generate ~custom_gate_type ~prev_challenges cs =
+    let generate ~prev_challenges cs =
       let open Tock.Keypair in
-      let keypair =
-        if Option.is_some custom_gate_type then
-          create_plus custom_gate_type ~prev_challenges cs
-        else create ~prev_challenges cs
-      in
+      let keypair = create ~prev_challenges cs in
       { pk = pk keypair; vk = vk keypair }
   end
 

src/lib/pickles/impls.mli

@@ -111,9 +111,7 @@ 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
-      -> prev_challenges:int
+         prev_challenges:int
       -> Kimchi_pasta_constraint_system.Pallas_constraint_system.t
       -> t
   end

src/lib/pickles/pickles.ml

@@ -1282,11 +1282,8 @@ module Make_str (_ : Wire_types.Concrete) = struct
             in
             let r =
               Common.time "wrap read or generate " (fun () ->
-                  Cache.Wrap.read_or_generate
-                    ~custom_gate_type
-                      (* Note: custom_gate_type not needed here, but needed when called from compile.ml *)
-                    ~prev_challenges:2 [] disk_key_prover disk_key_verifier typ
-                    Typ.unit main )
+                  Cache.Wrap.read_or_generate ~prev_challenges:2 []
+                    disk_key_prover disk_key_verifier typ Typ.unit main )
             in
             (r, disk_key_verifier)
           in

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

@@ -63,12 +63,11 @@ let constraint_constants =
 let main_body ~(feature_flags : _ Plonk_types.Features.t) () =
   if feature_flags.foreign_field_add then main_foreign_field_add ()
 
-let register_test name feature_flags1 feature_flags2 custom_gate_type1
-    custom_gate_type2 =
-  let _tag, _cache_handle, proof, Pickles.Provers.[ prove1; prove2 ] =
+let register_test name feature_flags custom_gate_type =
+  let _tag, _cache_handle, proof, Pickles.Provers.[ prove ] =
     Pickles.compile ~public_input:(Pickles.Inductive_rule.Input Typ.unit)
       ~auxiliary_typ:Typ.unit
-      ~branches:(module Nat.N2)
+      ~branches:(module Nat.N1)
       ~max_proofs_verified:(module Nat.N0)
       ~name:"optional_custom_gates"
       ~constraint_constants (* TODO(mrmr1993): This was misguided.. Delete. *)
@@ -77,62 +76,37 @@ let register_test name feature_flags1 feature_flags2 custom_gate_type1
           ; prevs = []
           ; main =
               (fun _ ->
-                main_body ~feature_flags:feature_flags1 () ;
+                main_body ~feature_flags () ;
                 { previous_proof_statements = []
                 ; public_output = ()
                 ; auxiliary_output = ()
                 } )
-          ; feature_flags = feature_flags1
-          ; custom_gate_type = custom_gate_type1
-          }
-        ; { identifier = "main2"
-          ; prevs = []
-          ; main =
-              (fun _ ->
-                main_body ~feature_flags:feature_flags2 () ;
-                { previous_proof_statements = []
-                ; public_output = ()
-                ; auxiliary_output = ()
-                } )
-          ; feature_flags = feature_flags2
-          ; custom_gate_type = custom_gate_type2
+          ; feature_flags
+          ; custom_gate_type
           }
         ] )
       ()
   in
   let module Proof = (val proof) in
-  let test_prove1 () =
+  let test_prove () =
     let public_input1, (), proof1 =
-      Async.Thread_safe.block_on_async_exn (fun () -> prove1 ())
+      Async.Thread_safe.block_on_async_exn (fun () -> prove ())
     in
     Or_error.ok_exn
       (Async.Thread_safe.block_on_async_exn (fun () ->
            Proof.verify [ (public_input1, proof1) ] ) )
   in
-  let test_prove2 () =
-    let public_input2, (), proof2 =
-      Async.Thread_safe.block_on_async_exn (fun () -> prove2 ())
-    in
-    Or_error.ok_exn
-      (Async.Thread_safe.block_on_async_exn (fun () ->
-           Proof.verify [ (public_input2, proof2) ] ) )
-  in
 
   let open Alcotest in
-  add_tests name
-    [ test_case "prove 1" `Quick test_prove1
-    ; test_case "prove 2" `Quick test_prove2
-    ]
+  add_tests name [ test_case "prove" `Quick test_prove ]
 
 let register_feature_test (name, specific_feature_flags, custom_gate_type) =
   (* Tests activating "on" logic *)
-  register_test name specific_feature_flags specific_feature_flags
-    custom_gate_type custom_gate_type ;
+  register_test name specific_feature_flags custom_gate_type ;
   (* Tests activating "maybe on" logic *)
   register_test
     (Printf.sprintf "%s (maybe)" name)
-    specific_feature_flags Plonk_types.Features.none_bool custom_gate_type
-    custom_gate_type
+    Plonk_types.Features.none_bool custom_gate_type
 
 (* User-supplied conditional gate in RPN
  *     w(0) = w(1) * w(3) + (1 - w(3)) * w(2)