Use tx segments in delegation program

src/lib/mina_lib/mina_lib.ml

@@ -1528,7 +1528,8 @@ let start t =
     ~graphql_control_port:t.config.graphql_control_port ~built_with_commit_sha
     ~get_next_producer_timing:(fun () -> t.next_producer_timing)
     ~get_snark_work_fee:(fun () -> snark_work_fee t)
-    ~get_peer:(fun () -> t.config.gossip_net_params.addrs_and_ports.peer) ;
+    ~get_peer:(fun () -> t.config.gossip_net_params.addrs_and_ports.peer)
+    ~signature_kind:t.signature_kind ;
   stop_long_running_daemon t ;
   Snark_worker.start t
 

src/lib/uptime_service/uptime_service.ml

@@ -203,7 +203,7 @@ let read_all_proofs_for_work_single_spec =
 let send_block_and_transaction_snark ~logger ~constraint_constants ~interruptor
     ~url ~snark_worker ~transition_frontier ~peer_id
     ~(submitter_keypair : Keypair.t) ~snark_work_fee ~graphql_control_port
-    ~built_with_commit_sha =
+    ~built_with_commit_sha ~signature_kind =
   match Broadcast_pipe.Reader.peek transition_frontier with
   | None ->
       (* expected during daemon boot, so not logging as error *)
@@ -327,12 +327,48 @@ let send_block_and_transaction_snark ~logger ~constraint_constants ~interruptor
                 send_uptime_data ~logger ~interruptor ~submitter_keypair ~url
                   ~state_hash ~produced:false block_data
             | Some single_spec -> (
-                match%bind
-                  make_interruptible
-                    (Uptime_snark_worker.perform_single snark_worker
-                       ( message
-                       , read_all_proofs_for_work_single_spec single_spec ) )
-                with
+                let module T = Transaction_snark.Make (struct
+                  let signature_kind = signature_kind
+
+                  let constraint_constants = constraint_constants
+
+                  let proof_level = Genesis_constants.Proof_level.Full
+                end) in
+                let s =
+                  match single_spec with
+                  | Transition (input, witness) -> (
+                      match witness.transaction with
+                      | Command (Zkapp_command zkapp_command) ->
+                          Ok (`Zk_app (input, witness, zkapp_command))
+                      | Command (Signed_command _) | Fee_transfer _ | Coinbase _
+                        ->
+                          Ok (`Transaction single_spec) )
+                  | Merge _ ->
+                      Error
+                        (Error.of_string "Undexpecetd merge operation in spec")
+                in
+                let work =
+                  match s with
+                  | Ok (`Zk_app (input, witness, zkapp_command)) ->
+                      let zkapp_command =
+                        Zkapp_command.read_all_proofs_from_disk zkapp_command
+                      in
+                      let witness =
+                        Transaction_witness.read_all_proofs_from_disk witness
+                      in
+                      make_interruptible
+                      @@ Uptime_snark_worker.perform_partitioned snark_worker
+                           (witness, input, zkapp_command, staged_ledger_hash)
+                  | Ok (`Transaction single_spec) ->
+                      make_interruptible
+                      @@ Uptime_snark_worker.perform_single snark_worker
+                           ( message
+                           , read_all_proofs_for_work_single_spec single_spec )
+                  | Error error ->
+                      Interruptible.return (Error error)
+                in
+
+                match%bind work with
                 | Error e ->
                     (* error in submitting to process *)
                     [%log error]
@@ -373,7 +409,8 @@ let send_block_and_transaction_snark ~logger ~constraint_constants ~interruptor
 let start ~logger ~uptime_url ~snark_worker_opt ~constraint_constants
     ~protocol_constants ~transition_frontier ~time_controller
     ~block_produced_bvar ~uptime_submitter_keypair ~get_next_producer_timing
-    ~get_snark_work_fee ~get_peer ~graphql_control_port ~built_with_commit_sha =
+    ~get_snark_work_fee ~get_peer ~graphql_control_port ~built_with_commit_sha
+    ~signature_kind =
   match uptime_url with
   | None ->
       [%log info] "Not running uptime service, no URL given" ;
@@ -475,7 +512,7 @@ let start ~logger ~uptime_url ~snark_worker_opt ~constraint_constants
                   send_block_and_transaction_snark ~logger ~interruptor ~url
                     ~constraint_constants ~snark_worker ~transition_frontier
                     ~peer_id ~submitter_keypair ~snark_work_fee
-                    ~graphql_control_port ~built_with_commit_sha
+                    ~graphql_control_port ~built_with_commit_sha ~signature_kind
                 in
                 match get_next_producer_time_opt () with
                 | None ->

src/lib/uptime_service/uptime_snark_worker.ml

@@ -8,6 +8,31 @@ open struct
   module Impl = Snark_worker.Impl
 end
 
+let extract_terminal_zk_segment ~(m : (module Transaction_snark.S)) ~witness
+    ~input ~zkapp_command ~staged_ledger_hash =
+  let staged_ledger_hash = Staged_ledger_hash.ledger_hash staged_ledger_hash in
+  let%bind.Result final_segment =
+    Work_partitioner.Snark_worker_shared.extract_zkapp_segment_works ~m ~input
+      ~witness ~zkapp_command
+    |> Result.map_error
+         ~f:
+           Work_partitioner.Snark_worker_shared.Failed_to_generate_inputs
+           .error_of_t
+    |> Result.map ~f:(function x ->
+           Work_partitioner.Snark_worker_shared.Zkapp_command_inputs
+           .read_all_proofs_from_disk x
+           |> Mina_stdlib.Nonempty_list.find ~f:(function _, _, s ->
+                  Ledger_hash.(s.target.second_pass_ledger = staged_ledger_hash) ) )
+  in
+  match final_segment with
+  | Some res ->
+      Ok res
+  | _ ->
+      Error
+        ( Error.of_string
+        @@ sprintf "Failed to find zkapp segment with target hash %s"
+             (Ledger_hash0.to_base58_check staged_ledger_hash) )
+
 module Worker = struct
   module T = struct
     module F = Rpc_parallel.Function
@@ -18,6 +43,14 @@ module Worker = struct
           , Sok_message.t * Snark_work_lib.Spec.Single.Stable.Latest.t
           , (Ledger_proof.t * Time.Span.t) Or_error.t )
           F.t
+      ; perform_partitioned :
+          ( 'w
+          , Transaction_witness.Stable.Latest.t
+            * Mina_state.Snarked_ledger_state.Stable.Latest.t
+            * Zkapp_command.Stable.Latest.t
+            * Staged_ledger_hash.t
+          , (Ledger_proof.t * Time.Span.t) Or_error.t )
+          F.t
       }
 
     module Worker_state = struct
@@ -43,6 +76,41 @@ module Worker = struct
       let perform_single (state : Worker_state.t) (message, single_spec) =
         Impl.perform_single ~message state single_spec
 
+      let perform_partitioned (state : Worker_state.t)
+          (witness, statement, zkapp_command, staged_ledger_hash) =
+        let zkapp_command =
+          Zkapp_command.write_all_proofs_to_disk
+            ~signature_kind:state.signature_kind
+            ~proof_cache_db:state.proof_cache_db zkapp_command
+        in
+        match state.proof_level_snark with
+        | Full (module S) ->
+            let%bind.Deferred.Or_error witness, spec, statement =
+              extract_terminal_zk_segment
+                ~m:(module S)
+                ~witness ~input:statement ~zkapp_command ~staged_ledger_hash
+              |> Deferred.return
+            in
+
+            Snark_worker.Impl.measure_runtime ~logger:state.logger
+              ~spec_json:
+                ( lazy
+                  ( "zkapp_segment_spec"
+                  , Transaction_snark.Zkapp_command_segment.Basic.Stable.Latest
+                    .to_yojson spec ) )
+              (fun () ->
+                let witness =
+                  Transaction_witness.Zkapp_command_segment_witness
+                  .write_all_proofs_to_disk witness
+                    ~proof_cache_db:state.proof_cache_db
+                    ~signature_kind:state.signature_kind
+                in
+                S.of_zkapp_command_segment_exn ~statement ~witness ~spec
+                |> Deferred.map ~f:(fun a -> Result.Ok a) )
+        | _ ->
+            Deferred.Or_error.error_string
+              "Unexpected prover mode in uptime snark worker"
+
       let functions =
         let f (i, o, f) =
           C.create_rpc
@@ -57,6 +125,16 @@ module Worker = struct
               , [%bin_type_class:
                   (Ledger_proof.Stable.Latest.t * Time.Span.t) Or_error.t]
               , perform_single )
+        ; perform_partitioned =
+            f
+              ( [%bin_type_class:
+                  Transaction_witness.Stable.Latest.t
+                  * Mina_state.Snarked_ledger_state.Stable.Latest.t
+                  * Zkapp_command.Stable.Latest.t
+                  * Staged_ledger_hash.Stable.Latest.t]
+              , [%bin_type_class:
+                  (Ledger_proof.Stable.Latest.t * Time.Span.t) Or_error.t]
+              , perform_partitioned )
         }
 
       let init_worker_state (logger, constraint_constants) =
@@ -123,3 +201,6 @@ let create ~logger ~constraint_constants ~pids : t Deferred.t =
 
 let perform_single { connection; _ } ((_message, _single_spec) as arg) =
   Worker.Connection.run connection ~f:Worker.functions.perform_single ~arg
+
+let perform_partitioned { connection; _ } arg =
+  Worker.Connection.run connection ~f:Worker.functions.perform_partitioned ~arg