Merge pull request #1117 from openmina/feat/snark_worker

Rust based Snark Worker

Author: binier

ledger/src/port_ocaml/hash.rs

@@ -22,6 +22,7 @@ impl Default for JaneStreetHasher {
     }
 }
 
+#[allow(clippy::precedence)]
 fn rotl32(x: u32, n: u32) -> u32 {
     (x) << n | (x) >> (32 - n)
 }

macros/src/action_event.rs

@@ -165,9 +165,8 @@ fn filter_fields(field_spec: &FieldsSpec, fields: &FieldsNamed) -> Result<Vec<To
         FieldsSpec::Some(f) => f
             .iter()
             .filter(|(name, _)| {
-                name.as_ref().map_or(true, |name| {
-                    fields.named.iter().any(|n| Some(name) == n.ident.as_ref())
-                })
+                name.as_ref()
+                    .is_none_or(|name| fields.named.iter().any(|n| Some(name) == n.ident.as_ref()))
             })
             .map(|(_, expr)| Ok(expr.clone()))
             .collect(),

node/common/src/service/builder.rs

@@ -146,6 +146,8 @@ impl NodeServiceCommonBuilder {
             ),
             ledger_manager,
             block_producer: self.block_producer,
+            // initialized in state machine.
+            snark_worker: None,
             archive: self.archive,
             p2p,
             stats: self.gather_stats.then(Stats::new),

node/common/src/service/service.rs

@@ -23,6 +23,7 @@ use super::{
     p2p::webrtc_with_libp2p::P2pServiceCtx,
     replay::ReplayerState,
     rpc::{RpcSender, RpcService},
+    snark_worker::SnarkWorker,
     snarks::SnarkBlockVerifyArgs,
     EventReceiver, EventSender,
 };
@@ -41,6 +42,7 @@ pub struct NodeService {
     pub snark_block_proof_verify: mpsc::TrackedUnboundedSender<SnarkBlockVerifyArgs>,
 
     pub ledger_manager: LedgerManager,
+    pub snark_worker: Option<SnarkWorker>,
     pub block_producer: Option<BlockProducerService>,
     pub archive: Option<ArchiveService>,
     pub p2p: P2pServiceCtx,
@@ -116,6 +118,7 @@ impl NodeService {
             event_receiver: mpsc::unbounded_channel().1.into(),
             snark_block_proof_verify: mpsc::unbounded_channel().0,
             ledger_manager: LedgerManager::spawn(Default::default()),
+            snark_worker: None,
             block_producer: None,
             archive: None,
             p2p: P2pServiceCtx::mocked(p2p_sec_key),

node/common/src/service/snark_worker.rs

@@ -1,40 +1,242 @@
+use ledger::proofs::provers::{TransactionProver, ZkappProver};
+use ledger::proofs::zkapp::ZkappParams;
+use ledger::scan_state::scan_state::transaction_snark::SokMessage;
 use mina_p2p_messages::v2;
-use node::external_snark_worker::{ExternalSnarkWorkerError, SnarkWorkSpec};
+use mina_signer::CompressedPubKey;
+use node::core::channels::mpsc;
+use node::event_source::ExternalSnarkWorkerEvent;
+use node::external_snark_worker::{
+    ExternalSnarkWorkerError, ExternalSnarkWorkerWorkError, SnarkWorkResult, SnarkWorkSpec,
+    SnarkWorkSpecError,
+};
+use node::snark::TransactionVerifier;
 
 use crate::NodeService;
 
-pub struct SnarkWorker {}
+use super::EventSender;
+
+pub struct SnarkWorker {
+    cmd_sender: mpsc::UnboundedSender<Cmd>,
+}
+
+enum Cmd {
+    Submit(Box<SnarkWorkSpec>),
+    Cancel,
+    Kill,
+}
 
 impl node::service::ExternalSnarkWorkerService for NodeService {
     fn start(
         &mut self,
-        _public_key: v2::NonZeroCurvePoint,
-        _fee: v2::CurrencyFeeStableV1,
+        pub_key: v2::NonZeroCurvePoint,
+        fee: v2::CurrencyFeeStableV1,
+        work_verifier: TransactionVerifier,
     ) -> Result<(), ExternalSnarkWorkerError> {
         if self.replayer.is_some() {
             return Ok(());
         }
-        todo!()
+        let (cmd_sender, cmd_receiver) = mpsc::unbounded_channel();
+        // TODO(binier): improve pub key conv
+        let sok_message = SokMessage::create(
+            (&fee).into(),
+            CompressedPubKey::from_address(&pub_key.to_string()).unwrap(),
+        );
+        self.snark_worker = Some(SnarkWorker { cmd_sender });
+        let event_sender = self.event_sender().clone();
+
+        node::core::thread::Builder::new()
+            .name("snark_worker".to_owned())
+            .spawn(move || worker_thread(cmd_receiver, event_sender, sok_message, work_verifier))
+            .map(|_| ())
+            .map_err(|err| ExternalSnarkWorkerError::Error(err.to_string()))
     }
 
     fn kill(&mut self) -> Result<(), ExternalSnarkWorkerError> {
         if self.replayer.is_some() {
             return Ok(());
         }
-        todo!()
+
+        if self
+            .snark_worker
+            .as_ref()
+            .and_then(|s| s.cmd_sender.send(Cmd::Kill).ok())
+            .is_none()
+        {
+            return Err(ExternalSnarkWorkerError::NotRunning);
+        }
+        Ok(())
     }
 
-    fn submit(&mut self, _spec: SnarkWorkSpec) -> Result<(), ExternalSnarkWorkerError> {
+    fn submit(&mut self, spec: SnarkWorkSpec) -> Result<(), ExternalSnarkWorkerError> {
         if self.replayer.is_some() {
             return Ok(());
         }
-        todo!()
+
+        if self
+            .snark_worker
+            .as_ref()
+            .and_then(|s| s.cmd_sender.send(Cmd::Submit(spec.into())).ok())
+            .is_none()
+        {
+            return Err(ExternalSnarkWorkerError::NotRunning);
+        }
+        Ok(())
     }
 
     fn cancel(&mut self) -> Result<(), ExternalSnarkWorkerError> {
         if self.replayer.is_some() {
             return Ok(());
         }
-        todo!()
+
+        // TODO(binier): for wasm threads, call terminate:
+        // https://developer.mozilla.org/en-US/docs/Web/API/Worker/terminate
+        if self
+            .snark_worker
+            .as_ref()
+            .and_then(|s| s.cmd_sender.send(Cmd::Cancel).ok())
+            .is_none()
+        {
+            return Err(ExternalSnarkWorkerError::NotRunning);
+        }
+        Ok(())
+    }
+}
+
+fn worker_thread(
+    mut cmd_receiver: mpsc::UnboundedReceiver<Cmd>,
+    event_sender: EventSender,
+    sok_message: SokMessage,
+    work_verifier: TransactionVerifier,
+) {
+    let _ = event_sender.send(ExternalSnarkWorkerEvent::Started.into());
+    let tx_prover = TransactionProver::make(Some(work_verifier.clone()));
+    let zkapp_prover = ZkappProver::make(Some(work_verifier));
+    while let Some(cmd) = cmd_receiver.blocking_recv() {
+        match cmd {
+            Cmd::Kill => {
+                let _ = event_sender.send(ExternalSnarkWorkerEvent::Killed.into());
+                return;
+            }
+            Cmd::Cancel => {
+                // can't cancel as it's a blocking thread. Once this
+                // is moved to another process, kill it.
+                let _ = event_sender.send(ExternalSnarkWorkerEvent::WorkCancelled.into());
+            }
+            Cmd::Submit(spec) => {
+                let event = match prove_spec(&tx_prover, &zkapp_prover, *spec, &sok_message) {
+                    Err(err) => ExternalSnarkWorkerEvent::WorkError(err),
+                    Ok(res) => ExternalSnarkWorkerEvent::WorkResult(res),
+                };
+
+                let _ = event_sender.send(event.into());
+            }
+        }
     }
 }
+
+fn prove_spec(
+    tx_prover: &TransactionProver,
+    zkapp_prover: &ZkappProver,
+    spec: SnarkWorkSpec,
+    sok_message: &SokMessage,
+) -> Result<SnarkWorkResult, ExternalSnarkWorkerWorkError> {
+    match spec {
+        SnarkWorkSpec::One(single) => prove_single(tx_prover, zkapp_prover, single, sok_message)
+            .map(v2::TransactionSnarkWorkTStableV2Proofs::One),
+        SnarkWorkSpec::Two((one, two)) => Ok(v2::TransactionSnarkWorkTStableV2Proofs::Two((
+            prove_single(tx_prover, zkapp_prover, one, sok_message)?,
+            prove_single(tx_prover, zkapp_prover, two, sok_message)?,
+        ))),
+    }
+    .map(Into::into)
+}
+
+fn invalid_bigint_err() -> ExternalSnarkWorkerWorkError {
+    ExternalSnarkWorkerWorkError::WorkSpecError(SnarkWorkSpecError::InvalidBigInt)
+}
+
+fn prove_single(
+    tx_prover: &TransactionProver,
+    zkapp_prover: &ZkappProver,
+    single: v2::SnarkWorkerWorkerRpcsVersionedGetWorkV2TResponseA0Single,
+    sok_message: &SokMessage,
+) -> Result<v2::LedgerProofProdStableV2, ExternalSnarkWorkerWorkError> {
+    use ledger::proofs::{merge::MergeParams, transaction::TransactionParams};
+
+    let (snarked_ledger_state, res) = match single {
+        v2::SnarkWorkerWorkerRpcsVersionedGetWorkV2TResponseA0Single::Transition(
+            snarked_ledger_state,
+            witness,
+        ) => {
+            if let v2::MinaTransactionTransactionStableV2::Command(cmd) = &witness.transaction {
+                if matches!(&**cmd, v2::MinaBaseUserCommandStableV2::ZkappCommand(_)) {
+                    return prove_zkapp(zkapp_prover, snarked_ledger_state, witness, sok_message);
+                }
+            }
+            let res = ledger::proofs::generate_tx_proof(TransactionParams {
+                statement: &snarked_ledger_state.0,
+                tx_witness: &witness,
+                message: sok_message,
+                tx_step_prover: &tx_prover.tx_step_prover,
+                tx_wrap_prover: &tx_prover.tx_wrap_prover,
+                only_verify_constraints: false,
+                expected_step_proof: None,
+                ocaml_wrap_witness: None,
+            });
+            (snarked_ledger_state.0, res)
+        }
+        v2::SnarkWorkerWorkerRpcsVersionedGetWorkV2TResponseA0Single::Merge(data) => {
+            let (snarked_ledger_state, proof_1, proof_2) = *data;
+            let res = ledger::proofs::generate_merge_proof(MergeParams {
+                statement: (&snarked_ledger_state.0)
+                    .try_into()
+                    .map_err(|_| invalid_bigint_err())?,
+                proofs: &[proof_1, proof_2],
+                message: sok_message,
+                step_prover: &tx_prover.merge_step_prover,
+                wrap_prover: &tx_prover.tx_wrap_prover,
+                only_verify_constraints: false,
+                expected_step_proof: None,
+                ocaml_wrap_witness: None,
+            });
+            (snarked_ledger_state.0, res)
+        }
+    };
+    res.map_err(|err| ExternalSnarkWorkerWorkError::Error(err.to_string()))
+        .map(|proof| {
+            v2::LedgerProofProdStableV2(v2::TransactionSnarkStableV2 {
+                statement: v2::MinaStateSnarkedLedgerStateWithSokStableV2 {
+                    source: snarked_ledger_state.source,
+                    target: snarked_ledger_state.target,
+                    connecting_ledger_left: snarked_ledger_state.connecting_ledger_left,
+                    connecting_ledger_right: snarked_ledger_state.connecting_ledger_right,
+                    supply_increase: snarked_ledger_state.supply_increase,
+                    fee_excess: snarked_ledger_state.fee_excess,
+                    sok_digest: (&sok_message.digest()).into(),
+                },
+                proof: v2::TransactionSnarkProofStableV2((&proof).into()),
+            })
+        })
+}
+
+fn prove_zkapp(
+    zkapp_prover: &ZkappProver,
+    snarked_ledger_state: v2::MinaStateSnarkedLedgerStateStableV2,
+    witness: v2::TransactionWitnessStableV2,
+    sok_message: &SokMessage,
+) -> Result<v2::LedgerProofProdStableV2, ExternalSnarkWorkerWorkError> {
+    ledger::proofs::generate_zkapp_proof(ZkappParams {
+        statement: &snarked_ledger_state.0,
+        tx_witness: &witness,
+        message: sok_message,
+        step_opt_signed_opt_signed_prover: &zkapp_prover.step_opt_signed_opt_signed_prover,
+        step_opt_signed_prover: &zkapp_prover.step_opt_signed_prover,
+        step_proof_prover: &zkapp_prover.step_proof_prover,
+        merge_step_prover: &zkapp_prover.merge_step_prover,
+        tx_wrap_prover: &zkapp_prover.tx_wrap_prover,
+        opt_signed_path: None,
+        proved_path: None,
+    })
+    .map(|proof| (&proof).into())
+    .map_err(|err| ExternalSnarkWorkerWorkError::Error(err.to_string()))
+}

node/src/event_source/event_source_effects.rs

@@ -97,7 +97,7 @@ pub fn event_source_effects<S: Service>(store: &mut Store<S>, action: EventSourc
                     MioEvent::IncomingDataDidReceive(addr, result) => {
                         store.dispatch(P2pNetworkSchedulerAction::IncomingDataDidReceive {
                             addr,
-                            result: result.map(From::from),
+                            result,
                         });
                     }
                     MioEvent::OutgoingDataDidSend(_, _result) => {}

node/src/external_snark_worker_effectful/external_snark_worker_effectful_effects.rs

@@ -9,7 +9,8 @@ pub fn external_snark_worker_effectful_effects<S: crate::Service>(
     let (action, _) = action.split();
     match action {
         ExternalSnarkWorkerEffectfulAction::Start { public_key, fee } => {
-            if let Err(err) = store.service.start(public_key, fee) {
+            let work_verifier = store.state().snark.work_verify.verifier_index.clone();
+            if let Err(err) = store.service.start(public_key, fee, work_verifier) {
                 store.dispatch(ExternalSnarkWorkerAction::Error {
                     error: err,
                     permanent: true,

node/src/external_snark_worker_effectful/external_snark_worker_service.rs

@@ -1,5 +1,6 @@
 use mina_p2p_messages::v2::{CurrencyFeeStableV1, NonZeroCurvePoint};
 use serde::{Deserialize, Serialize};
+use snark::TransactionVerifier;
 
 use crate::external_snark_worker::{
     ExternalSnarkWorkerError, ExternalSnarkWorkerWorkError, SnarkWorkResult, SnarkWorkSpec,
@@ -21,6 +22,7 @@ pub trait ExternalSnarkWorkerService {
         &mut self,
         public_key: NonZeroCurvePoint,
         fee: CurrencyFeeStableV1,
+        work_verifier: TransactionVerifier,
     ) -> Result<(), ExternalSnarkWorkerError>;
 
     /// Submits snark work

node/src/recorder/recorder.rs

@@ -160,7 +160,7 @@ fn graceful_shutdown(only_i: Option<usize>) {
     let files_iter = files
         .iter_mut()
         .enumerate()
-        .filter(|(i, _)| only_i.map_or(true, |only_i| only_i == *i))
+        .filter(|(i, _)| only_i.is_none_or(|only_i| only_i == *i))
         .filter_map(|(i, v)| Some((i, v.take()?)));
 
     for (i, file) in files_iter {

node/src/snark_pool/candidate/snark_pool_candidate_actions.rs

@@ -70,7 +70,7 @@ impl redux::EnablingCondition<crate::State> for SnarkPoolCandidateAction {
                         .snark_pool
                         .candidates
                         .get(*peer_id, &info.job_id)
-                        .map_or(true, |v| info > v)
+                        .is_none_or(|v| info > v)
             }
             SnarkPoolCandidateAction::WorkFetchAll => state.p2p.ready().is_some(),
             SnarkPoolCandidateAction::WorkFetchInit { peer_id, job_id } => {
@@ -104,7 +104,7 @@ impl redux::EnablingCondition<crate::State> for SnarkPoolCandidateAction {
                         .snark_pool
                         .candidates
                         .get(*peer_id, &job_id)
-                        .map_or(true, |v| match work.partial_cmp(v).unwrap() {
+                        .is_none_or(|v| match work.partial_cmp(v).unwrap() {
                             Ordering::Less => false,
                             Ordering::Greater => true,
                             Ordering::Equal => {