Merge pull request #983 from openmina/faster-vrf

sebastiencs · web-flow · commit 54456579df30 · 2024-12-11T21:26:40.000+01:00
Allocate `num-bigint` on the stack
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -87,6 +87,9 @@ ark-ec = { git = "https://github.com/openmina/algebra", rev = "d0343f5" } # bran
 ark-poly = { git = "https://github.com/openmina/algebra", rev = "d0343f5" } # branch: fix-openmina-webnode
 ark-serialize = { git = "https://github.com/openmina/algebra", rev = "d0343f5" } # branch: fix-openmina-webnode
 
+num-bigint = { git = "https://github.com/openmina/num-bigint", rev = "8bb5ee4" } # branch: on-stack
+num-rational = { git = "https://github.com/openmina/num-rational", rev = "336f11d" } # branch: on-stack
+
 [profile.test.package."*"]
 opt-level = 3
 debug-assertions = true
diff --git a/mina-p2p-messages/src/bigint.rs b/mina-p2p-messages/src/bigint.rs
@@ -47,7 +47,7 @@ impl BigInt {
     }
 
     pub fn from_decimal(s: &str) -> Result<Self, InvalidDecimalNumber> {
-        num_bigint::BigInt::parse_bytes(s.as_bytes(), 10)
+        num_bigint::BigInt::<4>::parse_bytes(s.as_bytes(), 10)
             .map(|num| {
                 let mut bytes = num.to_bytes_be().1;
                 bytes.reverse();
@@ -238,7 +238,7 @@ impl<'de> Deserialize<'de> for BigInt {
                         }),
                         None => {
                             // Try to parse as a decimal number
-                            num_bigint::BigInt::parse_bytes(v.as_bytes(), 10)
+                            num_bigint::BigInt::<4>::parse_bytes(v.as_bytes(), 10)
                                 .map(|num| {
                                     let mut bytes = num.to_bytes_be().1;
                                     bytes.reverse();
diff --git a/node/common/src/service/block_producer/vrf_evaluator.rs b/node/common/src/service/block_producer/vrf_evaluator.rs
@@ -8,6 +8,7 @@ use node::{
     core::channels::mpsc::{UnboundedReceiver, UnboundedSender},
     event_source::Event,
 };
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use vrf::{VrfEvaluationInput, VrfEvaluationOutput};
 
 use crate::NodeService;
@@ -18,29 +19,46 @@ pub fn vrf_evaluator(
     keypair: Keypair,
 ) {
     while let Some(vrf_evaluator_input) = vrf_evaluation_receiver.blocking_recv() {
-        let mut vrf_result = VrfEvaluationOutput::SlotLost(vrf_evaluator_input.global_slot);
-
-        for (index, (pub_key, stake)) in vrf_evaluator_input.delegator_table.iter() {
-            let vrf_input = VrfEvaluationInput::new(
-                keypair.clone(),
-                vrf_evaluator_input.epoch_seed.clone(),
-                pub_key.clone(),
-                vrf_evaluator_input.global_slot,
-                *index,
-                (*stake).into(),
-                vrf_evaluator_input.total_currency.into(),
-            );
-            vrf_result = vrf::evaluate_vrf(vrf_input).unwrap();
-
-            // the first delegate that won the slot
-            if let VrfEvaluationOutput::SlotWon(_) = vrf_result {
-                break;
-            }
-        }
+        // let bytes = serde_json::to_string(&vrf_evaluator_input).unwrap();
+        // openmina_core::http::download("vrf.json".to_string(), bytes.as_bytes().to_vec()).unwrap();
+
+        let keypair = &keypair;
+        let VrfEvaluatorInput {
+            epoch_seed,
+            delegator_table,
+            global_slot,
+            total_currency,
+            staking_ledger_hash: _,
+        } = &vrf_evaluator_input;
+
+        let vrf_result = delegator_table
+            .par_iter()
+            .find_map_any(|(index, (pub_key, stake))| {
+                let vrf_input = VrfEvaluationInput {
+                    producer_key: keypair.clone(),
+                    global_slot: *global_slot,
+                    epoch_seed: epoch_seed.clone(),
+                    account_pub_key: pub_key.clone(),
+                    delegator_index: *index,
+                    delegated_stake: (*stake).into(),
+                    total_currency: (*total_currency).into(),
+                };
+
+                let vrf_result = vrf::evaluate_vrf(vrf_input).unwrap();
+
+                // the first delegate that won the slot
+                if let VrfEvaluationOutput::SlotWon(_) = vrf_result {
+                    return Some(vrf_result);
+                }
+                None
+            })
+            .unwrap_or(VrfEvaluationOutput::SlotLost(*global_slot));
+
         let vrf_result_with_hash = VrfEvaluationOutputWithHash::new(
             vrf_result,
             vrf_evaluator_input.staking_ledger_hash.clone(),
         );
+
         // send the result back to the state machine
         let _ = event_sender.send(
             BlockProducerEvent::VrfEvaluator(BlockProducerVrfEvaluatorEvent::Evaluated(
@@ -60,3 +78,70 @@ impl node::block_producer_effectful::vrf_evaluator_effectful::BlockProducerVrfEv
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    // use mina_signer::keypair;
+    use node::account::AccountSecretKey;
+
+    use super::*;
+
+    #[test]
+    #[ignore]
+    fn test_vrf() {
+        let json = std::fs::read_to_string("/tmp/vrf.json").unwrap();
+        let vrf_evaluator_input: VrfEvaluatorInput = serde_json::from_str(&json).unwrap();
+
+        let private = "SOME_KEY";
+        let private = AccountSecretKey::from_str(private).unwrap();
+        let keypair: Keypair = private.into();
+
+        let VrfEvaluatorInput {
+            epoch_seed,
+            delegator_table,
+            global_slot,
+            total_currency,
+            staking_ledger_hash: _,
+        } = &vrf_evaluator_input;
+
+        let now = std::time::Instant::now();
+
+        let vrf_result = delegator_table
+            .par_iter()
+            .map(|(index, (pub_key, stake))| {
+                let vrf_input = VrfEvaluationInput {
+                    producer_key: keypair.clone(),
+                    global_slot: *global_slot,
+                    epoch_seed: epoch_seed.clone(),
+                    account_pub_key: pub_key.clone(),
+                    delegator_index: *index,
+                    delegated_stake: (*stake).into(),
+                    total_currency: (*total_currency).into(),
+                };
+                // let now = redux::Instant::now();
+                let vrf_result = vrf::evaluate_vrf(vrf_input).unwrap();
+                // let elapsed = now.elapsed();
+                // let slot = global_slot;
+                // eprintln!("vrf::evaluate_vrf: {elapsed:?} slot:{slot:?} index:{index:?}");
+                // openmina_core::info!(openmina_core::log::system_time(); "vrf::evaluate_vrf: {elapsed:?} slot:{slot:?} index:{index:?}");
+
+                // nevaluated.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
+
+                // the first delegate that won the slot
+                if let VrfEvaluationOutput::SlotWon(_) = vrf_result {
+                    return Some(vrf_result);
+                }
+                None
+            })
+            .collect::<Vec<_>>();
+
+        let elapsed = now.elapsed();
+        let slot = vrf_evaluator_input.global_slot;
+        let ndelegator = vrf_evaluator_input.delegator_table.len();
+        // let nevaluated = nevaluated.load(std::sync::atomic::Ordering::Relaxed);
+        eprintln!("TOTAL vrf::evaluate_vrf: {elapsed:?} slot:{slot:?} ndelegators:{ndelegator:?}");
+        dbg!(vrf_result);
+    }
+}
diff --git a/producer-dashboard/src/evaluator/mod.rs b/producer-dashboard/src/evaluator/mod.rs
@@ -51,15 +51,15 @@ impl Evaluator {
                     .db
                     .store_slot(global_slot, &SlotData::new_lost(global_slot, timestamp));
                 for (index, delegate) in &delegates {
-                    let vrf_input = VrfEvaluationInput::new(
-                        self.key.clone().into(),
-                        epoch_seed.clone(),
-                        pub_key.clone(),
+                    let vrf_input = VrfEvaluationInput {
+                        producer_key: self.key.clone().into(),
                         global_slot,
-                        (*index).into(),
-                        delegate.balance.clone().into(),
-                        total_currency.clone(),
-                    );
+                        epoch_seed: epoch_seed.clone(),
+                        account_pub_key: pub_key.clone(),
+                        delegator_index: (*index).into(),
+                        delegated_stake: delegate.balance.clone().into(),
+                        total_currency: total_currency.clone(),
+                    };
 
                     if let Ok(VrfEvaluationOutput::SlotWon(_)) = vrf::evaluate_vrf(vrf_input) {
                         println!("Won slot: {global_slot}");
diff --git a/vrf/Cargo.toml b/vrf/Cargo.toml
@@ -22,7 +22,8 @@ poseidon = { workspace = true }
 ark-ff = { workspace = true }
 ark-ec = { version = "0.3.0", features = [ "std" ] }
 ark-serialize = { version = "0.3.0", features = [ "std" ] }
-num = "0.4"
+num = {version = "0.4", features = [ "std" ] }
+# num = "0.4"
 itertools = "0.12"
 sha2 = "0.10"
 bs58 = "0.4.0"
diff --git a/vrf/src/lib.rs b/vrf/src/lib.rs
@@ -3,7 +3,7 @@ use ark_ff::PrimeField;
 use ledger::AccountIndex;
 use message::VrfMessage;
 use mina_p2p_messages::v2::EpochSeed;
-use num::{BigInt, ToPrimitive};
+use num::{rational::Ratio, BigInt, ToPrimitive};
 use openmina_node_account::AccountPublicKey;
 use output::VrfOutput;
 use serde::{Deserialize, Serialize};
@@ -52,6 +52,17 @@ pub enum VrfError {
     IvalidWitness,
 }
 
+/// 256 bits
+pub(crate) type BigInt256 = BigInt<4>;
+
+/// 2048 bits
+pub(crate) type BigInt2048 = BigInt<32>;
+pub(crate) type BigRational2048 = Ratio<BigInt2048>;
+
+/// 4096 bits
+pub(crate) type BigInt4096 = BigInt<64>;
+pub(crate) type BigRational4096 = Ratio<BigInt4096>;
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct VrfWonSlot {
     pub producer: AccountPublicKey,
@@ -79,35 +90,13 @@ impl VrfEvaluationOutput {
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct VrfEvaluationInput {
-    producer_key: Keypair,
-    global_slot: u32,
-    epoch_seed: EpochSeed,
-    account_pub_key: AccountPublicKey,
-    delegator_index: AccountIndex,
-    delegated_stake: BigInt,
-    total_currency: BigInt,
-}
-
-impl VrfEvaluationInput {
-    pub fn new(
-        producer_key: Keypair,
-        epoch_seed: EpochSeed,
-        account_pub_key: AccountPublicKey,
-        global_slot: u32,
-        delegator_index: AccountIndex,
-        delegated_stake: BigInt,
-        total_currency: BigInt,
-    ) -> Self {
-        Self {
-            producer_key,
-            global_slot,
-            epoch_seed,
-            delegator_index,
-            delegated_stake,
-            total_currency,
-            account_pub_key,
-        }
-    }
+    pub producer_key: Keypair,
+    pub global_slot: u32,
+    pub epoch_seed: EpochSeed,
+    pub account_pub_key: AccountPublicKey,
+    pub delegator_index: AccountIndex,
+    pub delegated_stake: BigInt,
+    pub total_currency: BigInt,
 }
 
 /// Generates the VRF output for the genesis block
diff --git a/vrf/src/output.rs b/vrf/src/output.rs
@@ -1,5 +1,6 @@
 use ark_ec::short_weierstrass_jacobian::GroupAffine;
 use ark_ff::{BigInteger, BigInteger256, PrimeField};
+use ledger::proofs::transaction::field_to_bits;
 use ledger::{AppendToInputs, ToInputs};
 use mina_p2p_messages::v2::ConsensusVrfOutputTruncatedStableV1;
 use num::{BigInt, BigRational, One, ToPrimitive};
@@ -8,7 +9,7 @@ use poseidon::hash::params::MINA_VRF_OUTPUT;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 
-use crate::{BaseField, ScalarField};
+use crate::{BaseField, BigInt2048, ScalarField};
 
 use super::serialize::{ark_deserialize, ark_serialize};
 
@@ -62,7 +63,8 @@ impl VrfOutput {
     }
 
     pub fn truncated(&self) -> ScalarField {
-        let bits = self.hash().to_bits();
+        let hash = self.hash();
+        let bits = field_to_bits::<_, 255>(hash);
 
         let repr = BigInteger256::from_bits_le(&bits[..bits.len() - 3]);
         ScalarField::from_repr(repr).unwrap()
@@ -89,7 +91,7 @@ impl VrfOutput {
         //                 Field.size_in_bits = 255
         let two_tpo_256 = BigInt::one() << 253u32;
 
-        let vrf_out = BigInt::from_bytes_be(
+        let vrf_out: BigInt2048 = BigInt2048::from_bytes_be(
             num::bigint::Sign::Plus,
             &self.truncated().into_repr().to_bytes_be(),
         );
diff --git a/vrf/src/threshold.rs b/vrf/src/threshold.rs
