Skip to content

Commit 30c6517

Browse files
AlenarAlenar
authored
Merge pull request #661 from input-output-hk/djo/simplify_epoch_offset
Simplify epoch offset computations
2 parents f3fd5aa + cbfadfd commit 30c6517

File tree

10 files changed

+138
-261
lines changed

10 files changed

+138
-261
lines changed

‎mithril-aggregator/src/command_args.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ async fn do_first_launch_initialization_if_needed(
110110
Epoch(0) => (Epoch(0), Epoch(1)),
111111
epoch => (
112112
epoch.offset_to_signer_retrieval_epoch()?,
113-
epoch.offset_to_next_signer_retrieval_epoch()?,
113+
epoch.offset_to_next_signer_retrieval_epoch(),
114114
),
115115
};
116116

‎mithril-aggregator/src/dependency.rs‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,9 +104,7 @@ impl DependencyManager {
104104
let work_epoch = current_epoch
105105
.offset_to_signer_retrieval_epoch()
106106
.expect("epoch.offset_by SIGNER_EPOCH_RETRIEVAL_OFFSET should not fail");
107-
let epoch_to_sign = current_epoch
108-
.offset_to_next_signer_retrieval_epoch()
109-
.expect("epoch.offset_by NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET should not fail");
107+
let epoch_to_sign = current_epoch.offset_to_next_signer_retrieval_epoch();
110108

111109
(work_epoch, epoch_to_sign)
112110
}

‎mithril-aggregator/src/multi_signer.rs‎

Lines changed: 82 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,23 @@
1-
use std::collections::HashMap;
2-
use std::sync::Arc;
3-
41
use async_trait::async_trait;
52
use chrono::prelude::*;
63
use hex::ToHex;
74
use slog_scope::{debug, trace, warn};
5+
use std::{collections::HashMap, sync::Arc};
86
use thiserror::Error;
97

10-
use mithril_common::crypto_helper::{
11-
key_decode_hex, key_encode_hex, ProtocolAggregateVerificationKey, ProtocolAggregationError,
12-
ProtocolClerk, ProtocolKeyRegistration, ProtocolMultiSignature, ProtocolParameters,
13-
ProtocolPartyId, ProtocolRegistrationError, ProtocolSignerVerificationKey,
14-
ProtocolSingleSignature, ProtocolStakeDistribution,
15-
};
16-
use mithril_common::entities::{self, SignerWithStake};
17-
use mithril_common::store::{StakeStore, StakeStorer, StoreError};
188
use mithril_common::{
19-
NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET, SIGNER_EPOCH_RECORDING_OFFSET,
20-
SIGNER_EPOCH_RETRIEVAL_OFFSET,
9+
crypto_helper::{
10+
key_decode_hex, key_encode_hex, ProtocolAggregateVerificationKey, ProtocolAggregationError,
11+
ProtocolClerk, ProtocolKeyRegistration, ProtocolMultiSignature, ProtocolParameters,
12+
ProtocolPartyId, ProtocolRegistrationError, ProtocolSignerVerificationKey,
13+
ProtocolSingleSignature, ProtocolStakeDistribution,
14+
},
15+
entities::{self, Epoch, SignerWithStake},
16+
store::{StakeStore, StakeStorer, StoreError},
2117
};
2218

23-
use crate::store::{SingleSignatureStorer, VerificationKeyStorer};
2419
use crate::{
20+
store::{SingleSignatureStorer, VerificationKeyStorer},
2521
ProtocolParametersStore, ProtocolParametersStorer, SingleSignatureStore, VerificationKeyStore,
2622
};
2723

@@ -191,13 +187,10 @@ pub trait MultiSigner: Sync + Send {
191187
}
192188

193189
/// Get signers with stake
194-
async fn get_signers_with_stake(&self)
195-
-> Result<Vec<entities::SignerWithStake>, ProtocolError>;
190+
async fn get_signers_with_stake(&self) -> Result<Vec<SignerWithStake>, ProtocolError>;
196191

197192
/// Get signers for the next epoch with their stake
198-
async fn get_next_signers_with_stake(
199-
&self,
200-
) -> Result<Vec<entities::SignerWithStake>, ProtocolError>;
193+
async fn get_next_signers_with_stake(&self) -> Result<Vec<SignerWithStake>, ProtocolError>;
201194

202195
/// Registers a single signature
203196
async fn register_single_signature(
@@ -322,18 +315,12 @@ impl MultiSignerImpl {
322315
}
323316
}
324317

325-
/// Get stake distribution with epoch offset
326-
async fn get_stake_distribution_with_epoch_offset(
318+
/// Get the [stake distribution][ProtocolStakeDistribution] for the given `epoch`
319+
async fn get_stake_distribution_at_epoch(
327320
&self,
328-
epoch_offset: i64,
321+
epoch: Epoch,
329322
) -> Result<ProtocolStakeDistribution, ProtocolError> {
330-
debug!("Get stake distribution with epoch offset"; "epoch_offset"=>epoch_offset);
331-
let epoch = self
332-
.current_beacon
333-
.as_ref()
334-
.ok_or_else(ProtocolError::UnavailableBeacon)?
335-
.epoch
336-
.offset_by(epoch_offset)?;
323+
debug!("Get stake distribution at epoch"; "epoch"=> #?epoch);
337324

338325
let stakes = self
339326
.stake_store
@@ -343,18 +330,12 @@ impl MultiSignerImpl {
343330
Ok(stakes.into_iter().collect::<ProtocolStakeDistribution>())
344331
}
345332

346-
/// Get protocol parameters with epoch offset
347-
async fn get_protocol_parameters_with_epoch_offset(
333+
/// Get the [protocol parameters][ProtocolParameters] for the given `epoch`
334+
async fn get_protocol_parameters_at_epoch(
348335
&self,
349-
epoch_offset: i64,
336+
epoch: Epoch,
350337
) -> Result<Option<ProtocolParameters>, ProtocolError> {
351-
debug!("Get protocol parameters with epoch offset"; "epoch_offset"=>epoch_offset);
352-
let epoch = self
353-
.current_beacon
354-
.as_ref()
355-
.ok_or_else(ProtocolError::UnavailableBeacon)?
356-
.epoch
357-
.offset_by(epoch_offset)?;
338+
debug!("Get protocol parameters at epoch"; "epoch"=> #?epoch);
358339

359340
match self
360341
.protocol_parameters_store
@@ -431,8 +412,13 @@ impl MultiSigner for MultiSignerImpl {
431412
/// Get protocol parameters
432413
async fn get_protocol_parameters(&self) -> Result<Option<ProtocolParameters>, ProtocolError> {
433414
debug!("Get protocol parameters");
434-
self.get_protocol_parameters_with_epoch_offset(SIGNER_EPOCH_RETRIEVAL_OFFSET)
435-
.await
415+
let epoch = self
416+
.current_beacon
417+
.as_ref()
418+
.ok_or_else(ProtocolError::UnavailableBeacon)?
419+
.epoch
420+
.offset_to_signer_retrieval_epoch()?;
421+
self.get_protocol_parameters_at_epoch(epoch).await
436422
}
437423

438424
/// Update protocol parameters
@@ -446,7 +432,7 @@ impl MultiSigner for MultiSignerImpl {
446432
.as_ref()
447433
.ok_or_else(ProtocolError::UnavailableBeacon)?
448434
.epoch
449-
.offset_by(SIGNER_EPOCH_RECORDING_OFFSET)?;
435+
.offset_to_recording_epoch();
450436

451437
self.protocol_parameters_store
452438
.save_protocol_parameters(epoch, protocol_parameters.to_owned().into())
@@ -459,24 +445,39 @@ impl MultiSigner for MultiSignerImpl {
459445
&self,
460446
) -> Result<Option<ProtocolParameters>, ProtocolError> {
461447
debug!("Get next protocol parameters");
462-
self.get_protocol_parameters_with_epoch_offset(NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET)
463-
.await
448+
let epoch = self
449+
.current_beacon
450+
.as_ref()
451+
.ok_or_else(ProtocolError::UnavailableBeacon)?
452+
.epoch
453+
.offset_to_next_signer_retrieval_epoch();
454+
self.get_protocol_parameters_at_epoch(epoch).await
464455
}
465456

466457
/// Get stake distribution
467458
async fn get_stake_distribution(&self) -> Result<ProtocolStakeDistribution, ProtocolError> {
468459
debug!("Get stake distribution");
469-
self.get_stake_distribution_with_epoch_offset(SIGNER_EPOCH_RETRIEVAL_OFFSET)
470-
.await
460+
let epoch = self
461+
.current_beacon
462+
.as_ref()
463+
.ok_or_else(ProtocolError::UnavailableBeacon)?
464+
.epoch
465+
.offset_to_signer_retrieval_epoch()?;
466+
self.get_stake_distribution_at_epoch(epoch).await
471467
}
472468

473469
/// Get next stake distribution
474470
async fn get_next_stake_distribution(
475471
&self,
476472
) -> Result<ProtocolStakeDistribution, ProtocolError> {
477473
debug!("Get next stake distribution");
478-
self.get_stake_distribution_with_epoch_offset(NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET)
479-
.await
474+
let epoch = self
475+
.current_beacon
476+
.as_ref()
477+
.ok_or_else(ProtocolError::UnavailableBeacon)?
478+
.epoch
479+
.offset_to_next_signer_retrieval_epoch();
480+
self.get_stake_distribution_at_epoch(epoch).await
480481
}
481482

482483
/// Update stake distribution
@@ -490,7 +491,7 @@ impl MultiSigner for MultiSignerImpl {
490491
.as_ref()
491492
.ok_or_else(ProtocolError::UnavailableBeacon)?
492493
.epoch
493-
.offset_by(SIGNER_EPOCH_RECORDING_OFFSET)?;
494+
.offset_to_recording_epoch();
494495
let stakes = HashMap::from_iter(stakes.iter().cloned());
495496
self.stake_store.save_stakes(epoch, stakes).await?;
496497

@@ -523,7 +524,7 @@ impl MultiSigner for MultiSignerImpl {
523524
.as_ref()
524525
.ok_or_else(ProtocolError::UnavailableBeacon)?
525526
.epoch
526-
.offset_by(SIGNER_EPOCH_RETRIEVAL_OFFSET)?;
527+
.offset_to_signer_retrieval_epoch()?;
527528
let signers = self
528529
.verification_key_store
529530
.get_verification_keys(epoch)
@@ -537,16 +538,14 @@ impl MultiSigner for MultiSignerImpl {
537538
}
538539
}
539540

540-
async fn get_signers_with_stake(
541-
&self,
542-
) -> Result<Vec<entities::SignerWithStake>, ProtocolError> {
541+
async fn get_signers_with_stake(&self) -> Result<Vec<SignerWithStake>, ProtocolError> {
543542
debug!("Get signers with stake");
544543
let epoch = self
545544
.current_beacon
546545
.as_ref()
547546
.ok_or_else(ProtocolError::UnavailableBeacon)?
548547
.epoch
549-
.offset_by(SIGNER_EPOCH_RETRIEVAL_OFFSET)?;
548+
.offset_to_signer_retrieval_epoch()?;
550549
let signers = self
551550
.verification_key_store
552551
.get_verification_keys(epoch)
@@ -558,7 +557,7 @@ impl MultiSigner for MultiSignerImpl {
558557
.iter()
559558
.filter_map(|(party_id, stake)| {
560559
signers.get(party_id).map(|signer| {
561-
entities::SignerWithStake::new(
560+
SignerWithStake::new(
562561
party_id.to_owned(),
563562
signer.verification_key.to_owned(),
564563
signer.verification_key_signature.to_owned(),
@@ -571,16 +570,14 @@ impl MultiSigner for MultiSignerImpl {
571570
.collect())
572571
}
573572

574-
async fn get_next_signers_with_stake(
575-
&self,
576-
) -> Result<Vec<entities::SignerWithStake>, ProtocolError> {
573+
async fn get_next_signers_with_stake(&self) -> Result<Vec<SignerWithStake>, ProtocolError> {
577574
debug!("Get next signers with stake");
578575
let epoch = self
579576
.current_beacon
580577
.as_ref()
581578
.ok_or_else(ProtocolError::UnavailableBeacon)?
582579
.epoch
583-
.offset_by(NEXT_SIGNER_EPOCH_RETRIEVAL_OFFSET)?;
580+
.offset_to_next_signer_retrieval_epoch();
584581
let signers = self
585582
.verification_key_store
586583
.get_verification_keys(epoch)
@@ -592,7 +589,7 @@ impl MultiSigner for MultiSignerImpl {
592589
.iter()
593590
.filter_map(|(party_id, stake)| {
594591
signers.get(party_id).map(|signer| {
595-
entities::SignerWithStake::new(
592+
SignerWithStake::new(
596593
party_id.to_owned(),
597594
signer.verification_key.to_owned(),
598595
signer.verification_key_signature.to_owned(),
@@ -724,31 +721,30 @@ impl MultiSigner for MultiSignerImpl {
724721
#[cfg(test)]
725722
mod tests {
726723
use super::*;
727-
use crate::store::{SingleSignatureStore, VerificationKeyStore};
728-
use crate::ProtocolParametersStore;
729-
730-
use mithril_common::crypto_helper::tests_setup::*;
731-
use mithril_common::fake_data;
732-
use mithril_common::store::adapter::MemoryAdapter;
733-
use mithril_common::store::StakeStore;
734-
735-
use std::collections::HashMap;
736-
use std::sync::Arc;
724+
use crate::{
725+
store::{SingleSignatureStore, VerificationKeyStore},
726+
ProtocolParametersStore,
727+
};
728+
use mithril_common::{
729+
crypto_helper::tests_setup::*,
730+
fake_data,
731+
store::{adapter::MemoryAdapter, StakeStore},
732+
};
733+
use std::{collections::HashMap, sync::Arc};
737734

738735
async fn setup_multi_signer() -> MultiSignerImpl {
739736
let beacon = fake_data::beacon();
740-
let verification_key_store = VerificationKeyStore::new(Box::new(
741-
MemoryAdapter::<entities::Epoch, HashMap<entities::PartyId, entities::Signer>>::new(
742-
None,
743-
)
744-
.unwrap(),
745-
), None);
737+
let verification_key_store = VerificationKeyStore::new(
738+
Box::new(
739+
MemoryAdapter::<Epoch, HashMap<entities::PartyId, entities::Signer>>::new(None)
740+
.unwrap(),
741+
),
742+
None,
743+
);
746744
let stake_store = StakeStore::new(
747745
Box::new(
748-
MemoryAdapter::<entities::Epoch, HashMap<entities::PartyId, entities::Stake>>::new(
749-
None,
750-
)
751-
.unwrap(),
746+
MemoryAdapter::<Epoch, HashMap<entities::PartyId, entities::Stake>>::new(None)
747+
.unwrap(),
752748
),
753749
None,
754750
);
@@ -764,16 +760,13 @@ mod tests {
764760
);
765761
let protocol_parameters_store = ProtocolParametersStore::new(
766762
Box::new(
767-
MemoryAdapter::<entities::Epoch, entities::ProtocolParameters>::new(Some(vec![
763+
MemoryAdapter::<Epoch, entities::ProtocolParameters>::new(Some(vec![
768764
(
769765
beacon.epoch.offset_to_signer_retrieval_epoch().unwrap(),
770766
fake_data::protocol_parameters(),
771767
),
772768
(
773-
beacon
774-
.epoch
775-
.offset_to_next_signer_retrieval_epoch()
776-
.unwrap(),
769+
beacon.epoch.offset_to_next_signer_retrieval_epoch(),
777770
fake_data::protocol_parameters(),
778771
),
779772
]))
@@ -851,7 +844,7 @@ mod tests {
851844

852845
offset_epoch(
853846
&mut multi_signer,
854-
SIGNER_EPOCH_RECORDING_OFFSET - SIGNER_EPOCH_RETRIEVAL_OFFSET,
847+
Epoch::SIGNER_RECORDING_OFFSET as i64 - Epoch::SIGNER_RETRIEVAL_OFFSET,
855848
)
856849
.await;
857850

@@ -887,7 +880,7 @@ mod tests {
887880

888881
offset_epoch(
889882
&mut multi_signer,
890-
SIGNER_EPOCH_RECORDING_OFFSET - SIGNER_EPOCH_RETRIEVAL_OFFSET,
883+
Epoch::SIGNER_RECORDING_OFFSET as i64 - Epoch::SIGNER_RETRIEVAL_OFFSET,
891884
)
892885
.await;
893886

@@ -943,7 +936,7 @@ mod tests {
943936
for (signer_with_stake, _, _) in &signers {
944937
verification_key_store
945938
.save_verification_key(
946-
start_epoch.offset_to_recording_epoch().unwrap(),
939+
start_epoch.offset_to_recording_epoch(),
947940
signer_with_stake.to_owned().into(),
948941
)
949942
.await
@@ -952,7 +945,7 @@ mod tests {
952945

953946
offset_epoch(
954947
&mut multi_signer,
955-
SIGNER_EPOCH_RECORDING_OFFSET - SIGNER_EPOCH_RETRIEVAL_OFFSET,
948+
Epoch::SIGNER_RECORDING_OFFSET as i64 - Epoch::SIGNER_RETRIEVAL_OFFSET,
956949
)
957950
.await;
958951
// We have to update the current message AFTER we reached the epoch for

0 commit comments

Comments
 (0)