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

@@ -200,10 +200,37 @@ let read_all_proofs_for_work_single_spec =
     ~f_proof:Ledger_proof.Cached.read_proof_from_disk
     ~f_witness:Transaction_witness.read_all_proofs_from_disk
 
+let extract_terminal_zk_segment ~(m : (module Transaction_snark.S))
+    ~(witness : Transaction_witness.t) ~input ~zkapp_command ~staged_ledger_hash
+    =
+  let staged_ledger_hash = Staged_ledger_hash.ledger_hash staged_ledger_hash in
+  let witness = Transaction_witness.read_all_proofs_from_disk witness 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 (witness, spec, statement) ->
+      Ok (witness, statement, spec)
+  | _ ->
+      Error
+        ( Error.of_string
+        @@ sprintf "Failed to find zkapp segment with target hash %s"
+             (Ledger_hash0.to_base58_check staged_ledger_hash) )
+
 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 +354,47 @@ 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) ->
+                          let%map.Result final_segment =
+                            extract_terminal_zk_segment
+                              ~m:(module T)
+                              ~witness ~input ~zkapp_command ~staged_ledger_hash
+                          in
+                          `Zk_app final_segment
+                      | 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 (witness, statement, spec)) ->
+                      make_interruptible
+                      @@ Uptime_snark_worker.perform_partitioned snark_worker
+                           (witness, statement, spec)
+                  | 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 +435,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 +538,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

@@ -18,6 +18,13 @@ 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.Zkapp_command_segment_witness.Stable.Latest.t
+            * Mina_state.Snarked_ledger_state.With_sok.t
+            * Transaction_snark.Zkapp_command_segment.Basic.t
+          , (Ledger_proof.t * Time.Span.t) Or_error.t )
+          F.t
       }
 
     module Worker_state = struct
@@ -43,6 +50,28 @@ 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, spec)
+          =
+        match state.proof_level_snark with
+        | Full (module S) ->
+            let witness =
+              Transaction_witness.Zkapp_command_segment_witness
+              .write_all_proofs_to_disk ~signature_kind:state.signature_kind
+                ~proof_cache_db:state.proof_cache_db witness
+            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 () ->
+                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 +86,18 @@ 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.Zkapp_command_segment_witness.Stable
+                  .Latest
+                  .t
+                  * Mina_state.Snarked_ledger_state.With_sok.Stable.Latest.t
+                  * Transaction_snark.Zkapp_command_segment.Basic.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 +164,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