Customisable pickles polish_token changes - pickles

src/lib/crypto/kimchi_bindings/stubs/dune

@@ -98,7 +98,7 @@
   (backend bisect_ppx))
  (inline_tests (flags -verbose -show-counts))
  (preprocess
-  (pps ppx_version ppx_inline_test)))
+  (pps ppx_version ppx_inline_test )))
 
 (library
  (public_name pasta_bindings)

src/lib/crypto/kimchi_bindings/stubs/src/linearization.rs

@@ -1,15 +1,17 @@
 use kimchi::{
     circuits::{
+        gate::GateType,
         constraints::FeatureFlags,
-        expr::Linearization,
+        expr::{Linearization, PolishToken},
         lookup::lookups::{LookupFeatures, LookupPatterns},
     },
     linearization::{constraints_expr, linearization_columns},
 };
 
 /// Converts the linearization of the kimchi circuit polynomial into a printable string.
 pub fn linearization_strings<F: ark_ff::PrimeField + ark_ff::SquareRootField>(
-    omit_custom_gate: bool,
+    // omit_custom_gate: bool,
+    custom_gate_type: Option<&Vec<PolishToken<F>>>,
     uses_custom_gates: bool,
 ) -> (String, Vec<(String, String)>) {
     let features = if uses_custom_gates {
@@ -36,7 +38,7 @@ pub fn linearization_strings<F: ark_ff::PrimeField + ark_ff::SquareRootField>(
     };
     let evaluated_cols = linearization_columns::<F>(features.as_ref());
     let (linearization, _powers_of_alpha) =
-        constraints_expr::<F>(omit_custom_gate, None, features.as_ref(), true);
+        constraints_expr::<F>(/* omit_custom_gate */ false, custom_gate_type, features.as_ref(), true);
 
     let Linearization {
         constant_term,
@@ -56,22 +58,90 @@ pub fn linearization_strings<F: ark_ff::PrimeField + ark_ff::SquareRootField>(
     (constant, other_terms)
 }
 
+#[ocaml::func]
+pub fn fp_linearization_strings_plus() -> (String, Vec<(String, String)>) {
+    // Define conditional gate in RPN
+    //     w(0) = w(1) * w(3) + (1 - w(3)) * w(2)
+    use kimchi::circuits::expr::{PolishToken::*, *};
+    use kimchi::circuits::gate::CurrOrNext::Curr;
+    let conditional_gate = Some(vec![
+        Cell(Variable {
+            col: Column::Index(GateType::ForeignFieldAdd),
+            row: Curr,
+        }),
+        Cell(Variable {
+            col: Column::Witness(3),
+            row: Curr,
+        }),
+        Dup,
+        Mul,
+        Cell(Variable {
+            col: Column::Witness(3),
+            row: Curr,
+        }),
+        Sub,
+        Alpha,
+        Pow(1),
+        Cell(Variable {
+            col: Column::Witness(0),
+            row: Curr,
+        }),
+        Cell(Variable {
+            col: Column::Witness(3),
+            row: Curr,
+        }),
+        Cell(Variable {
+            col: Column::Witness(1),
+            row: Curr,
+        }),
+        Mul,
+        Literal(mina_curves::pasta::Fp::from(1u32)),
+        Cell(Variable {
+            col: Column::Witness(3),
+            row: Curr,
+        }),
+        Sub,
+        Cell(Variable {
+            col: Column::Witness(2),
+            row: Curr,
+        }),
+        Mul,
+        Add,
+        Sub,
+        Mul,
+        Add,
+        Mul,
+    ]);
+
+    linearization_strings::<mina_curves::pasta::Fp>(conditional_gate.as_ref(), true)
+}
+
+#[ocaml::func]
+pub fn fq_linearization_strings_plus() -> (String, Vec<(String, String)>) {
+    linearization_strings::<mina_curves::pasta::Fq>(None, false)
+}
+
+
 #[ocaml::func]
 pub fn fp_linearization_strings_minus() -> (String, Vec<(String, String)>) {
-    linearization_strings::<mina_curves::pasta::Fp>(true, true)
+    // linearization_strings::<mina_curves::pasta::Fp>(true, true)
+    linearization_strings::<mina_curves::pasta::Fp>(None, true)
 }
 
 #[ocaml::func]
 pub fn fq_linearization_strings_minus() -> (String, Vec<(String, String)>) {
-    linearization_strings::<mina_curves::pasta::Fq>(true, false)
+    // linearization_strings::<mina_curves::pasta::Fq>(true, false)
+    linearization_strings::<mina_curves::pasta::Fq>(None, false)
 }
 
 #[ocaml::func]
 pub fn fp_linearization_strings() -> (String, Vec<(String, String)>) {
-    linearization_strings::<mina_curves::pasta::Fp>(false, true)
+    // linearization_strings::<mina_curves::pasta::Fp>(false, true)
+    linearization_strings::<mina_curves::pasta::Fp>(None, true)
 }
 
 #[ocaml::func]
 pub fn fq_linearization_strings() -> (String, Vec<(String, String)>) {
-    linearization_strings::<mina_curves::pasta::Fq>(false, false)
+    // linearization_strings::<mina_curves::pasta::Fq>(false, false)
+    linearization_strings::<mina_curves::pasta::Fq>(None, false)
 }

src/lib/pickles/cache.ml

@@ -83,8 +83,10 @@ module Step = struct
                 (Kimchi_bindings.Protocol.VerifierIndex.Fp.write (Some true) x)
               header path ) )
 
-  let read_or_generate ~(custom_gate_type : bool) ~prev_challenges cache k_p k_v
-      typ return_typ main =
+  let read_or_generate
+      ~(custom_gate_type :
+         Backend.Tick.Field.t Kimchi_types.polish_token array option )
+      ~prev_challenges cache k_p k_v typ return_typ main =
     let s_p = storable in
     let s_v = vk_storable in
     let open Impls.Step in
@@ -183,8 +185,7 @@ module Wrap = struct
                 (Kimchi_bindings.Protocol.Index.Fq.write (Some true) t.index)
               header path ) )
 
-  let read_or_generate ~(custom_gate_type : bool) ~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
@@ -201,7 +202,7 @@ module Wrap = struct
              let r =
                Common.time "wrapkeygen" (fun () ->
                    constraint_system ~input_typ:typ ~return_typ main
-                   |> Keypair.generate ~custom_gate_type ~prev_challenges )
+                   |> Keypair.generate ~prev_challenges )
              in
              ignore
                ( Key_cache.Sync.write cache s_p k (Keypair.pk r)

src/lib/pickles/cache.mli

@@ -21,7 +21,8 @@ module Step : sig
   end
 
   val read_or_generate :
-       custom_gate_type:bool
+       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,8 +62,7 @@ module Wrap : sig
   end
 
   val read_or_generate :
-       custom_gate_type:bool
-    -> 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

@@ -420,16 +420,16 @@ 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 * bool =
+          -> Opt.Flag.t Plonk_types.Features.Full.t
+             * Backend.Tick.Field.t Kimchi_types.polish_token array option =
        fun rules ->
         match rules with
         | [] ->
-            (Plonk_types.Features.Full.none, false)
+            (Plonk_types.Features.Full.none, None)
         | [ r ] ->
             ( Plonk_types.Features.map r.feature_flags ~f:(function
                 | true ->
@@ -440,7 +440,12 @@ struct
             , r.custom_gate_type )
         | r :: rules ->
             let feature_flags, custom_gate_type = go rules in
-            assert (Bool.(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 +687,6 @@ struct
       let r =
         Common.time "wrap read or generate " (fun () ->
             Cache.Wrap.read_or_generate (* Due to Wrap_hack *)
-              ~custom_gate_type
-                (* 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 ())
               main )
@@ -949,7 +952,7 @@ module Side_loaded = struct
                   { constraints = 0 }
               }
             in
-            Verify.Instance.T (max_proofs_verified, m, None, vk, x, p) )
+            Verify.Instance.T (max_proofs_verified, m, None, None, vk, x, p) )
         |> Verify.verify_heterogenous )
 
   let verify ~typ ts = verify_promise ~typ ts |> Promise.to_deferred
@@ -1133,6 +1136,10 @@ let compile_with_wrap_main_override_promise :
           ; zk_rows = compiled.zk_rows
           }
   in
+  let custom_gate_type =
+    let compiled = Types_map.lookup_compiled self.id in
+    compiled.custom_gate_type
+  in
   let module P = struct
     type statement = value
 
@@ -1152,7 +1159,7 @@ let compile_with_wrap_main_override_promise :
     let verification_key = wrap_vk
 
     let verify_promise ts =
-      verify_promise ?chunking_data
+      verify_promise ?chunking_data ?custom_gate_type
         ( module struct
           include Max_proofs_verified
         end )
@@ -1318,7 +1325,7 @@ struct
           ; auxiliary_output = ()
           } )
     ; feature_flags = Plonk_types.Features.none_bool
-    ; custom_gate_type = false
+    ; custom_gate_type = None
     }
 
   let override_wrap_main =
@@ -1387,7 +1394,7 @@ struct
             ~choices:(fun ~self:_ ->
               [ { identifier = "main"
                 ; feature_flags = Plonk_types.Features.none_bool
-                ; custom_gate_type = false
+                ; custom_gate_type = None
                 ; prevs = [ tag; tag ]
                 ; main =
                     (fun { public_input = () } ->

src/lib/pickles/compile.mli

@@ -49,6 +49,7 @@ type chunking_data = Verify.Instance.chunking_data =
 
 val verify_promise :
      ?chunking_data:chunking_data
+  -> ?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
@@ -124,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:bool
+    -> 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.ml

@@ -44,7 +44,7 @@ module Step = struct
     let generate ~custom_gate_type ~prev_challenges cs =
       let open Tick.Keypair in
       let keypair =
-        if custom_gate_type then
+        if Option.is_some custom_gate_type then
           create_plus custom_gate_type ~prev_challenges cs
         else create ~prev_challenges cs
       in
@@ -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 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

@@ -21,7 +21,8 @@ module Step : sig
     val create : pk:Proving_key.t -> vk:Verification_key.t -> t
 
     val generate :
-         custom_gate_type:bool
+         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,8 +111,7 @@ module Wrap : sig
     val create : pk:Proving_key.t -> vk:Verification_key.t -> t
 
     val generate :
-         custom_gate_type:bool
-      -> prev_challenges:int
+         prev_challenges:int
       -> Kimchi_pasta_constraint_system.Pallas_constraint_system.t
       -> t
   end

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 : bool
+  ; 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 : bool
+  ; custom_gate_type :
+      Backend.Tick.Field.t Kimchi_types.polish_token array option
   }
 
 module T