Merge pull request #14412 from MinaProtocol/feature/avoid-cached-vk-r…

…ecomputation [Pickles] Avoid recomputation of cached verification keys
MinaProtocol · Oct 26, 2023 · 196d51b · 196d51b
2 parents 23ea73a + 83e7f97
commit 196d51b
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 17 deletions.
diff --git a/src/lib/pickles/cache.ml b/src/lib/pickles/cache.ml
@@ -94,8 +94,7 @@ module Step = struct
                 Key_cache.Sync.read cache s_p (Lazy.force k_p) )
           with
         | Ok (pk, dirty) ->
-            Common.time "step keypair create" (fun () ->
-                (Keypair.create ~pk ~vk:(Backend.Tick.Keypair.vk pk), dirty) )
+            Common.time "step keypair create" (fun () -> (pk, dirty))
         | Error _e ->
             let r =
               Common.time "stepkeygen" (fun () ->
@@ -106,7 +105,7 @@ module Step = struct
             ignore
               ( Key_cache.Sync.write cache s_p (Lazy.force k_p) (Keypair.pk r)
                 : unit Or_error.t ) ;
-            (r, `Generated_something) )
+            (Keypair.pk r, `Generated_something) )
     in
     let vk =
       lazy
@@ -119,7 +118,7 @@ module Step = struct
              (vk, `Cache_hit)
          | Error _e ->
              let pk, c = Lazy.force pk in
-             let vk = Keypair.vk pk in
+             let vk = Backend.Tick.Keypair.vk pk in
              ignore (Key_cache.Sync.write cache s_v k_v vk : unit Or_error.t) ;
              (vk, c) )
     in
@@ -194,7 +193,7 @@ module Wrap = struct
                Key_cache.Sync.read cache s_p k )
          with
          | Ok (pk, d) ->
-             (Keypair.create ~pk ~vk:(Backend.Tock.Keypair.vk pk), d)
+             (pk, d)
          | Error _e ->
              let r =
                Common.time "wrapkeygen" (fun () ->
@@ -204,7 +203,7 @@ module Wrap = struct
              ignore
                ( Key_cache.Sync.write cache s_p k (Keypair.pk r)
                  : unit Or_error.t ) ;
-             (r, `Generated_something) )
+             (Keypair.pk r, `Generated_something) )
     in
     let vk =
       lazy
@@ -247,9 +246,8 @@ module Wrap = struct
          | Ok (vk, d) ->
              (vk, d)
          | Error _e ->
-             let kp, _dirty = Lazy.force pk in
-             let vk = Keypair.vk kp in
-             let pk = Keypair.pk kp in
+             let pk, _dirty = Lazy.force pk in
+             let vk = Backend.Tock.Keypair.vk pk in
              let vk : Vk.t =
                { index = vk
                ; commitments =

diff --git a/src/lib/pickles/cache.mli b/src/lib/pickles/cache.mli
@@ -28,7 +28,7 @@ module Step : sig
     -> ('a, 'b) Impls.Step.Typ.t
     -> ('c, 'd) Impls.Step.Typ.t
     -> ('a -> unit -> 'c)
-    -> ( Impls.Step.Keypair.t
+    -> ( Impls.Step.Proving_key.t
        * [> `Cache_hit | `Generated_something | `Locally_generated ] )
        lazy_t
        * ( Kimchi_bindings.Protocol.VerifierIndex.Fp.t
@@ -67,7 +67,7 @@ module Wrap : sig
     -> ('a, 'b) Impls.Wrap.Typ.t
     -> ('c, 'd) Impls.Wrap.Typ.t
     -> ('a -> unit -> 'c)
-    -> ( Impls.Wrap.Keypair.t
+    -> ( Impls.Wrap.Proving_key.t
        * [> `Cache_hit | `Generated_something | `Locally_generated ] )
        lazy_t
        * ( Verification_key.Stable.V2.t

diff --git a/src/lib/pickles/compile.ml b/src/lib/pickles/compile.ml
@@ -277,7 +277,7 @@ struct
 
   module Lazy_keys = struct
     type t =
-      (Impls.Step.Keypair.t * Dirty.t) Lazy.t
+      (Impls.Step.Proving_key.t * Dirty.t) Lazy.t
       * (Kimchi_bindings.Protocol.VerifierIndex.Fp.t * Dirty.t) Lazy.t
 
     (* TODO Think this is right.. *)
@@ -733,7 +733,7 @@ struct
                  ~f:(fun x -> [| x |])
                  wrap_vk.commitments )
             ~public_input ~auxiliary_typ ~feature_flags
-            (Impls.Step.Keypair.pk (fst (Lazy.force step_pk)))
+            (fst (Lazy.force step_pk))
             wrap_vk.index
         in
         let step_vk = fst (Lazy.force step_vk) in
@@ -781,7 +781,7 @@ struct
                    wrap_vk.commitments )
               wrap_main ~typ ~step_vk ~step_plonk_indices:(Lazy.force step_vks)
               ~actual_wrap_domains
-              (Impls.Wrap.Keypair.pk (fst (Lazy.force wrap_pk)))
+              (fst (Lazy.force wrap_pk))
               proof
           in
           ( return_value

diff --git a/src/lib/pickles/pickles.ml b/src/lib/pickles/pickles.ml
@@ -1076,7 +1076,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
 
         module Lazy_keys = struct
           type t =
-            (Impls.Step.Keypair.t * Dirty.t) Lazy.t
+            (Impls.Step.Proving_key.t * Dirty.t) Lazy.t
             * (Kimchi_bindings.Protocol.VerifierIndex.Fp.t * Dirty.t) Lazy.t
 
           (* TODO Think this is right.. *)
@@ -1297,7 +1297,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
                     ((* TODO *) Plonk_verification_key_evals.map
                        ~f:(fun x -> [| x |])
                        wrap_vk.commitments )
-                  (Impls.Step.Keypair.pk (fst (Lazy.force step_pk)))
+                  (fst (Lazy.force step_pk))
                   wrap_vk.index
               in
               let pairing_vk = fst (Lazy.force step_vk) in
@@ -1785,7 +1785,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
                     wrap_main A_value.to_field_elements ~pairing_vk
                     ~step_domains:b.domains
                     ~pairing_plonk_indices:(Lazy.force step_vks) ~wrap_domains
-                    (Impls.Wrap.Keypair.pk (fst (Lazy.force wrap_pk)))
+                    (fst (Lazy.force wrap_pk))
                     proof
                 in
                 Proof.T