Merge pull request #16244 from MinaProtocol/al/snark-work-graphql

svv232 · web-flow · commit c1634c3956b7 · 2024-12-09T16:10:41.000-05:00
GraphQL: add a query to fetch specific pending snark work
diff --git a/graphql_schema.json b/graphql_schema.json
@@ -6660,6 +6660,107 @@
           "enumValues": null,
           "possibleTypes": null
         },
+        {
+          "kind": "OBJECT",
+          "name": "WorkBundleSpec",
+          "description":
+            "Witnesses and statements for snark work bundles. Includes optional fees and prover public keys for ones that have proofs in the snark pool",
+          "fields": [
+            {
+              "name": "spec",
+              "description": "Snark work specification in json format",
+              "args": [],
+              "type": {
+                "kind": "NON_NULL",
+                "name": null,
+                "ofType": {
+                  "kind": "SCALAR",
+                  "name": "String",
+                  "ofType": null
+                }
+              },
+              "isDeprecated": false,
+              "deprecationReason": null
+            },
+            {
+              "name": "snarkFee",
+              "description": "Fee if proof for the spec exists in snark pool",
+              "args": [],
+              "type": { "kind": "SCALAR", "name": "Fee", "ofType": null },
+              "isDeprecated": false,
+              "deprecationReason": null
+            },
+            {
+              "name": "snarkProverKey",
+              "description":
+                "Prover public key if proof for the spec exists in snark pool",
+              "args": [],
+              "type": {
+                "kind": "SCALAR",
+                "name": "PublicKey",
+                "ofType": null
+              },
+              "isDeprecated": false,
+              "deprecationReason": null
+            },
+            {
+              "name": "workIds",
+              "description": "Unique identifier for a snark work",
+              "args": [],
+              "type": {
+                "kind": "NON_NULL",
+                "name": null,
+                "ofType": {
+                  "kind": "LIST",
+                  "name": null,
+                  "ofType": {
+                    "kind": "NON_NULL",
+                    "name": null,
+                    "ofType": {
+                      "kind": "SCALAR",
+                      "name": "Int",
+                      "ofType": null
+                    }
+                  }
+                }
+              },
+              "isDeprecated": false,
+              "deprecationReason": null
+            }
+          ],
+          "inputFields": null,
+          "interfaces": [],
+          "enumValues": null,
+          "possibleTypes": null
+        },
+        {
+          "kind": "OBJECT",
+          "name": "PendingSnarkWorkSpec",
+          "description":
+            "Snark work witnesses and statements that are yet to be proven or included in blocks",
+          "fields": [
+            {
+              "name": "workBundleSpec",
+              "description": "Work bundle spec with one or two snark work",
+              "args": [],
+              "type": {
+                "kind": "NON_NULL",
+                "name": null,
+                "ofType": {
+                  "kind": "OBJECT",
+                  "name": "WorkBundleSpec",
+                  "ofType": null
+                }
+              },
+              "isDeprecated": false,
+              "deprecationReason": null
+            }
+          ],
+          "inputFields": null,
+          "interfaces": [],
+          "enumValues": null,
+          "possibleTypes": null
+        },
         {
           "kind": "OBJECT",
           "name": "WorkDescription",
@@ -15282,6 +15383,58 @@
               "isDeprecated": false,
               "deprecationReason": null
             },
+            {
+              "name": "snarkWorkRange",
+              "description":
+                "Find any sequence of snark work between two indexes in all available snark work. Returns both completed and uncompleted work.",
+              "args": [
+                {
+                  "name": "endingIndex",
+                  "description":
+                    "The last index to be taken from all available snark work (exclusive). If not specified or greater than the available snark work list,all elements from index [startingIndex] will be returned. An empty list will be returned if startingIndex is not a valid index or if startingIndex >= endingIndex.",
+                  "type": {
+                    "kind": "SCALAR",
+                    "name": "UInt32",
+                    "ofType": null
+                  },
+                  "defaultValue": null
+                },
+                {
+                  "name": "startingIndex",
+                  "description":
+                    "The first index to be taken from all available snark work",
+                  "type": {
+                    "kind": "NON_NULL",
+                    "name": null,
+                    "ofType": {
+                      "kind": "SCALAR",
+                      "name": "UInt32",
+                      "ofType": null
+                    }
+                  },
+                  "defaultValue": null
+                }
+              ],
+              "type": {
+                "kind": "NON_NULL",
+                "name": null,
+                "ofType": {
+                  "kind": "LIST",
+                  "name": null,
+                  "ofType": {
+                    "kind": "NON_NULL",
+                    "name": null,
+                    "ofType": {
+                      "kind": "OBJECT",
+                      "name": "PendingSnarkWorkSpec",
+                      "ofType": null
+                    }
+                  }
+                }
+              },
+              "isDeprecated": false,
+              "deprecationReason": null
+            },
             {
               "name": "snarkedLedgerAccountMembership",
               "description":
diff --git a/src/lib/mina_graphql/mina_graphql.ml b/src/lib/mina_graphql/mina_graphql.ml
@@ -2216,6 +2216,52 @@ module Queries = struct
         Work_selector.pending_work_statements ~snark_pool ~fee_opt
           snark_job_state )
 
+  let snark_work_range =
+    field "snarkWorkRange"
+      ~doc:
+        "Find any sequence of snark work between two indexes in all available \
+         snark work. Returns both completed and uncompleted work."
+      ~args:
+        Arg.
+          [ arg "startingIndex"
+              ~doc:"The first index to be taken from all available snark work"
+              ~typ:(non_null Types.Input.UInt32.arg_typ)
+          ; arg "endingIndex"
+              ~doc:
+                "The last index to be taken from all available snark work \
+                 (exclusive). If not specified or greater than the available \
+                 snark work list,all elements from index [startingIndex] will \
+                 be returned. An empty list will be returned if startingIndex \
+                 is not a valid index or if startingIndex >= endingIndex."
+              ~typ:Types.Input.UInt32.arg_typ
+          ]
+      ~typ:(non_null @@ list @@ non_null Types.pending_work_spec)
+      ~resolve:(fun { ctx = mina; _ } () start_idx end_idx ->
+        let snark_job_state = Mina_lib.snark_job_state mina in
+        let snark_pool = Mina_lib.snark_pool mina in
+        let all_work = Work_selector.all_work ~snark_pool snark_job_state in
+        let work_size = all_work |> List.length |> Unsigned.UInt32.of_int in
+        let less_than uint1 uint2 = Unsigned.UInt32.compare uint1 uint2 < 0 in
+        let to_bundle_specs =
+          List.map ~f:(fun (spec, fee_prover) ->
+              { Types.Snark_work_bundle.spec; fee_prover } )
+        in
+        match end_idx with
+        | None when less_than start_idx work_size ->
+            (* drop handles case when start_idx is greater than pending work and is O(start_idx)*)
+            let start = Unsigned.UInt32.to_int start_idx in
+            List.drop all_work start |> to_bundle_specs
+        | Some end_idx
+          when less_than start_idx end_idx && less_than start_idx work_size ->
+            let pos = Unsigned.UInt32.to_int start_idx in
+            let len =
+              Unsigned.UInt32.(
+                min (sub end_idx start_idx) (sub work_size start_idx) |> to_int)
+            in
+            List.sub ~pos ~len all_work |> to_bundle_specs
+        | _ ->
+            [] )
+
   module SnarkedLedgerMembership = struct
     let resolve_membership :
            mapper:(Ledger.path -> Account.t -> 'a)
@@ -2784,6 +2830,7 @@ module Queries = struct
     ; trust_status_all
     ; snark_pool
     ; pending_snark_work
+    ; snark_work_range
     ; SnarkedLedgerMembership.snarked_ledger_account_membership
     ; SnarkedLedgerMembership.encoded_snarked_ledger_account_membership
     ; genesis_constants
diff --git a/src/lib/mina_graphql/types.ml b/src/lib/mina_graphql/types.ml
@@ -622,6 +622,55 @@ let pending_work =
           ~resolve:(fun _ w -> One_or_two.to_list w)
       ] )
 
+module Snark_work_bundle = struct
+  type t =
+    { spec : Work_selector.work One_or_two.t
+    ; fee_prover : (Currency.Fee.t * Public_key.Compressed.t) option
+    }
+
+  let spec =
+    obj "WorkBundleSpec"
+      ~doc:
+        "Witnesses and statements for snark work bundles. Includes optional \
+         fees and prover public keys for ones that have proofs in the snark \
+         pool" ~fields:(fun _ ->
+        [ field "spec" ~typ:(non_null string)
+            ~doc:"Snark work specification in json format"
+            ~args:Arg.[]
+            ~resolve:(fun _ { spec; _ } ->
+              One_or_two.to_yojson Work_selector.work_to_yojson spec
+              |> Yojson.Safe.to_string )
+        ; field "snarkFee" ~typ:fee
+            ~doc:"Fee if proof for the spec exists in snark pool"
+            ~args:Arg.[]
+            ~resolve:(fun _ w -> Option.map w.fee_prover ~f:fst)
+        ; field "snarkProverKey" ~typ:public_key
+            ~doc:"Prover public key if proof for the spec exists in snark pool"
+            ~args:Arg.[]
+            ~resolve:(fun _ w -> Option.map w.fee_prover ~f:snd)
+        ; field "workIds" ~doc:"Unique identifier for a snark work"
+            ~typ:(non_null (list (non_null int)))
+            ~args:Arg.[]
+            ~resolve:(fun _ { spec; _ } ->
+              One_or_two.map spec ~f:(fun w ->
+                  Transaction_snark.Statement.hash
+                    (Snark_work_lib.Work.Single.Spec.statement w) )
+              |> One_or_two.to_list )
+        ] )
+end
+
+let pending_work_spec =
+  obj "PendingSnarkWorkSpec"
+    ~doc:
+      "Snark work witnesses and statements that are yet to be proven or \
+       included in blocks" ~fields:(fun _ ->
+      [ field "workBundleSpec"
+          ~args:Arg.[]
+          ~doc:"Work bundle spec with one or two snark work"
+          ~typ:(non_null Snark_work_bundle.spec)
+          ~resolve:(fun _ w -> w)
+      ] )
+
 let state_stack =
   let module M = Pending_coinbase.State_stack in
   obj "StateStack" ~fields:(fun _ ->
diff --git a/src/lib/transaction_snark_work/transaction_snark_work.ml b/src/lib/transaction_snark_work/transaction_snark_work.ml
@@ -122,3 +122,5 @@ end
 let forget = Fn.id
 
 let fee { fee; _ } = fee
+
+let prover { prover; _ } = prover
diff --git a/src/lib/transaction_snark_work/transaction_snark_work.mli b/src/lib/transaction_snark_work/transaction_snark_work.mli
@@ -60,6 +60,8 @@ type t = Mina_wire_types.Transaction_snark_work.V2.t =
 
 val fee : t -> Fee.t
 
+val prover : t -> Public_key.Compressed.t
+
 val info : t -> Info.t
 
 val statement : t -> Statement.t
diff --git a/src/lib/work_selector/inputs.ml b/src/lib/work_selector/inputs.ml
@@ -21,6 +21,8 @@ module Test_inputs = struct
 
     let fee = Fn.id
 
+    let prover _ = Key_gen.Sample_keypairs.genesis_winner |> fst
+
     module Statement = struct
       type t = Transaction_snark.Statement.t One_or_two.t
     end
diff --git a/src/lib/work_selector/intf.ml b/src/lib/work_selector/intf.ml
@@ -27,6 +27,8 @@ module type Inputs_intf = sig
 
     val fee : t -> Fee.t
 
+    val prover : t -> Signature_lib.Public_key.Compressed.t
+
     module Statement : sig
       type t = Transaction_snark.Statement.t One_or_two.t
     end
diff --git a/src/lib/work_selector/work_lib.ml b/src/lib/work_selector/work_lib.ml
@@ -147,6 +147,19 @@ module Make (Inputs : Intf.Inputs_intf) = struct
     List.filter statements ~f:(fun st ->
         Option.is_none (Inputs.Snark_pool.get_completed_work snark_pool st) )
 
+  let all_work ~snark_pool (state : State.t) =
+    O1trace.sync_thread "work_lib_all_unseen_works" (fun () ->
+        List.map state.available_jobs ~f:(fun job ->
+            let statement = One_or_two.map ~f:Work_spec.statement job in
+            let fee_prover_opt =
+              Option.map
+                (Inputs.Snark_pool.get_completed_work snark_pool statement)
+                ~f:(fun (p : Inputs.Transaction_snark_work.t) ->
+                  ( Inputs.Transaction_snark_work.fee p
+                  , Inputs.Transaction_snark_work.prover p ) )
+            in
+            (job, fee_prover_opt) ) )
+
   let all_completed_work ~snark_pool statements =
     List.filter_map statements ~f:(fun st ->
         Inputs.Snark_pool.get_completed_work snark_pool st )
diff --git a/src/lib/work_selector/work_selector.ml b/src/lib/work_selector/work_selector.ml
@@ -3,6 +3,7 @@ module State = Lib.State
 
 type work =
   (Transaction_witness.t, Ledger_proof.t) Snark_work_lib.Work.Single.Spec.t
+[@@deriving yojson]
 
 type snark_pool = Network_pool.Snark_pool.t
 
@@ -24,4 +25,6 @@ let remove = Lib.State.remove
 
 let pending_work_statements = Lib.pending_work_statements
 
+let all_work = Lib.all_work
+
 let completed_work_statements = Lib.completed_work_statements
diff --git a/src/lib/work_selector/work_selector.mli b/src/lib/work_selector/work_selector.mli
@@ -3,6 +3,7 @@ module State :
 
 type work =
   (Transaction_witness.t, Ledger_proof.t) Snark_work_lib.Work.Single.Spec.t
+[@@deriving yojson]
 
 type snark_pool = Network_pool.Snark_pool.t
 
@@ -32,5 +33,12 @@ val pending_work_statements :
   -> State.t
   -> Transaction_snark.Statement.t One_or_two.t list
 
+val all_work :
+     snark_pool:snark_pool
+  -> State.t
+  -> ( work One_or_two.t
+     * (Currency.Fee.t * Signature_lib.Public_key.Compressed.t) option )
+     list
+
 val completed_work_statements :
   snark_pool:snark_pool -> State.t -> Transaction_snark_work.unchecked list
