From 120fc306284eb84be75522a7ea421a89ec4eb700 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 17 Feb 2021 14:39:27 +0100
Subject: [PATCH 01/19] Make changes

---
 frame/babe/src/lib.rs                | 32 ++++++++++++++++++++++------
 primitives/consensus/babe/src/lib.rs | 10 ++++++++-
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index b42b4f177ff62..e63564ea76ba9 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -42,7 +42,8 @@ use sp_timestamp::OnTimestampSet;
 
 use sp_consensus_babe::{
 	digests::{NextConfigDescriptor, NextEpochDescriptor, PreDigest},
-	BabeAuthorityWeight, ConsensusLog, Epoch, EquivocationProof, Slot, BABE_ENGINE_ID,
+	BabeAuthorityWeight, BabeEpochConfiguration, ConsensusLog, Epoch,
+	EquivocationProof, Slot, BABE_ENGINE_ID,
 };
 use sp_consensus_vrf::schnorrkel;
 
@@ -187,7 +188,7 @@ decl_storage! {
 		pub Randomness get(fn randomness): schnorrkel::Randomness;
 
 		/// Next epoch configuration, if changed.
-		NextEpochConfig: Option<NextConfigDescriptor>;
+		PendingEpochConfigChange: Option<NextConfigDescriptor>;
 
 		/// Next epoch randomness.
 		NextRandomness: schnorrkel::Randomness;
@@ -224,10 +225,16 @@ decl_storage! {
 		/// on block finalization. Querying this storage entry outside of block
 		/// execution context should always yield zero.
 		Lateness get(fn lateness): T::BlockNumber;
+
+		EpochConfig: BabeEpochConfiguration;
+		NextEpochConfig: Option<BabeEpochConfiguration>;
 	}
 	add_extra_genesis {
 		config(authorities): Vec<(AuthorityId, BabeAuthorityWeight)>;
-		build(|config| Module::<T>::initialize_authorities(&config.authorities))
+		config(epoch_config): BabeEpochConfiguration;
+		build(|config| Module::<T>::initialize_authorities_and_epoch_config(
+			&config.authorities, &config.epoch_config
+		))
 	}
 }
 
@@ -436,7 +443,7 @@ impl<T: Config> Module<T> {
 	pub fn plan_config_change(
 		config: NextConfigDescriptor,
 	) {
-		NextEpochConfig::put(config);
+		PendingEpochConfigChange::put(config);
 	}
 
 	/// DANGEROUS: Enact an epoch change. Should be done on every block where `should_epoch_change` has returned `true`,
@@ -483,8 +490,16 @@ impl<T: Config> Module<T> {
 		};
 		Self::deposit_consensus(ConsensusLog::NextEpochData(next_epoch));
 
-		if let Some(next_config) = NextEpochConfig::take() {
-			Self::deposit_consensus(ConsensusLog::NextConfigData(next_config));
+		if let Some(pending_epoch_config_change) = PendingEpochConfigChange::take() {
+			if let Some(next_config) = NextEpochConfig::get() {
+				EpochConfig::put(next_config);
+			}
+
+			let next_epoch_config: BabeEpochConfiguration =
+				pending_epoch_config_change.clone().into();
+			NextEpochConfig::put(next_epoch_config);
+
+			Self::deposit_consensus(ConsensusLog::NextConfigData(pending_epoch_config_change));
 		}
 	}
 
@@ -503,6 +518,7 @@ impl<T: Config> Module<T> {
 			duration: T::EpochDuration::get(),
 			authorities: Self::authorities(),
 			randomness: Self::randomness(),
+			config: EpochConfig::get(),
 		}
 	}
 
@@ -520,6 +536,7 @@ impl<T: Config> Module<T> {
 			duration: T::EpochDuration::get(),
 			authorities: NextAuthorities::get(),
 			randomness: NextRandomness::get(),
+			config: NextEpochConfig::get().unwrap_or_else(EpochConfig::get),
 		}
 	}
 
@@ -668,12 +685,13 @@ impl<T: Config> Module<T> {
 		this_randomness
 	}
 
-	fn initialize_authorities(authorities: &[(AuthorityId, BabeAuthorityWeight)]) {
+	fn initialize_authorities_and_epoch_config(authorities: &[(AuthorityId, BabeAuthorityWeight)], epoch_config: &BabeEpochConfiguration) {
 		if !authorities.is_empty() {
 			assert!(Authorities::get().is_empty(), "Authorities are already initialized!");
 			Authorities::put(authorities);
 			NextAuthorities::put(authorities);
 		}
+		EpochConfig::put(epoch_config);
 	}
 
 	fn do_report_equivocation(
diff --git a/primitives/consensus/babe/src/lib.rs b/primitives/consensus/babe/src/lib.rs
index 6987796c114a0..2bcf6f7e79277 100644
--- a/primitives/consensus/babe/src/lib.rs
+++ b/primitives/consensus/babe/src/lib.rs
@@ -225,6 +225,12 @@ pub enum AllowedSlots {
 	PrimaryAndSecondaryVRFSlots,
 }
 
+impl Default for AllowedSlots {
+	fn default() -> Self {
+		Self::PrimarySlots
+	}
+}
+
 impl AllowedSlots {
 	/// Whether plain secondary slots are allowed.
 	pub fn is_secondary_plain_slots_allowed(&self) -> bool {
@@ -247,7 +253,7 @@ impl sp_consensus::SlotData for BabeGenesisConfiguration {
 }
 
 /// Configuration data used by the BABE consensus engine.
-#[derive(Clone, PartialEq, Eq, Encode, Decode, RuntimeDebug)]
+#[derive(Clone, PartialEq, Eq, Encode, Decode, RuntimeDebug, Default)]
 pub struct BabeEpochConfiguration {
 	/// A constant value that is used in the threshold calculation formula.
 	/// Expressed as a rational where the first member of the tuple is the
@@ -362,6 +368,8 @@ pub struct Epoch {
 	pub authorities: Vec<(AuthorityId, BabeAuthorityWeight)>,
 	/// Randomness for this epoch.
 	pub randomness: [u8; VRF_OUTPUT_LENGTH],
+	/// Configuration of the epoch.
+	pub config: BabeEpochConfiguration,
 }
 
 sp_api::decl_runtime_apis! {

From 4cbac3ebc11a1391a42af188dfed5441d6f8e213 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 17 Feb 2021 15:17:17 +0100
Subject: [PATCH 02/19] Add serialize/deserialize, copy babe epoch config
 defaults from node runtime

---
 Cargo.lock                            |  2 ++
 frame/babe/src/lib.rs                 |  3 ++-
 primitives/consensus/babe/Cargo.toml  |  1 +
 primitives/consensus/babe/src/lib.rs  | 25 ++++++++++++++++---------
 primitives/consensus/slots/Cargo.toml |  1 +
 primitives/consensus/slots/src/lib.rs |  3 ++-
 6 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e3c9131a86e3b..34384621eb8de 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8349,6 +8349,7 @@ version = "0.9.0"
 dependencies = [
  "merlin",
  "parity-scale-codec",
+ "serde",
  "sp-api",
  "sp-application-crypto",
  "sp-consensus",
@@ -8378,6 +8379,7 @@ name = "sp-consensus-slots"
 version = "0.9.0"
 dependencies = [
  "parity-scale-codec",
+ "serde",
  "sp-arithmetic",
  "sp-runtime",
 ]
diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index e63564ea76ba9..609f26b169b3d 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -800,7 +800,8 @@ impl<T: Config> OneSessionHandler<T::AccountId> for Module<T> {
 		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
 	{
 		let authorities = validators.map(|(_, k)| (k, 1)).collect::<Vec<_>>();
-		Self::initialize_authorities(&authorities);
+		let epoch_config = BabeEpochConfiguration::default();
+		Self::initialize_authorities_and_epoch_config(&authorities, &epoch_config);
 	}
 
 	fn on_new_session<'a, I: 'a>(_changed: bool, validators: I, queued_validators: I)
diff --git a/primitives/consensus/babe/Cargo.toml b/primitives/consensus/babe/Cargo.toml
index fb02014eeef51..3f36723c16f5e 100644
--- a/primitives/consensus/babe/Cargo.toml
+++ b/primitives/consensus/babe/Cargo.toml
@@ -26,6 +26,7 @@ sp-inherents = { version = "3.0.0", default-features = false, path = "../../inhe
 sp-keystore = { version = "0.9.0", default-features = false, path = "../../keystore", optional = true }
 sp-runtime = { version = "3.0.0", default-features = false, path = "../../runtime" }
 sp-timestamp = { version = "3.0.0", default-features = false, path = "../../timestamp" }
+serde = { version = "1.0.123", features = ["derive"] }
 
 [features]
 default = ["std"]
diff --git a/primitives/consensus/babe/src/lib.rs b/primitives/consensus/babe/src/lib.rs
index 2bcf6f7e79277..849a00e4f1406 100644
--- a/primitives/consensus/babe/src/lib.rs
+++ b/primitives/consensus/babe/src/lib.rs
@@ -29,6 +29,7 @@ pub use sp_consensus_vrf::schnorrkel::{
 };
 
 use codec::{Decode, Encode};
+use serde::{Serialize, Deserialize};
 #[cfg(feature = "std")]
 use sp_keystore::vrf::{VRFTranscriptData, VRFTranscriptValue};
 use sp_runtime::{traits::Header, ConsensusEngineId, RuntimeDebug};
@@ -47,6 +48,9 @@ mod app {
 /// The prefix used by BABE for its VRF keys.
 pub const BABE_VRF_PREFIX: &[u8] = b"substrate-babe-vrf";
 
+/// 1 in 4 blocks (on average, not counting collisions) will be primary BABE blocks.
+pub const PRIMARY_PROBABILITY: (u64, u64) = (1, 4);
+
 /// BABE VRFInOut context.
 pub static BABE_VRF_INOUT_CONTEXT: &[u8] = b"BabeVRFInOutContext";
 
@@ -215,7 +219,7 @@ pub struct BabeGenesisConfiguration {
 }
 
 /// Types of allowed slots.
-#[derive(Clone, Copy, PartialEq, Eq, Encode, Decode, RuntimeDebug)]
+#[derive(Clone, Copy, PartialEq, Eq, Encode, Decode, RuntimeDebug, Serialize, Deserialize)]
 pub enum AllowedSlots {
 	/// Only allow primary slots.
 	PrimarySlots,
@@ -225,12 +229,6 @@ pub enum AllowedSlots {
 	PrimaryAndSecondaryVRFSlots,
 }
 
-impl Default for AllowedSlots {
-	fn default() -> Self {
-		Self::PrimarySlots
-	}
-}
-
 impl AllowedSlots {
 	/// Whether plain secondary slots are allowed.
 	pub fn is_secondary_plain_slots_allowed(&self) -> bool {
@@ -253,7 +251,7 @@ impl sp_consensus::SlotData for BabeGenesisConfiguration {
 }
 
 /// Configuration data used by the BABE consensus engine.
-#[derive(Clone, PartialEq, Eq, Encode, Decode, RuntimeDebug, Default)]
+#[derive(Clone, PartialEq, Eq, Encode, Decode, RuntimeDebug, Serialize, Deserialize)]
 pub struct BabeEpochConfiguration {
 	/// A constant value that is used in the threshold calculation formula.
 	/// Expressed as a rational where the first member of the tuple is the
@@ -268,6 +266,15 @@ pub struct BabeEpochConfiguration {
 	pub allowed_slots: AllowedSlots,
 }
 
+impl Default for BabeEpochConfiguration {
+	fn default() -> Self {
+		Self {
+			c: PRIMARY_PROBABILITY,
+			allowed_slots: AllowedSlots::PrimaryAndSecondaryPlainSlots,
+		}
+	}
+}
+
 /// Verifies the equivocation proof by making sure that: both headers have
 /// different hashes, are targetting the same slot, and have valid signatures by
 /// the same authority.
@@ -356,7 +363,7 @@ impl OpaqueKeyOwnershipProof {
 }
 
 /// BABE epoch information
-#[derive(Decode, Encode, PartialEq, Eq, Clone, Debug)]
+#[derive(Decode, Encode, PartialEq, Eq, Clone, Debug, Serialize, Deserialize)]
 pub struct Epoch {
 	/// The epoch index.
 	pub epoch_index: u64,
diff --git a/primitives/consensus/slots/Cargo.toml b/primitives/consensus/slots/Cargo.toml
index 46dbaca1a6ad4..95d22360a1190 100644
--- a/primitives/consensus/slots/Cargo.toml
+++ b/primitives/consensus/slots/Cargo.toml
@@ -16,6 +16,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 codec = { package = "parity-scale-codec", version = "2.0.0", default-features = false, features = ["derive"] }
 sp-runtime = { version = "3.0.0", default-features = false, path = "../../runtime" }
 sp-arithmetic = { version = "3.0.0", default-features = false, path = "../../arithmetic" }
+serde = { version = "1.0.123", features = ["derive"] }
 
 [features]
 default = ["std"]
diff --git a/primitives/consensus/slots/src/lib.rs b/primitives/consensus/slots/src/lib.rs
index 545d18af1f9be..3e453ccdf4db6 100644
--- a/primitives/consensus/slots/src/lib.rs
+++ b/primitives/consensus/slots/src/lib.rs
@@ -20,9 +20,10 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 use codec::{Decode, Encode};
+use serde::{Serialize, Deserialize};
 
 /// Unit type wrapper that represents a slot.
-#[derive(Debug, Encode, Decode, Eq, Clone, Copy, Default, Ord)]
+#[derive(Debug, Encode, Decode, Eq, Clone, Copy, Default, Ord, Serialize, Deserialize)]
 pub struct Slot(u64);
 
 impl core::ops::Deref for Slot {

From cb2203457375b54b1f65c11b4c91d6852e90e6ed Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 17 Feb 2021 15:56:09 +0100
Subject: [PATCH 03/19] Fix line widths and turn default features off for serde

---
 frame/babe/src/lib.rs                 | 5 ++++-
 primitives/consensus/babe/Cargo.toml  | 3 ++-
 primitives/consensus/slots/Cargo.toml | 3 ++-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index 609f26b169b3d..ff8d9c47552d5 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -685,7 +685,10 @@ impl<T: Config> Module<T> {
 		this_randomness
 	}
 
-	fn initialize_authorities_and_epoch_config(authorities: &[(AuthorityId, BabeAuthorityWeight)], epoch_config: &BabeEpochConfiguration) {
+	fn initialize_authorities_and_epoch_config(
+		authorities: &[(AuthorityId, BabeAuthorityWeight)],
+		epoch_config: &BabeEpochConfiguration,
+	) {
 		if !authorities.is_empty() {
 			assert!(Authorities::get().is_empty(), "Authorities are already initialized!");
 			Authorities::put(authorities);
diff --git a/primitives/consensus/babe/Cargo.toml b/primitives/consensus/babe/Cargo.toml
index 3f36723c16f5e..c6645ae706288 100644
--- a/primitives/consensus/babe/Cargo.toml
+++ b/primitives/consensus/babe/Cargo.toml
@@ -26,7 +26,7 @@ sp-inherents = { version = "3.0.0", default-features = false, path = "../../inhe
 sp-keystore = { version = "0.9.0", default-features = false, path = "../../keystore", optional = true }
 sp-runtime = { version = "3.0.0", default-features = false, path = "../../runtime" }
 sp-timestamp = { version = "3.0.0", default-features = false, path = "../../timestamp" }
-serde = { version = "1.0.123", features = ["derive"] }
+serde = { version = "1.0.123", features = ["derive"], default-features = false }
 
 [features]
 default = ["std"]
@@ -44,4 +44,5 @@ std = [
 	"sp-keystore",
 	"sp-runtime/std",
 	"sp-timestamp/std",
+	"serde/std",
 ]
diff --git a/primitives/consensus/slots/Cargo.toml b/primitives/consensus/slots/Cargo.toml
index 95d22360a1190..9167e90523dff 100644
--- a/primitives/consensus/slots/Cargo.toml
+++ b/primitives/consensus/slots/Cargo.toml
@@ -16,7 +16,7 @@ targets = ["x86_64-unknown-linux-gnu"]
 codec = { package = "parity-scale-codec", version = "2.0.0", default-features = false, features = ["derive"] }
 sp-runtime = { version = "3.0.0", default-features = false, path = "../../runtime" }
 sp-arithmetic = { version = "3.0.0", default-features = false, path = "../../arithmetic" }
-serde = { version = "1.0.123", features = ["derive"] }
+serde = { version = "1.0.123", features = ["derive"], default-features = false }
 
 [features]
 default = ["std"]
@@ -24,4 +24,5 @@ std = [
 	"codec/std",
 	"sp-runtime/std",
 	"sp-arithmetic/std",
+	"serde/std"
 ]

From a85fdcdd6034828d39105960865cd5469b40c4d3 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 17 Feb 2021 16:25:44 +0100
Subject: [PATCH 04/19] Remove ser/deser from Epoch, fix node-cli

---
 bin/node/cli/src/chain_spec.rs       |  1 +
 bin/node/runtime/src/lib.rs          | 10 ++++++++--
 primitives/consensus/babe/src/lib.rs |  2 +-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/bin/node/cli/src/chain_spec.rs b/bin/node/cli/src/chain_spec.rs
index 7de9cfd0b6aa5..35edba1dc2b9d 100644
--- a/bin/node/cli/src/chain_spec.rs
+++ b/bin/node/cli/src/chain_spec.rs
@@ -305,6 +305,7 @@ pub fn testnet_genesis(
 		}),
 		pallet_babe: Some(BabeConfig {
 			authorities: vec![],
+			epoch_config: node_runtime::BABE_GENESIS_EPOCH_CONFIG,
 		}),
 		pallet_im_online: Some(ImOnlineConfig {
 			keys: vec![],
diff --git a/bin/node/runtime/src/lib.rs b/bin/node/runtime/src/lib.rs
index 86e3075c3ae5d..9c864638edcbf 100644
--- a/bin/node/runtime/src/lib.rs
+++ b/bin/node/runtime/src/lib.rs
@@ -118,6 +118,12 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	transaction_version: 2,
 };
 
+pub const BABE_GENESIS_EPOCH_CONFIG: sp_consensus_babe::BabeEpochConfiguration =
+	sp_consensus_babe::BabeEpochConfiguration {
+		c: PRIMARY_PROBABILITY,
+		allowed_slots: sp_consensus_babe::AllowedSlots::PrimaryAndSecondaryPlainSlots
+	};
+
 /// Native version.
 #[cfg(any(feature = "std", test))]
 pub fn native_version() -> NativeVersion {
@@ -1196,10 +1202,10 @@ impl_runtime_apis! {
 			sp_consensus_babe::BabeGenesisConfiguration {
 				slot_duration: Babe::slot_duration(),
 				epoch_length: EpochDuration::get(),
-				c: PRIMARY_PROBABILITY,
+				c: BABE_GENESIS_EPOCH_CONFIG.c,
 				genesis_authorities: Babe::authorities(),
 				randomness: Babe::randomness(),
-				allowed_slots: sp_consensus_babe::AllowedSlots::PrimaryAndSecondaryPlainSlots,
+				allowed_slots: BABE_GENESIS_EPOCH_CONFIG.allowed_slots,
 			}
 		}
 
diff --git a/primitives/consensus/babe/src/lib.rs b/primitives/consensus/babe/src/lib.rs
index 849a00e4f1406..78974264cf70a 100644
--- a/primitives/consensus/babe/src/lib.rs
+++ b/primitives/consensus/babe/src/lib.rs
@@ -363,7 +363,7 @@ impl OpaqueKeyOwnershipProof {
 }
 
 /// BABE epoch information
-#[derive(Decode, Encode, PartialEq, Eq, Clone, Debug, Serialize, Deserialize)]
+#[derive(Decode, Encode, PartialEq, Eq, Clone, Debug)]
 pub struct Epoch {
 	/// The epoch index.
 	pub epoch_index: u64,

From 11826f0096af7d9e0e2f9ceba2d573a6ce41b82a Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Mon, 22 Feb 2021 13:16:04 +0100
Subject: [PATCH 05/19] Apply suggestions

---
 frame/babe/src/lib.rs                | 21 ++++++++++-----------
 primitives/consensus/babe/src/lib.rs |  4 ++--
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index ff8d9c47552d5..374d988e6f905 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -187,7 +187,7 @@ decl_storage! {
 		// variable to its underlying value.
 		pub Randomness get(fn randomness): schnorrkel::Randomness;
 
-		/// Next epoch configuration, if changed.
+		/// Pending epoch configuration change that will be applied when the next epoch is enacted.
 		PendingEpochConfigChange: Option<NextConfigDescriptor>;
 
 		/// Next epoch randomness.
@@ -226,15 +226,19 @@ decl_storage! {
 		/// execution context should always yield zero.
 		Lateness get(fn lateness): T::BlockNumber;
 
+		/// The configuration for the current epoch.
 		EpochConfig: BabeEpochConfiguration;
+
+		/// The configuration for the next epoch. Uses `EpochConfig` instead if `None`.
 		NextEpochConfig: Option<BabeEpochConfiguration>;
 	}
 	add_extra_genesis {
 		config(authorities): Vec<(AuthorityId, BabeAuthorityWeight)>;
 		config(epoch_config): BabeEpochConfiguration;
-		build(|config| Module::<T>::initialize_authorities_and_epoch_config(
-			&config.authorities, &config.epoch_config
-		))
+		build(|config| {
+			Module::<T>::initialize_authorities(&config.authorities);
+			EpochConfig::put(config.epoch_config.clone());
+		})
 	}
 }
 
@@ -685,16 +689,12 @@ impl<T: Config> Module<T> {
 		this_randomness
 	}
 
-	fn initialize_authorities_and_epoch_config(
-		authorities: &[(AuthorityId, BabeAuthorityWeight)],
-		epoch_config: &BabeEpochConfiguration,
-	) {
+	fn initialize_authorities(authorities: &[(AuthorityId, BabeAuthorityWeight)]) {
 		if !authorities.is_empty() {
 			assert!(Authorities::get().is_empty(), "Authorities are already initialized!");
 			Authorities::put(authorities);
 			NextAuthorities::put(authorities);
 		}
-		EpochConfig::put(epoch_config);
 	}
 
 	fn do_report_equivocation(
@@ -803,8 +803,7 @@ impl<T: Config> OneSessionHandler<T::AccountId> for Module<T> {
 		where I: Iterator<Item=(&'a T::AccountId, AuthorityId)>
 	{
 		let authorities = validators.map(|(_, k)| (k, 1)).collect::<Vec<_>>();
-		let epoch_config = BabeEpochConfiguration::default();
-		Self::initialize_authorities_and_epoch_config(&authorities, &epoch_config);
+		Self::initialize_authorities(&authorities);
 	}
 
 	fn on_new_session<'a, I: 'a>(_changed: bool, validators: I, queued_validators: I)
diff --git a/primitives/consensus/babe/src/lib.rs b/primitives/consensus/babe/src/lib.rs
index 78974264cf70a..208910898a16f 100644
--- a/primitives/consensus/babe/src/lib.rs
+++ b/primitives/consensus/babe/src/lib.rs
@@ -49,7 +49,7 @@ mod app {
 pub const BABE_VRF_PREFIX: &[u8] = b"substrate-babe-vrf";
 
 /// 1 in 4 blocks (on average, not counting collisions) will be primary BABE blocks.
-pub const PRIMARY_PROBABILITY: (u64, u64) = (1, 4);
+pub const DEFAULT_PRIMARY_PROBABILITY: (u64, u64) = (1, 4);
 
 /// BABE VRFInOut context.
 pub static BABE_VRF_INOUT_CONTEXT: &[u8] = b"BabeVRFInOutContext";
@@ -269,7 +269,7 @@ pub struct BabeEpochConfiguration {
 impl Default for BabeEpochConfiguration {
 	fn default() -> Self {
 		Self {
-			c: PRIMARY_PROBABILITY,
+			c: DEFAULT_PRIMARY_PROBABILITY,
 			allowed_slots: AllowedSlots::PrimaryAndSecondaryPlainSlots,
 		}
 	}

From 6c920d6f160b1a5d7d12dcbc83898e1bf322276f Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Tue, 23 Feb 2021 15:41:33 +0100
Subject: [PATCH 06/19] Add comment to BABE_GENESIS_EPOCH_CONFIG in bin

---
 bin/node/runtime/src/lib.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/bin/node/runtime/src/lib.rs b/bin/node/runtime/src/lib.rs
index 9b0ccd65928c4..9cc40c3796b4c 100644
--- a/bin/node/runtime/src/lib.rs
+++ b/bin/node/runtime/src/lib.rs
@@ -118,6 +118,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	transaction_version: 2,
 };
 
+/// The BABE epoch configuration at genesis.
 pub const BABE_GENESIS_EPOCH_CONFIG: sp_consensus_babe::BabeEpochConfiguration =
 	sp_consensus_babe::BabeEpochConfiguration {
 		c: PRIMARY_PROBABILITY,

From 099673ac11046491959ccc3cd26a7770dd47a020 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 24 Feb 2021 14:38:29 +0100
Subject: [PATCH 07/19] Apply suggestions

---
 frame/babe/src/lib.rs                | 14 ++++++++------
 primitives/consensus/babe/Cargo.toml |  4 ++--
 primitives/consensus/babe/src/lib.rs | 19 +++++--------------
 3 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index 374d988e6f905..f6ebb41bdc990 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -226,18 +226,18 @@ decl_storage! {
 		/// execution context should always yield zero.
 		Lateness get(fn lateness): T::BlockNumber;
 
-		/// The configuration for the current epoch.
-		EpochConfig: BabeEpochConfiguration;
+		/// The configuration for the current epoch. Should never be `None`.
+		EpochConfig: Option<BabeEpochConfiguration>;
 
 		/// The configuration for the next epoch. Uses `EpochConfig` instead if `None`.
 		NextEpochConfig: Option<BabeEpochConfiguration>;
 	}
 	add_extra_genesis {
 		config(authorities): Vec<(AuthorityId, BabeAuthorityWeight)>;
-		config(epoch_config): BabeEpochConfiguration;
+		config(epoch_config): Option<BabeEpochConfiguration>;
 		build(|config| {
 			Module::<T>::initialize_authorities(&config.authorities);
-			EpochConfig::put(config.epoch_config.clone());
+			EpochConfig::put(config.epoch_config.clone().expect("EpochConfig is never None; qed"));
 		})
 	}
 }
@@ -522,7 +522,7 @@ impl<T: Config> Module<T> {
 			duration: T::EpochDuration::get(),
 			authorities: Self::authorities(),
 			randomness: Self::randomness(),
-			config: EpochConfig::get(),
+			config: EpochConfig::get().expect("EpochConfig is never None; qed"),
 		}
 	}
 
@@ -540,7 +540,9 @@ impl<T: Config> Module<T> {
 			duration: T::EpochDuration::get(),
 			authorities: NextAuthorities::get(),
 			randomness: NextRandomness::get(),
-			config: NextEpochConfig::get().unwrap_or_else(EpochConfig::get),
+			config: NextEpochConfig::get().unwrap_or_else(|| {
+				EpochConfig::get().expect("EpochConfig is never None; qed")
+			}),
 		}
 	}
 
diff --git a/primitives/consensus/babe/Cargo.toml b/primitives/consensus/babe/Cargo.toml
index c6645ae706288..a8ab03dcdaa49 100644
--- a/primitives/consensus/babe/Cargo.toml
+++ b/primitives/consensus/babe/Cargo.toml
@@ -26,7 +26,7 @@ sp-inherents = { version = "3.0.0", default-features = false, path = "../../inhe
 sp-keystore = { version = "0.9.0", default-features = false, path = "../../keystore", optional = true }
 sp-runtime = { version = "3.0.0", default-features = false, path = "../../runtime" }
 sp-timestamp = { version = "3.0.0", default-features = false, path = "../../timestamp" }
-serde = { version = "1.0.123", features = ["derive"], default-features = false }
+serde = { version = "1.0.123", features = ["derive"], optional = true }
 
 [features]
 default = ["std"]
@@ -44,5 +44,5 @@ std = [
 	"sp-keystore",
 	"sp-runtime/std",
 	"sp-timestamp/std",
-	"serde/std",
+	"serde",
 ]
diff --git a/primitives/consensus/babe/src/lib.rs b/primitives/consensus/babe/src/lib.rs
index 208910898a16f..1b416c996fcf0 100644
--- a/primitives/consensus/babe/src/lib.rs
+++ b/primitives/consensus/babe/src/lib.rs
@@ -29,6 +29,7 @@ pub use sp_consensus_vrf::schnorrkel::{
 };
 
 use codec::{Decode, Encode};
+#[cfg(feature = "std")]
 use serde::{Serialize, Deserialize};
 #[cfg(feature = "std")]
 use sp_keystore::vrf::{VRFTranscriptData, VRFTranscriptValue};
@@ -48,9 +49,6 @@ mod app {
 /// The prefix used by BABE for its VRF keys.
 pub const BABE_VRF_PREFIX: &[u8] = b"substrate-babe-vrf";
 
-/// 1 in 4 blocks (on average, not counting collisions) will be primary BABE blocks.
-pub const DEFAULT_PRIMARY_PROBABILITY: (u64, u64) = (1, 4);
-
 /// BABE VRFInOut context.
 pub static BABE_VRF_INOUT_CONTEXT: &[u8] = b"BabeVRFInOutContext";
 
@@ -219,7 +217,8 @@ pub struct BabeGenesisConfiguration {
 }
 
 /// Types of allowed slots.
-#[derive(Clone, Copy, PartialEq, Eq, Encode, Decode, RuntimeDebug, Serialize, Deserialize)]
+#[derive(Clone, Copy, PartialEq, Eq, Encode, Decode, RuntimeDebug)]
+#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
 pub enum AllowedSlots {
 	/// Only allow primary slots.
 	PrimarySlots,
@@ -251,7 +250,8 @@ impl sp_consensus::SlotData for BabeGenesisConfiguration {
 }
 
 /// Configuration data used by the BABE consensus engine.
-#[derive(Clone, PartialEq, Eq, Encode, Decode, RuntimeDebug, Serialize, Deserialize)]
+#[derive(Clone, PartialEq, Eq, Encode, Decode, RuntimeDebug)]
+#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
 pub struct BabeEpochConfiguration {
 	/// A constant value that is used in the threshold calculation formula.
 	/// Expressed as a rational where the first member of the tuple is the
@@ -266,15 +266,6 @@ pub struct BabeEpochConfiguration {
 	pub allowed_slots: AllowedSlots,
 }
 
-impl Default for BabeEpochConfiguration {
-	fn default() -> Self {
-		Self {
-			c: DEFAULT_PRIMARY_PROBABILITY,
-			allowed_slots: AllowedSlots::PrimaryAndSecondaryPlainSlots,
-		}
-	}
-}
-
 /// Verifies the equivocation proof by making sure that: both headers have
 /// different hashes, are targetting the same slot, and have valid signatures by
 /// the same authority.

From f5d3b18b04cce394463a950868a1eaac1d7cd3e1 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 24 Feb 2021 17:45:43 +0100
Subject: [PATCH 08/19] Add a sketchy migration function

---
 frame/babe/src/lib.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index f6ebb41bdc990..596b107bdbaca 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -847,3 +847,32 @@ fn compute_randomness(
 
 	sp_io::hashing::blake2_256(&s)
 }
+
+pub mod migrations {
+	use super::*;
+	use frame_support::pallet_prelude::{ValueQuery, StorageValue};
+
+	struct __OldNextEpochConfig;
+	impl frame_support::traits::StorageInstance for __OldNextEpochConfig {
+		fn pallet_prefix() -> &'static str { "BabeApi" }
+		const STORAGE_PREFIX: &'static str = "NextEpochConfig";
+	}
+
+	type OldNextEpochConfig = StorageValue<
+		__OldNextEpochConfig, Option<NextConfigDescriptor>, ValueQuery
+	>;
+
+	pub fn add_epoch_configurations(
+		current_epoch_config: BabeEpochConfiguration,
+		next_epoch_config: BabeEpochConfiguration,
+	) {
+		if let Some(pending_change) = OldNextEpochConfig::get() {
+			PendingEpochConfigChange::put(pending_change);
+		}
+
+		OldNextEpochConfig::kill();
+
+		EpochConfig::put(current_epoch_config);
+		NextEpochConfig::put(next_epoch_config);
+	}
+}

From 44ad4ffacc5fd4f01e28c24317915444820fd2bd Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 24 Feb 2021 18:40:30 +0100
Subject: [PATCH 09/19] Add a migration test

---
 frame/babe/src/tests.rs | 55 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/frame/babe/src/tests.rs b/frame/babe/src/tests.rs
index 8576389af31ff..6167cedbe6c19 100644
--- a/frame/babe/src/tests.rs
+++ b/frame/babe/src/tests.rs
@@ -25,7 +25,7 @@ use frame_support::{
 };
 use mock::*;
 use pallet_session::ShouldEndSession;
-use sp_consensus_babe::{AllowedSlots, Slot};
+use sp_consensus_babe::{AllowedSlots, Slot, BabeEpochConfiguration};
 use sp_core::crypto::Pair;
 
 const EMPTY_RANDOMNESS: [u8; 32] = [
@@ -255,6 +255,11 @@ fn can_enact_next_config() {
 #[test]
 fn can_fetch_current_and_next_epoch_data() {
 	new_test_ext(5).execute_with(|| {
+		EpochConfig::put(BabeEpochConfiguration {
+			c: (1, 4),
+			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
+		});
+
 		// genesis authorities should be used for the first and second epoch
 		assert_eq!(
 			Babe::current_epoch().authorities,
@@ -773,3 +778,51 @@ fn valid_equivocation_reports_dont_pay_fees() {
 		assert_eq!(post_info.pays_fee, Pays::Yes);
 	})
 }
+
+#[test]
+fn add_epoch_configurations_migration_works() {
+	use frame_support::storage::migration::{
+		put_storage_value, get_storage_value,
+	};
+
+	new_test_ext(1).execute_with(|| {
+		put_storage_value(
+			b"BabeApi",
+			b"NextEpochConfig",
+			&[],
+			Some(NextConfigDescriptor::V1 {
+				c: (2, 4),
+				allowed_slots: AllowedSlots::PrimarySlots
+			})
+		);
+
+		assert!(get_storage_value::<Option<NextConfigDescriptor>>(
+			b"BabeApi",
+			b"NextEpochConfig",
+			&[],
+		).is_some());
+
+		let current_epoch = BabeEpochConfiguration {
+			c: (1, 4),
+			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
+		};
+
+		let next_epoch = BabeEpochConfiguration {
+			c: (2, 4),
+			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
+		};
+
+		crate::migrations::add_epoch_configurations(
+			current_epoch.clone(), next_epoch.clone()
+		);
+
+		assert!(get_storage_value::<Option<NextConfigDescriptor>>(
+			b"BabeApi",
+			b"NextEpochConfig",
+			&[],
+		).is_none());
+
+		assert_eq!(EpochConfig::get(), Some(current_epoch));
+		assert_eq!(NextEpochConfig::get(), Some(next_epoch));
+	});
+}

From 8d9cecb401aca019c37c0b6f45c2187fd5eb8c86 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 24 Feb 2021 18:43:31 +0100
Subject: [PATCH 10/19] Check for PendingEpochConfigChange as well

---
 frame/babe/src/tests.rs | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/frame/babe/src/tests.rs b/frame/babe/src/tests.rs
index 6167cedbe6c19..daced9afe3a6f 100644
--- a/frame/babe/src/tests.rs
+++ b/frame/babe/src/tests.rs
@@ -786,14 +786,16 @@ fn add_epoch_configurations_migration_works() {
 	};
 
 	new_test_ext(1).execute_with(|| {
+		let next_config_descriptor = NextConfigDescriptor::V1 {
+			c: (3, 4),
+			allowed_slots: AllowedSlots::PrimarySlots
+		};
+
 		put_storage_value(
 			b"BabeApi",
 			b"NextEpochConfig",
 			&[],
-			Some(NextConfigDescriptor::V1 {
-				c: (2, 4),
-				allowed_slots: AllowedSlots::PrimarySlots
-			})
+			Some(next_config_descriptor.clone())
 		);
 
 		assert!(get_storage_value::<Option<NextConfigDescriptor>>(
@@ -824,5 +826,6 @@ fn add_epoch_configurations_migration_works() {
 
 		assert_eq!(EpochConfig::get(), Some(current_epoch));
 		assert_eq!(NextEpochConfig::get(), Some(next_epoch));
+		assert_eq!(PendingEpochConfigChange::get(), Some(next_config_descriptor));
 	});
 }

From a339dfe32b33a3ff4234ddefd61c39cc1149819e Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Thu, 25 Feb 2021 10:46:45 +0100
Subject: [PATCH 11/19] Make epoch_config in node-cli

---
 bin/node/cli/src/chain_spec.rs | 2 +-
 frame/babe/src/lib.rs          | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/bin/node/cli/src/chain_spec.rs b/bin/node/cli/src/chain_spec.rs
index 35edba1dc2b9d..f5b2dd6695e3a 100644
--- a/bin/node/cli/src/chain_spec.rs
+++ b/bin/node/cli/src/chain_spec.rs
@@ -305,7 +305,7 @@ pub fn testnet_genesis(
 		}),
 		pallet_babe: Some(BabeConfig {
 			authorities: vec![],
-			epoch_config: node_runtime::BABE_GENESIS_EPOCH_CONFIG,
+			epoch_config: Some(node_runtime::BABE_GENESIS_EPOCH_CONFIG),
 		}),
 		pallet_im_online: Some(ImOnlineConfig {
 			keys: vec![],
diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index 596b107bdbaca..f41c10303b651 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -237,7 +237,7 @@ decl_storage! {
 		config(epoch_config): Option<BabeEpochConfiguration>;
 		build(|config| {
 			Module::<T>::initialize_authorities(&config.authorities);
-			EpochConfig::put(config.epoch_config.clone().expect("EpochConfig is never None; qed"));
+			EpochConfig::put(config.epoch_config.clone().expect("epoch_config must not be None"));
 		})
 	}
 }

From fed698dbb44ab54613535f8808b060a39a16f805 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Thu, 25 Feb 2021 12:28:35 +0100
Subject: [PATCH 12/19] Move updating EpochConfig out of the if

---
 frame/babe/src/lib.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index f41c10303b651..cfb9d943b8547 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -494,11 +494,11 @@ impl<T: Config> Module<T> {
 		};
 		Self::deposit_consensus(ConsensusLog::NextEpochData(next_epoch));
 
-		if let Some(pending_epoch_config_change) = PendingEpochConfigChange::take() {
-			if let Some(next_config) = NextEpochConfig::get() {
-				EpochConfig::put(next_config);
-			}
+		if let Some(next_config) = NextEpochConfig::get() {
+			EpochConfig::put(next_config);
+		}
 
+		if let Some(pending_epoch_config_change) = PendingEpochConfigChange::take() {
 			let next_epoch_config: BabeEpochConfiguration =
 				pending_epoch_config_change.clone().into();
 			NextEpochConfig::put(next_epoch_config);

From 5834911e46b36aa3c5f26ffdc83c92cfda11bf0f Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Thu, 25 Feb 2021 16:28:53 +0100
Subject: [PATCH 13/19] Fix executor tests

---
 bin/node/testing/src/genesis.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/bin/node/testing/src/genesis.rs b/bin/node/testing/src/genesis.rs
index 75d0d18e6ef81..bdbd87366d20a 100644
--- a/bin/node/testing/src/genesis.rs
+++ b/bin/node/testing/src/genesis.rs
@@ -23,7 +23,7 @@ use sp_keyring::{Ed25519Keyring, Sr25519Keyring};
 use node_runtime::{
 	GenesisConfig, BalancesConfig, SessionConfig, StakingConfig, SystemConfig,
 	GrandpaConfig, IndicesConfig, ContractsConfig, SocietyConfig, wasm_binary_unwrap,
-	AccountId, StakerStatus,
+	AccountId, StakerStatus, BabeConfig, BABE_GENESIS_EPOCH_CONFIG,
 };
 use node_runtime::constants::currency::*;
 use sp_core::ChangesTrieConfiguration;
@@ -100,7 +100,10 @@ pub fn config_endowed(
 		pallet_contracts: Some(ContractsConfig {
 			current_schedule: Default::default(),
 		}),
-		pallet_babe: Some(Default::default()),
+		pallet_babe: Some(BabeConfig {
+			authorities: vec![],
+			epoch_config: Some(BABE_GENESIS_EPOCH_CONFIG),
+		}),
 		pallet_grandpa: Some(GrandpaConfig {
 			authorities: vec![],
 		}),

From a06fb96e4295d4ac20f7e159453613fa4ae11fbe Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Mon, 1 Mar 2021 15:22:17 +0100
Subject: [PATCH 14/19] Calculate weight for add_epoch_configurations

---
 frame/babe/src/lib.rs | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index bf057fe3c4d1f..49ebaa53e6b2c 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -868,17 +868,28 @@ pub mod migrations {
 		__OldNextEpochConfig, Option<NextConfigDescriptor>, ValueQuery
 	>;
 
-	pub fn add_epoch_configurations(
+	pub fn add_epoch_configurations<T: Config>(
 		current_epoch_config: BabeEpochConfiguration,
 		next_epoch_config: BabeEpochConfiguration,
-	) {
+	) -> Weight {
+		let mut writes = 0;
+		let mut reads = 0;
+
 		if let Some(pending_change) = OldNextEpochConfig::get() {
 			PendingEpochConfigChange::put(pending_change);
+
+			writes += 1;
 		}
 
+		reads += 1;
+
 		OldNextEpochConfig::kill();
 
 		EpochConfig::put(current_epoch_config);
 		NextEpochConfig::put(next_epoch_config);
+
+		writes += 3;
+
+		T::DbWeight::get().writes(writes) + T::DbWeight::get().reads(reads)
 	}
 }

From 7fd659f84ed2adb06f2bd70dbd8f398889f4dad7 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Mon, 1 Mar 2021 15:35:03 +0100
Subject: [PATCH 15/19] Fix babe test

---
 frame/babe/src/tests.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/frame/babe/src/tests.rs b/frame/babe/src/tests.rs
index daced9afe3a6f..9da77ec5a19ad 100644
--- a/frame/babe/src/tests.rs
+++ b/frame/babe/src/tests.rs
@@ -814,7 +814,7 @@ fn add_epoch_configurations_migration_works() {
 			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
 		};
 
-		crate::migrations::add_epoch_configurations(
+		crate::migrations::add_epoch_configurations::<Test>(
 			current_epoch.clone(), next_epoch.clone()
 		);
 

From bdfe0397449748fe1df81486a74744acd4f31fe7 Mon Sep 17 00:00:00 2001
From: Ashley <ashley.ruglys@gmail.com>
Date: Wed, 3 Mar 2021 09:28:13 +0100
Subject: [PATCH 16/19] Apply suggestions from code review
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>
---
 frame/babe/src/lib.rs | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index 47b29d97a1695..bbbac0436dcac 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -227,10 +227,11 @@ decl_storage! {
 		/// execution context should always yield zero.
 		Lateness get(fn lateness): T::BlockNumber;
 
-		/// The configuration for the current epoch. Should never be `None`.
+		/// The configuration for the current epoch. Should never be `None` as it is initialized in genesis.
 		EpochConfig: Option<BabeEpochConfiguration>;
 
-		/// The configuration for the next epoch. Uses `EpochConfig` instead if `None`.
+		/// The configuration for the next epoch, `None` if the config will not change
+		/// (you can fallback to `EpochConfig` instead in that case).
 		NextEpochConfig: Option<BabeEpochConfiguration>;
 	}
 	add_extra_genesis {
@@ -529,7 +530,7 @@ impl<T: Config> Module<T> {
 			duration: T::EpochDuration::get(),
 			authorities: Self::authorities(),
 			randomness: Self::randomness(),
-			config: EpochConfig::get().expect("EpochConfig is never None; qed"),
+			config: EpochConfig::get().expect("EpochConfig is initialized in genesis; we never `take` or `kill` it; qed"),
 		}
 	}
 
@@ -548,7 +549,7 @@ impl<T: Config> Module<T> {
 			authorities: NextAuthorities::get(),
 			randomness: NextRandomness::get(),
 			config: NextEpochConfig::get().unwrap_or_else(|| {
-				EpochConfig::get().expect("EpochConfig is never None; qed")
+				EpochConfig::get().expect("EpochConfig is initialized in genesis; we never `take` or `kill` it; qed"),
 			}),
 		}
 	}

From 74fbbb441c2bd8c3da6d709826bee67f8b3a2a57 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Wed, 3 Mar 2021 09:48:48 +0100
Subject: [PATCH 17/19] Add more asserts to tests, remove unused changes to
 primitives/slots

---
 Cargo.lock                            |  1 -
 frame/babe/src/lib.rs                 | 11 ++++---
 frame/babe/src/tests.rs               | 43 +++++++++++++++++++--------
 primitives/consensus/slots/Cargo.toml |  2 --
 primitives/consensus/slots/src/lib.rs |  3 +-
 5 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 13c78d198c637..6baab384f052c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8499,7 +8499,6 @@ name = "sp-consensus-slots"
 version = "0.9.0"
 dependencies = [
  "parity-scale-codec",
- "serde",
  "sp-arithmetic",
  "sp-runtime",
 ]
diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index bbbac0436dcac..3de2a0d825d5b 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -549,7 +549,7 @@ impl<T: Config> Module<T> {
 			authorities: NextAuthorities::get(),
 			randomness: NextRandomness::get(),
 			config: NextEpochConfig::get().unwrap_or_else(|| {
-				EpochConfig::get().expect("EpochConfig is initialized in genesis; we never `take` or `kill` it; qed"),
+				EpochConfig::get().expect("EpochConfig is initialized in genesis; we never `take` or `kill` it; qed")
 			}),
 		}
 	}
@@ -874,9 +874,8 @@ pub mod migrations {
 		__OldNextEpochConfig, Option<NextConfigDescriptor>, ValueQuery
 	>;
 
-	pub fn add_epoch_configurations<T: Config>(
-		current_epoch_config: BabeEpochConfiguration,
-		next_epoch_config: BabeEpochConfiguration,
+	pub fn add_epoch_configuration<T: Config>(
+		epoch_config: BabeEpochConfiguration,
 	) -> Weight {
 		let mut writes = 0;
 		let mut reads = 0;
@@ -891,8 +890,8 @@ pub mod migrations {
 
 		OldNextEpochConfig::kill();
 
-		EpochConfig::put(current_epoch_config);
-		NextEpochConfig::put(next_epoch_config);
+		EpochConfig::put(epoch_config.clone());
+		NextEpochConfig::put(epoch_config);
 
 		writes += 3;
 
diff --git a/frame/babe/src/tests.rs b/frame/babe/src/tests.rs
index ea7d5a9365dd1..85f64147cc1b9 100644
--- a/frame/babe/src/tests.rs
+++ b/frame/babe/src/tests.rs
@@ -231,11 +231,31 @@ fn can_enact_next_config() {
 		assert_eq!(Babe::epoch_index(), 0);
 		go_to_block(2, 7);
 
+		let current_config = BabeEpochConfiguration {
+			c: (0, 4),
+			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
+		};
+
+		let next_config = BabeEpochConfiguration {
+			c: (1, 4),
+			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
+		};
+
+		let next_next_config = BabeEpochConfiguration {
+			c: (2, 4),
+			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
+		};
+
+		EpochConfig::put(current_config);
+		NextEpochConfig::put(next_config.clone());
+
+		assert_eq!(NextEpochConfig::get(), Some(next_config.clone()));
+
 		Babe::plan_config_change(
 			Origin::root(),
 			NextConfigDescriptor::V1 {
-				c: (1, 4),
-				allowed_slots: AllowedSlots::PrimarySlots,
+				c: next_next_config.c,
+				allowed_slots: next_next_config.allowed_slots,
 			},
 		).unwrap();
 
@@ -243,10 +263,13 @@ fn can_enact_next_config() {
 		Babe::on_finalize(9);
 		let header = System::finalize();
 
+		assert_eq!(EpochConfig::get(), Some(next_config));
+		assert_eq!(NextEpochConfig::get(), Some(next_next_config.clone()));
+
 		let consensus_log = sp_consensus_babe::ConsensusLog::NextConfigData(
-			sp_consensus_babe::digests::NextConfigDescriptor::V1 {
-				c: (1, 4),
-				allowed_slots: AllowedSlots::PrimarySlots,
+			NextConfigDescriptor::V1 {
+				c: next_next_config.c,
+				allowed_slots: next_next_config.allowed_slots,
 			}
 		);
 		let consensus_digest = DigestItem::Consensus(BABE_ENGINE_ID, consensus_log.encode());
@@ -845,13 +868,8 @@ fn add_epoch_configurations_migration_works() {
 			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
 		};
 
-		let next_epoch = BabeEpochConfiguration {
-			c: (2, 4),
-			allowed_slots: sp_consensus_babe::AllowedSlots::PrimarySlots,
-		};
-
-		crate::migrations::add_epoch_configurations::<Test>(
-			current_epoch.clone(), next_epoch.clone()
+		crate::migrations::add_epoch_configuration::<Test>(
+			current_epoch.clone()
 		);
 
 		assert!(get_storage_value::<Option<NextConfigDescriptor>>(
@@ -861,7 +879,6 @@ fn add_epoch_configurations_migration_works() {
 		).is_none());
 
 		assert_eq!(EpochConfig::get(), Some(current_epoch));
-		assert_eq!(NextEpochConfig::get(), Some(next_epoch));
 		assert_eq!(PendingEpochConfigChange::get(), Some(next_config_descriptor));
 	});
 }
diff --git a/primitives/consensus/slots/Cargo.toml b/primitives/consensus/slots/Cargo.toml
index 9167e90523dff..46dbaca1a6ad4 100644
--- a/primitives/consensus/slots/Cargo.toml
+++ b/primitives/consensus/slots/Cargo.toml
@@ -16,7 +16,6 @@ targets = ["x86_64-unknown-linux-gnu"]
 codec = { package = "parity-scale-codec", version = "2.0.0", default-features = false, features = ["derive"] }
 sp-runtime = { version = "3.0.0", default-features = false, path = "../../runtime" }
 sp-arithmetic = { version = "3.0.0", default-features = false, path = "../../arithmetic" }
-serde = { version = "1.0.123", features = ["derive"], default-features = false }
 
 [features]
 default = ["std"]
@@ -24,5 +23,4 @@ std = [
 	"codec/std",
 	"sp-runtime/std",
 	"sp-arithmetic/std",
-	"serde/std"
 ]
diff --git a/primitives/consensus/slots/src/lib.rs b/primitives/consensus/slots/src/lib.rs
index 3e453ccdf4db6..545d18af1f9be 100644
--- a/primitives/consensus/slots/src/lib.rs
+++ b/primitives/consensus/slots/src/lib.rs
@@ -20,10 +20,9 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 use codec::{Decode, Encode};
-use serde::{Serialize, Deserialize};
 
 /// Unit type wrapper that represents a slot.
-#[derive(Debug, Encode, Decode, Eq, Clone, Copy, Default, Ord, Serialize, Deserialize)]
+#[derive(Debug, Encode, Decode, Eq, Clone, Copy, Default, Ord)]
 pub struct Slot(u64);
 
 impl core::ops::Deref for Slot {

From a49a87f4ba31f7f3521824c6e6b9c95826bdfd72 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Tue, 9 Mar 2021 09:03:27 +0100
Subject: [PATCH 18/19] Allow setting the migration pallet prefix

---
 frame/babe/src/lib.rs   | 20 ++++++++++++--------
 frame/babe/src/tests.rs |  6 ++++++
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index 3de2a0d825d5b..46a2cb3fe82d8 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -864,23 +864,27 @@ pub mod migrations {
 	use super::*;
 	use frame_support::pallet_prelude::{ValueQuery, StorageValue};
 
-	struct __OldNextEpochConfig;
-	impl frame_support::traits::StorageInstance for __OldNextEpochConfig {
-		fn pallet_prefix() -> &'static str { "BabeApi" }
+	pub trait HasPalletPrefix: Config {
+		fn pallet_prefix() -> &'static str;
+	}
+
+	struct __OldNextEpochConfig<T>(sp_std::marker::PhantomData<T>);
+	impl<T: HasPalletPrefix> frame_support::traits::StorageInstance for __OldNextEpochConfig<T> {
+		fn pallet_prefix() -> &'static str { T::pallet_prefix() }
 		const STORAGE_PREFIX: &'static str = "NextEpochConfig";
 	}
 
-	type OldNextEpochConfig = StorageValue<
-		__OldNextEpochConfig, Option<NextConfigDescriptor>, ValueQuery
+	type OldNextEpochConfig<T> = StorageValue<
+		__OldNextEpochConfig<T>, Option<NextConfigDescriptor>, ValueQuery
 	>;
 
-	pub fn add_epoch_configuration<T: Config>(
+	pub fn add_epoch_configuration<T: HasPalletPrefix>(
 		epoch_config: BabeEpochConfiguration,
 	) -> Weight {
 		let mut writes = 0;
 		let mut reads = 0;
 
-		if let Some(pending_change) = OldNextEpochConfig::get() {
+		if let Some(pending_change) = OldNextEpochConfig::<T>::get() {
 			PendingEpochConfigChange::put(pending_change);
 
 			writes += 1;
@@ -888,7 +892,7 @@ pub mod migrations {
 
 		reads += 1;
 
-		OldNextEpochConfig::kill();
+		OldNextEpochConfig::<T>::kill();
 
 		EpochConfig::put(epoch_config.clone());
 		NextEpochConfig::put(epoch_config);
diff --git a/frame/babe/src/tests.rs b/frame/babe/src/tests.rs
index 85f64147cc1b9..74d767e22fb3d 100644
--- a/frame/babe/src/tests.rs
+++ b/frame/babe/src/tests.rs
@@ -844,6 +844,12 @@ fn add_epoch_configurations_migration_works() {
 		put_storage_value, get_storage_value,
 	};
 
+	impl crate::migrations::HasPalletPrefix for Test {
+		fn pallet_prefix() -> &'static str {
+			"BabeApi"
+		}
+	}
+
 	new_test_ext(1).execute_with(|| {
 		let next_config_descriptor = NextConfigDescriptor::V1 {
 			c: (3, 4),

From 35ff1cfe1e68623368f88052a6a77b70721915e9 Mon Sep 17 00:00:00 2001
From: Ashley Ruglys <ashley.ruglys@gmail.com>
Date: Tue, 9 Mar 2021 12:17:28 +0100
Subject: [PATCH 19/19] Rename to BabePalletPrefix

---
 frame/babe/src/lib.rs   |  9 ++++++---
 frame/babe/src/tests.rs | 10 +++++-----
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/frame/babe/src/lib.rs b/frame/babe/src/lib.rs
index 46a2cb3fe82d8..29c815444a3a9 100644
--- a/frame/babe/src/lib.rs
+++ b/frame/babe/src/lib.rs
@@ -864,12 +864,13 @@ pub mod migrations {
 	use super::*;
 	use frame_support::pallet_prelude::{ValueQuery, StorageValue};
 
-	pub trait HasPalletPrefix: Config {
+	/// Something that can return the storage prefix of the `Babe` pallet.
+	pub trait BabePalletPrefix: Config {
 		fn pallet_prefix() -> &'static str;
 	}
 
 	struct __OldNextEpochConfig<T>(sp_std::marker::PhantomData<T>);
-	impl<T: HasPalletPrefix> frame_support::traits::StorageInstance for __OldNextEpochConfig<T> {
+	impl<T: BabePalletPrefix> frame_support::traits::StorageInstance for __OldNextEpochConfig<T> {
 		fn pallet_prefix() -> &'static str { T::pallet_prefix() }
 		const STORAGE_PREFIX: &'static str = "NextEpochConfig";
 	}
@@ -878,7 +879,9 @@ pub mod migrations {
 		__OldNextEpochConfig<T>, Option<NextConfigDescriptor>, ValueQuery
 	>;
 
-	pub fn add_epoch_configuration<T: HasPalletPrefix>(
+	/// A storage migration that adds the current epoch configuration for Babe
+	/// to storage.
+	pub fn add_epoch_configuration<T: BabePalletPrefix>(
 		epoch_config: BabeEpochConfiguration,
 	) -> Weight {
 		let mut writes = 0;
diff --git a/frame/babe/src/tests.rs b/frame/babe/src/tests.rs
index 74d767e22fb3d..6515e5fdaaf99 100644
--- a/frame/babe/src/tests.rs
+++ b/frame/babe/src/tests.rs
@@ -844,9 +844,9 @@ fn add_epoch_configurations_migration_works() {
 		put_storage_value, get_storage_value,
 	};
 
-	impl crate::migrations::HasPalletPrefix for Test {
+	impl crate::migrations::BabePalletPrefix for Test {
 		fn pallet_prefix() -> &'static str {
-			"BabeApi"
+			"Babe"
 		}
 	}
 
@@ -857,14 +857,14 @@ fn add_epoch_configurations_migration_works() {
 		};
 
 		put_storage_value(
-			b"BabeApi",
+			b"Babe",
 			b"NextEpochConfig",
 			&[],
 			Some(next_config_descriptor.clone())
 		);
 
 		assert!(get_storage_value::<Option<NextConfigDescriptor>>(
-			b"BabeApi",
+			b"Babe",
 			b"NextEpochConfig",
 			&[],
 		).is_some());
@@ -879,7 +879,7 @@ fn add_epoch_configurations_migration_works() {
 		);
 
 		assert!(get_storage_value::<Option<NextConfigDescriptor>>(
-			b"BabeApi",
+			b"Babe",
 			b"NextEpochConfig",
 			&[],
 		).is_none());