[#16112] Improve lagrange_commitments_whole_domain

Author: volhovm

src/lib/crypto/kimchi_backend/pasta/pallas_based_plonk.ml

@@ -34,10 +34,11 @@ end
 module R1CS_constraint_system =
   Kimchi_pasta_constraint_system.Pallas_constraint_system
 
-let lagrange srs domain_log2 : _ Kimchi_types.poly_comm array =
+let lagrange (srs : Kimchi_bindings.Protocol.SRS.Fq.t) domain_log2 :
+    _ Kimchi_types.poly_comm array =
   let domain_size = Int.pow 2 domain_log2 in
-  Array.init domain_size ~f:(fun i ->
-      Kimchi_bindings.Protocol.SRS.Fq.lagrange_commitment srs domain_size i )
+  Kimchi_bindings.Protocol.SRS.Fq.lagrange_commitments_whole_domain srs
+    domain_size
 
 let with_lagrange f (vk : Verification_key.t) =
   f (lagrange vk.srs vk.domain.log_size_of_group) vk

src/lib/crypto/kimchi_backend/pasta/vesta_based_plonk.ml

@@ -35,8 +35,8 @@ module R1CS_constraint_system =
 
 let lagrange srs domain_log2 : _ Kimchi_types.poly_comm array =
   let domain_size = Int.pow 2 domain_log2 in
-  Array.init domain_size ~f:(fun i ->
-      Kimchi_bindings.Protocol.SRS.Fp.lagrange_commitment srs domain_size i )
+  Kimchi_bindings.Protocol.SRS.Fp.lagrange_commitments_whole_domain srs
+    domain_size
 
 let with_lagrange f (vk : Verification_key.t) =
   f (lagrange vk.srs vk.domain.log_size_of_group) vk

src/lib/crypto/kimchi_bindings/stubs/kimchi_bindings.ml

@@ -107,6 +107,12 @@ module Protocol = struct
         -> Pasta_bindings.Fq.t Kimchi_types.or_infinity Kimchi_types.poly_comm
         = "caml_fp_srs_lagrange_commitment"
 
+      external lagrange_commitments_whole_domain :
+           t
+        -> int
+        -> Pasta_bindings.Fq.t Kimchi_types.or_infinity Kimchi_types.poly_comm
+           array = "caml_fp_srs_lagrange_commitments_whole_domain"
+
       external add_lagrange_basis : t -> int -> unit
         = "caml_fp_srs_add_lagrange_basis"
 
@@ -156,6 +162,12 @@ module Protocol = struct
         -> Pasta_bindings.Fp.t Kimchi_types.or_infinity Kimchi_types.poly_comm
         = "caml_fq_srs_lagrange_commitment"
 
+      external lagrange_commitments_whole_domain :
+           t
+        -> int
+        -> Pasta_bindings.Fp.t Kimchi_types.or_infinity Kimchi_types.poly_comm
+           array = "caml_fq_srs_lagrange_commitments_whole_domain"
+
       external add_lagrange_basis : t -> int -> unit
         = "caml_fq_srs_add_lagrange_basis"
 

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

@@ -363,6 +363,7 @@ fn generate_kimchi_bindings(mut w: impl std::io::Write, env: &mut Env) {
                 decl_func!(w, env, caml_fp_srs_write => "write");
                 decl_func!(w, env, caml_fp_srs_read => "read");
                 decl_func!(w, env, caml_fp_srs_lagrange_commitment => "lagrange_commitment");
+                decl_func!(w, env, caml_fp_srs_lagrange_commitments_whole_domain => "lagrange_commitments_whole_domain");
                 decl_func!(w, env, caml_fp_srs_add_lagrange_basis=> "add_lagrange_basis");
                 decl_func!(w, env, caml_fp_srs_commit_evaluations => "commit_evaluations");
                 decl_func!(w, env, caml_fp_srs_b_poly_commitment => "b_poly_commitment");
@@ -378,6 +379,7 @@ fn generate_kimchi_bindings(mut w: impl std::io::Write, env: &mut Env) {
                 decl_func!(w, env, caml_fq_srs_write => "write");
                 decl_func!(w, env, caml_fq_srs_read => "read");
                 decl_func!(w, env, caml_fq_srs_lagrange_commitment => "lagrange_commitment");
+                decl_func!(w, env, caml_fq_srs_lagrange_commitments_whole_domain => "lagrange_commitments_whole_domain");
                 decl_func!(w, env, caml_fq_srs_add_lagrange_basis=> "add_lagrange_basis");
                 decl_func!(w, env, caml_fq_srs_commit_evaluations => "commit_evaluations");
                 decl_func!(w, env, caml_fq_srs_b_poly_commitment => "b_poly_commitment");

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

@@ -71,6 +71,18 @@ macro_rules! impl_srs {
                 Ok(Some($name::new(srs)))
             }
 
+            #[ocaml_gen::func]
+            #[ocaml::func]
+            /// This is same as _lagrange_commitments, but returns the result for every
+            /// i <= domain_size.
+            pub fn [<$name:snake _lagrange_commitments_whole_domain>](
+                srs: $name,
+                domain_size: ocaml::Int,
+            ) -> Result<Vec<CamlPolyComm<$CamlG>>, ocaml::Error> {
+                Ok(srs.get_lagrange_basis_from_domain_size(domain_size as usize).clone().into_iter().map(|x| x.into()).collect())
+            }
+
+
             #[ocaml_gen::func]
             #[ocaml::func]
             pub fn [<$name:snake _lagrange_commitment>](