Fixing BABE epochs to change between blocks (#3583)

* always fetch epoch from runtime

* node integration tests don't test light nodes

* give stand-in full node a FULL role

* rejig babe APIs

* introduce next-epoch-descriptor type

* overhaul srml-BABE epoch logic

* ensure VRF outputs end up in the right epoch-randomness

* rewrite `do_initialize` to remove unnecessary loop

* begin accounting for next epoch in epoch function

* slots passes header to epoch_data

* pass slot_number to SlotWorker::epoch_data

* begin extracting epoch-change logic into its own module

* aux methods for block weight

* aux methods for genesis configuration

* comment-out most, refactor header-check pipeline

* mostly flesh out verifier again

* reinstantiate babe BlockImport implementation

* reinstate import-queue instantiation

* reintroduce slot-worker implementation

* reinstate pretty much all the rest

* move fork-choice logic to BlockImport

* fix some, but not all errors

* patch test-runtime

* make is_descendent of slightly more generic

* get skeleton compiling when passing is_descendent_of

* make descendent-of-builder more succinct

* restore ordering of authority_index / slot_number

* start fiddling with tests

* fix warnings

* improve initialization architecture and handle genesis

* tests use correct block-import

* fix BABE tests

* fix some compiler errors

* fix node-cli compilation

* all crates compile

* bump runtime versions and fix some warnings

* tweak fork-tree search implementation

* do backtracking search in fork-tree

* node-cli integration tests now work

* fix broken assumption in test_connectivity

* babe tests fail for the right reasons.

* test genesis epoch logic for epoch_changes

* test that epochs can change between blocks

* First BABE SRML test

* Testing infrastructure for BABE

Also includes a trivial additional test.

* Apply suggestions from code review

Co-Authored-By: Bastian Köcher <bkchr@users.noreply.github.com>

* A little more test progress

* More work on BABE testing

* Try to get the tests working

* Implement `UintAuthorityId`-based test mocks

* Fix compilation errors

* Adjust to upstream changes

* Block numbers are ignored in BABE epoch calculation

* authority_index() should ignore invalid authorities

* Fix compile error

* Add tests that session transitions happen

* Check if BABE produces logs

It currently does not.

* Fix test suite

This was really nasty, due to a type confusion that showed up as an
off-by-1 buffer error.

* Add additional tests

Most of these were derived from the current output, so they are only
useful to guard against regressions.

* Make the tests more readable

Also bump impl_version.

* Fix excessive line width

* Remove unused imports

* Update srml/babe/src/lib.rs

Co-Authored-By: André Silva <andre.beat@gmail.com>

* try to fix imports

* Fix build errors in test suite

* tests did not pass

* Try to get at least one digest to be output

Currently, the code emits either no digests (if I don’t call
`Session::rotate_session()` or two digests (if I do), which is wrong.

* More tests

They still don’t work, but this should help debugging.

* fix silly error

* Don’t even try to compile a broken test

* remove broken check_epoch test and add one for genesis epoch

* Check that the length of the pre-digests is correct

* Bump `impl_version`

* use epoch_for_descendent_of even for genesis

* account for competing block 1s

* finish srml-babe docs

Co-Authored-By: André Silva <andre.beat@gmail.com>

* address grumbles

substrate/core/consensus/babe/src/aux_schema.rs

@@ -22,11 +22,16 @@ use codec::{Decode, Encode};
 use client::backend::AuxStore;
 use client::error::{Result as ClientResult, Error as ClientError};
 use sr_primitives::traits::Block as BlockT;
+use babe_primitives::BabeBlockWeight;
 
-use super::{EpochChanges, SharedEpochChanges};
+use super::{epoch_changes::EpochChangesFor, SharedEpochChanges};
 
 const BABE_EPOCH_CHANGES: &[u8] = b"babe_epoch_changes";
 
+fn block_weight_key<H: Encode>(block_hash: H) -> Vec<u8> {
+	(b"block_weight", block_hash).encode()
+}
+
 fn load_decode<B, T>(backend: &B, key: &[u8]) -> ClientResult<Option<T>>
 	where
 		B: AuxStore,
@@ -45,7 +50,7 @@ fn load_decode<B, T>(backend: &B, key: &[u8]) -> ClientResult<Option<T>>
 pub(crate) fn load_epoch_changes<Block: BlockT, B: AuxStore>(
 	backend: &B,
 ) -> ClientResult<SharedEpochChanges<Block>> {
-	let epoch_changes = load_decode::<_, EpochChanges<Block>>(backend, BABE_EPOCH_CHANGES)?
+	let epoch_changes = load_decode::<_, EpochChangesFor<Block>>(backend, BABE_EPOCH_CHANGES)?
 		.map(Into::into)
 		.unwrap_or_else(|| {
 			info!(target: "babe",
@@ -59,7 +64,7 @@ pub(crate) fn load_epoch_changes<Block: BlockT, B: AuxStore>(
 
 /// Update the epoch changes on disk after a change.
 pub(crate) fn write_epoch_changes<Block: BlockT, F, R>(
-	epoch_changes: &EpochChanges<Block>,
+	epoch_changes: &EpochChangesFor<Block>,
 	write_aux: F,
 ) -> R where
 	F: FnOnce(&[(&'static [u8], &[u8])]) -> R,
@@ -69,3 +74,28 @@ pub(crate) fn write_epoch_changes<Block: BlockT, F, R>(
 		&[(BABE_EPOCH_CHANGES, encoded_epoch_changes.as_slice())],
 	)
 }
+
+/// Write the cumulative chain-weight of a block ot aux storage.
+pub(crate) fn write_block_weight<H: Encode, F, R>(
+	block_hash: H,
+	block_weight: &BabeBlockWeight,
+	write_aux: F,
+) -> R where
+	F: FnOnce(&[(Vec<u8>, &[u8])]) -> R,
+{
+
+	let key = block_weight_key(block_hash);
+	block_weight.using_encoded(|s|
+		write_aux(
+			&[(key, s)],
+		)
+	)
+}
+
+/// Load the cumulative chain-weight associated with a block.
+pub(crate) fn load_block_weight<H: Encode, B: AuxStore>(
+	backend: &B,
+	block_hash: H,
+) -> ClientResult<Option<BabeBlockWeight>> {
+	load_decode(backend, block_weight_key(block_hash).as_slice())
+}

substrate/core/consensus/babe/src/epoch_changes.rs

@@ -0,0 +1,637 @@
+// Copyright 2019 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Handling epoch changes in BABE.
+//!
+//! This exposes the `SharedEpochChanges`, which is a wrapper around a
+//! persistent DAG superimposed over the forks of the blockchain.
+
+use std::sync::Arc;
+use babe_primitives::{Epoch, SlotNumber, NextEpochDescriptor};
+use fork_tree::ForkTree;
+use parking_lot::{Mutex, MutexGuard};
+use sr_primitives::traits::{Block as BlockT, NumberFor, One, Zero};
+use codec::{Encode, Decode};
+use client::error::Error as ClientError;
+use client::utils as client_utils;
+use client::blockchain::HeaderBackend;
+use primitives::H256;
+use std::ops::Add;
+
+/// A builder for `is_descendent_of` functions.
+pub trait IsDescendentOfBuilder<Hash> {
+	/// The error returned by the function.
+	type Error: std::error::Error;
+	/// A function that can tell you if the second parameter is a descendent of
+	/// the first.
+	type IsDescendentOf: Fn(&Hash, &Hash) -> Result<bool, Self::Error>;
+
+	/// Build an `is_descendent_of` function.
+	///
+	/// The `current` parameter can be `Some` with the details a fresh block whose
+	/// details aren't yet stored, but its parent is.
+	///
+	/// The format of `current` when `Some` is `(current, current_parent)`.
+	fn build_is_descendent_of(&self, current: Option<(Hash, Hash)>)
+		-> Self::IsDescendentOf;
+}
+
+/// Produce a descendent query object given the client.
+pub(crate) fn descendent_query<H, Block>(client: &H) -> HeaderBackendDescendentBuilder<&H, Block> {
+	HeaderBackendDescendentBuilder(client, std::marker::PhantomData)
+}
+
+/// Wrapper to get around unconstrained type errors when implementing
+/// `IsDescendentOfBuilder` for header backends.
+pub(crate) struct HeaderBackendDescendentBuilder<H, Block>(H, std::marker::PhantomData<Block>);
+
+// TODO: relying on Hash = H256 is awful.
+// https://github.com/paritytech/substrate/issues/3624
+impl<'a, H, Block> IsDescendentOfBuilder<H256>
+	for HeaderBackendDescendentBuilder<&'a H, Block> where
+	H: HeaderBackend<Block>,
+	Block: BlockT<Hash = H256>,
+{
+	type Error = ClientError;
+	type IsDescendentOf = Box<dyn Fn(&H256, &H256) -> Result<bool, ClientError> + 'a>;
+
+	fn build_is_descendent_of(&self, current: Option<(H256, H256)>)
+		-> Self::IsDescendentOf
+	{
+		Box::new(client_utils::is_descendent_of(self.0, current))
+	}
+}
+
+/// An unimported genesis epoch.
+pub struct UnimportedGenesis(Epoch);
+
+/// The viable epoch under which a block can be verified.
+///
+/// If this is the first non-genesis block in the chain, then it will
+/// hold an `UnimportedGenesis` epoch.
+pub enum ViableEpoch {
+	Genesis(UnimportedGenesis),
+	Regular(Epoch),
+}
+
+impl From<Epoch> for ViableEpoch {
+	fn from(epoch: Epoch) -> ViableEpoch {
+		ViableEpoch::Regular(epoch)
+	}
+}
+
+impl AsRef<Epoch> for ViableEpoch {
+	fn as_ref(&self) -> &Epoch {
+		match *self {
+			ViableEpoch::Genesis(UnimportedGenesis(ref e)) => e,
+			ViableEpoch::Regular(ref e) => e,
+		}
+	}
+}
+
+impl ViableEpoch {
+	/// Extract the underlying epoch, disregarding the fact that a genesis
+	/// epoch may be unimported.
+	pub fn into_inner(self) -> Epoch {
+		match self {
+			ViableEpoch::Genesis(UnimportedGenesis(e)) => e,
+			ViableEpoch::Regular(e) => e,
+		}
+	}
+
+	/// Increment the epoch, yielding an `IncrementedEpoch` to be imported
+	/// into the fork-tree.
+	pub fn increment(&self, next_descriptor: NextEpochDescriptor) -> IncrementedEpoch {
+		let next = self.as_ref().increment(next_descriptor);
+		let to_persist = match *self {
+			ViableEpoch::Genesis(UnimportedGenesis(ref epoch_0)) =>
+				PersistedEpoch::Genesis(epoch_0.clone(), next),
+			ViableEpoch::Regular(_) => PersistedEpoch::Regular(next),
+		};
+
+		IncrementedEpoch(to_persist)
+	}
+}
+
+/// The datatype encoded on disk.
+// This really shouldn't be public, but the encode/decode derives force it to be.
+#[derive(Clone, Encode, Decode)]
+pub enum PersistedEpoch {
+	// epoch_0, epoch_1,
+	Genesis(Epoch, Epoch),
+	// epoch_n
+	Regular(Epoch),
+}
+
+/// A fresh, incremented epoch to import into the underlying fork-tree.
+///
+/// Create this with `ViableEpoch::increment`.
+#[must_use = "Freshly-incremented epoch must be imported with `EpochChanges::import`"]
+pub struct IncrementedEpoch(PersistedEpoch);
+
+impl AsRef<Epoch> for IncrementedEpoch {
+	fn as_ref(&self) -> &Epoch {
+		match self.0 {
+			PersistedEpoch::Genesis(_, ref epoch_1) => epoch_1,
+			PersistedEpoch::Regular(ref epoch_n) => epoch_n,
+		}
+	}
+}
+
+/// Tree of all epoch changes across all *seen* forks. Data stored in tree is
+/// the hash and block number of the block signaling the epoch change, and the
+/// epoch that was signalled at that block.
+///
+/// BABE special-cases the first epoch, epoch_0, by saying that it starts at
+/// slot number of the first block in the chain. When bootstrapping a chain,
+/// there can be multiple competing block #1s, so we have to ensure that the overlayed
+/// DAG doesn't get confused.
+///
+/// The first block of every epoch should be producing a descriptor for the next
+/// epoch - this is checked in higher-level code. So the first block of epoch_0 contains
+/// a descriptor for epoch_1. We special-case these and bundle them together in the
+/// same DAG entry, pinned to a specific block #1.
+///
+/// Further epochs (epoch_2, ..., epoch_n) each get their own entry.
+#[derive(Clone, Encode, Decode)]
+pub struct EpochChanges<Hash, Number> {
+	inner: ForkTree<Hash, Number, PersistedEpoch>,
+}
+
+// create a fake header hash which hasn't been included in the chain.
+fn fake_head_hash<H: AsRef<[u8]> + AsMut<[u8]> + Clone>(parent_hash: &H) -> H {
+	let mut h = parent_hash.clone();
+	// dirty trick: flip the first bit of the parent hash to create a hash
+	// which has not been in the chain before (assuming a strong hash function).
+	h.as_mut()[0] ^= 0b10000000;
+	h
+}
+
+impl<Hash, Number> EpochChanges<Hash, Number> where
+	Hash: PartialEq + AsRef<[u8]> + AsMut<[u8]> + Copy,
+	Number: Ord + One + Zero + Add<Output=Number> + Copy,
+{
+	/// Create a new epoch-change tracker.
+	fn new() -> Self {
+		EpochChanges { inner: ForkTree::new() }
+	}
+
+	/// Prune out finalized epochs, except for the ancestor of the finalized block.
+	pub fn prune_finalized<D: IsDescendentOfBuilder<Hash>>(
+		&mut self,
+		descendent_of_builder: D,
+		_hash: &Hash,
+		_number: Number,
+	) -> Result<(), fork_tree::Error<D::Error>> {
+		let _is_descendent_of = descendent_of_builder
+			.build_is_descendent_of(None);
+
+		// TODO:
+		// https://github.com/paritytech/substrate/issues/3651
+		//
+		// prune any epochs which could not be _live_ as of the children of the
+		// finalized block.
+		// i.e. re-root the fork tree to the oldest ancestor of (hash, number)
+		// where epoch.end_slot() >= slot(hash)
+
+		Ok(())
+	}
+
+	/// Finds the epoch for a child of the given block, assuming the given slot number.
+	///
+	/// If the returned epoch is an `UnimportedGenesis` epoch, it should be imported into the
+	/// tree.
+	pub fn epoch_for_child_of<D: IsDescendentOfBuilder<Hash>, G>(
+		&self,
+		descendent_of_builder: D,
+		parent_hash: &Hash,
+		parent_number: Number,
+		slot_number: SlotNumber,
+		make_genesis: G,
+	) -> Result<Option<ViableEpoch>, fork_tree::Error<D::Error>>
+		where G: FnOnce(SlotNumber) -> Epoch
+	{
+		// find_node_where will give you the node in the fork-tree which is an ancestor
+		// of the `parent_hash` by default. if the last epoch was signalled at the parent_hash,
+		// then it won't be returned. we need to create a new fake chain head hash which
+		// "descends" from our parent-hash.
+		let fake_head_hash = fake_head_hash(parent_hash);
+
+		let is_descendent_of = descendent_of_builder
+			.build_is_descendent_of(Some((fake_head_hash, *parent_hash)));
+
+		if parent_number == Zero::zero() {
+			// need to insert the genesis epoch.
+			let genesis_epoch = make_genesis(slot_number);
+			return Ok(Some(ViableEpoch::Genesis(UnimportedGenesis(genesis_epoch))));
+		}
+
+		// We want to find the deepest node in the tree which is an ancestor
+		// of our block and where the start slot of the epoch was before the
+		// slot of our block. The genesis special-case doesn't need to look
+		// at epoch_1 -- all we're doing here is figuring out which node
+		// we need.
+		let predicate = |epoch: &PersistedEpoch| match *epoch {
+			PersistedEpoch::Genesis(ref epoch_0, _) =>
+				epoch_0.start_slot <= slot_number,
+			PersistedEpoch::Regular(ref epoch_n) =>
+				epoch_n.start_slot <= slot_number,
+		};
+
+		self.inner.find_node_where(
+			&fake_head_hash,
+			&(parent_number + One::one()),
+			&is_descendent_of,
+			&predicate,
+		)
+			.map(|n| n.map(|node| ViableEpoch::Regular(match node.data {
+				// Ok, we found our node.
+				// and here we figure out which of the internal epochs
+				// of a genesis node to use based on their start slot.
+				PersistedEpoch::Genesis(ref epoch_0, ref epoch_1) =>
+					if epoch_1.start_slot <= slot_number {
+						epoch_1.clone()
+					} else {
+						epoch_0.clone()
+					},
+				PersistedEpoch::Regular(ref epoch_n) => epoch_n.clone(),
+			})))
+	}
+
+	/// Import a new epoch-change, signalled at the given block.
+	///
+	/// This assumes that the given block is prospective (i.e. has not been
+	/// imported yet), but its parent has. This is why the parent hash needs
+	/// to be provided.
+	pub fn import<D: IsDescendentOfBuilder<Hash>>(
+		&mut self,
+		descendent_of_builder: D,
+		hash: Hash,
+		number: Number,
+		parent_hash: Hash,
+		epoch: IncrementedEpoch,
+	) -> Result<(), fork_tree::Error<D::Error>> {
+		let is_descendent_of = descendent_of_builder
+			.build_is_descendent_of(Some((hash, parent_hash)));
+
+		let res = self.inner.import(
+			hash,
+			number,
+			epoch.0,
+			&is_descendent_of,
+		);
+
+		match res {
+			Ok(_) | Err(fork_tree::Error::Duplicate) => Ok(()),
+			Err(e) => Err(e),
+		}
+	}
+}
+
+/// Type alias to produce the epoch-changes tree from a block type.
+pub type EpochChangesFor<Block> = EpochChanges<<Block as BlockT>::Hash, NumberFor<Block>>;
+
+/// A shared epoch changes tree.
+#[derive(Clone)]
+pub struct SharedEpochChanges<Block: BlockT> {
+	inner: Arc<Mutex<EpochChangesFor<Block>>>,
+}
+
+impl<Block: BlockT> SharedEpochChanges<Block> {
+	/// Create a new instance of the `SharedEpochChanges`.
+	pub fn new() -> Self {
+		SharedEpochChanges {
+			inner: Arc::new(Mutex::new(EpochChanges::<_, _>::new()))
+		}
+	}
+
+	/// Lock the shared epoch changes,
+	pub fn lock(&self) -> MutexGuard<EpochChangesFor<Block>> {
+		self.inner.lock()
+	}
+}
+
+impl<Block: BlockT> From<EpochChangesFor<Block>> for SharedEpochChanges<Block> {
+	fn from(epoch_changes: EpochChangesFor<Block>) -> Self {
+		SharedEpochChanges {
+			inner: Arc::new(Mutex::new(epoch_changes))
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[derive(Debug, PartialEq)]
+	pub struct TestError;
+
+	impl std::fmt::Display for TestError {
+		fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+			write!(f, "TestError")
+		}
+	}
+
+	impl std::error::Error for TestError {}
+
+	impl<'a, F: 'a , H: 'a + PartialEq + std::fmt::Debug> IsDescendentOfBuilder<H> for &'a F
+		where F: Fn(&H, &H) -> Result<bool, TestError>
+	{
+		type Error = TestError;
+		type IsDescendentOf = Box<dyn Fn(&H, &H) -> Result<bool, TestError> + 'a>;
+
+		fn build_is_descendent_of(&self, current: Option<(H, H)>)
+			-> Self::IsDescendentOf
+		{
+			let f = *self;
+			Box::new(move |base, head| {
+				let mut head = head;
+
+				if let Some((ref c_head, ref c_parent)) = current {
+					if head == c_head {
+						if base == c_parent {
+							return Ok(true);
+						} else {
+							head = c_parent;
+						}
+					}
+				}
+
+				f(base, head)
+			})
+		}
+	}
+
+	type Hash = [u8; 1];
+
+	#[test]
+	fn genesis_epoch_is_created_but_not_imported() {
+		//
+		// A - B
+		//  \
+		//   — C
+		//
+		let is_descendent_of = |base: &Hash, block: &Hash| -> Result<bool, TestError> {
+			match (base, *block) {
+				(b"A", b) => Ok(b == *b"B" || b == *b"C" || b == *b"D"),
+				(b"B", b) | (b"C", b) => Ok(b == *b"D"),
+				(b"0", _) => Ok(true),
+				_ => Ok(false),
+			}
+		};
+
+		let make_genesis = |slot| Epoch {
+			epoch_index: 0,
+			start_slot: slot,
+			duration: 100,
+			authorities: Vec::new(),
+			randomness: [0; 32],
+		};
+
+		let epoch_changes = EpochChanges::new();
+		let genesis_epoch = epoch_changes.epoch_for_child_of(
+			&is_descendent_of,
+			b"0",
+			0,
+			10101,
+			&make_genesis,
+		).unwrap().unwrap();
+
+		match genesis_epoch {
+			ViableEpoch::Genesis(_) => {},
+			_ => panic!("should be unimported genesis"),
+		};
+		assert_eq!(genesis_epoch.as_ref(), &make_genesis(10101));
+
+		let genesis_epoch_2 = epoch_changes.epoch_for_child_of(
+			&is_descendent_of,
+			b"0",
+			0,
+			10102,
+			&make_genesis,
+		).unwrap().unwrap();
+
+		match genesis_epoch_2 {
+			ViableEpoch::Genesis(_) => {},
+			_ => panic!("should be unimported genesis"),
+		};
+		assert_eq!(genesis_epoch_2.as_ref(), &make_genesis(10102));
+	}
+
+	#[test]
+	fn epoch_changes_between_blocks() {
+		//
+		// A - B
+		//  \
+		//   — C
+		//
+		let is_descendent_of = |base: &Hash, block: &Hash| -> Result<bool, TestError> {
+			match (base, *block) {
+				(b"A", b) => Ok(b == *b"B" || b == *b"C" || b == *b"D"),
+				(b"B", b) | (b"C", b) => Ok(b == *b"D"),
+				(b"0", _) => Ok(true),
+				_ => Ok(false),
+			}
+		};
+
+		let make_genesis = |slot| Epoch {
+			epoch_index: 0,
+			start_slot: slot,
+			duration: 100,
+			authorities: Vec::new(),
+			randomness: [0; 32],
+		};
+
+		let mut epoch_changes = EpochChanges::new();
+		let genesis_epoch = epoch_changes.epoch_for_child_of(
+			&is_descendent_of,
+			b"0",
+			0,
+			100,
+			&make_genesis,
+		).unwrap().unwrap();
+
+		assert_eq!(genesis_epoch.as_ref(), &make_genesis(100));
+
+		let import_epoch_1 = genesis_epoch.increment(NextEpochDescriptor {
+			authorities: Vec::new(),
+			randomness: [1; 32],
+		});
+		let epoch_1 = import_epoch_1.as_ref().clone();
+
+		epoch_changes.import(
+			&is_descendent_of,
+			*b"A",
+			1,
+			*b"0",
+			import_epoch_1,
+		).unwrap();
+		let genesis_epoch = genesis_epoch.into_inner();
+
+		assert!(is_descendent_of(b"0", b"A").unwrap());
+
+		let end_slot = genesis_epoch.end_slot();
+		assert_eq!(end_slot, epoch_1.start_slot);
+
+		{
+			// x is still within the genesis epoch.
+			let x = epoch_changes.epoch_for_child_of(
+				&is_descendent_of,
+				b"A",
+				1,
+				end_slot - 1,
+				&make_genesis,
+			).unwrap().unwrap().into_inner();
+
+			assert_eq!(x, genesis_epoch);
+		}
+
+		{
+			// x is now at the next epoch, because the block is now at the
+			// start slot of epoch 1.
+			let x = epoch_changes.epoch_for_child_of(
+				&is_descendent_of,
+				b"A",
+				1,
+				end_slot,
+				&make_genesis,
+			).unwrap().unwrap().into_inner();
+
+			assert_eq!(x, epoch_1);
+		}
+
+		{
+			// x is now at the next epoch, because the block is now after
+			// start slot of epoch 1.
+			let x = epoch_changes.epoch_for_child_of(
+				&is_descendent_of,
+				b"A",
+				1,
+				epoch_1.end_slot() - 1,
+				&make_genesis,
+			).unwrap().unwrap().into_inner();
+
+			assert_eq!(x, epoch_1);
+		}
+	}
+
+	#[test]
+	fn two_block_ones_dont_conflict() {
+		//     X - Y
+		//   /
+		// 0 - A - B
+		//
+		let is_descendent_of = |base: &Hash, block: &Hash| -> Result<bool, TestError> {
+			match (base, *block) {
+				(b"A", b) => Ok(b == *b"B"),
+				(b"X", b) => Ok(b == *b"Y"),
+				(b"0", _) => Ok(true),
+				_ => Ok(false),
+			}
+		};
+
+		let duration = 100;
+
+		let make_genesis = |slot| Epoch {
+			epoch_index: 0,
+			start_slot: slot,
+			duration,
+			authorities: Vec::new(),
+			randomness: [0; 32],
+		};
+
+		let mut epoch_changes = EpochChanges::new();
+		let next_descriptor = NextEpochDescriptor {
+			authorities: Vec::new(),
+			randomness: [0; 32],
+		};
+
+		// insert genesis epoch for A
+		{
+			let genesis_epoch_a = epoch_changes.epoch_for_child_of(
+				&is_descendent_of,
+				b"0",
+				0,
+				100,
+				&make_genesis,
+			).unwrap().unwrap();
+
+			epoch_changes.import(
+				&is_descendent_of,
+				*b"A",
+				1,
+				*b"0",
+				genesis_epoch_a.increment(next_descriptor.clone()),
+			).unwrap();
+
+		}
+
+		// insert genesis epoch for X
+		{
+			let genesis_epoch_x = epoch_changes.epoch_for_child_of(
+				&is_descendent_of,
+				b"0",
+				0,
+				1000,
+				&make_genesis,
+			).unwrap().unwrap();
+
+			epoch_changes.import(
+				&is_descendent_of,
+				*b"X",
+				1,
+				*b"0",
+				genesis_epoch_x.increment(next_descriptor.clone()),
+			).unwrap();
+		}
+
+		// now check that the genesis epochs for our respective block 1s
+		// respect the chain structure.
+		{
+			let epoch_for_a_child = epoch_changes.epoch_for_child_of(
+				&is_descendent_of,
+				b"A",
+				1,
+				101,
+				&make_genesis,
+			).unwrap().unwrap();
+
+			assert_eq!(epoch_for_a_child.into_inner(), make_genesis(100));
+
+			let epoch_for_x_child = epoch_changes.epoch_for_child_of(
+				&is_descendent_of,
+				b"X",
+				1,
+				1001,
+				&make_genesis,
+			).unwrap().unwrap();
+
+			assert_eq!(epoch_for_x_child.into_inner(), make_genesis(1000));
+
+			let epoch_for_x_child_before_genesis = epoch_changes.epoch_for_child_of(
+				&is_descendent_of,
+				b"X",
+				1,
+				101,
+				&make_genesis,
+			).unwrap();
+
+			// even though there is a genesis epoch at that slot, it's not in
+			// this chain.
+			assert!(epoch_for_x_child_before_genesis.is_none());
+		}
+	}
+}

substrate/core/consensus/babe/src/tests.rs

@@ -21,19 +21,22 @@
 #![allow(deprecated)]
 use super::*;
 
-use babe_primitives::AuthorityPair;
+use babe_primitives::{AuthorityPair, SlotNumber};
 use client::block_builder::BlockBuilder;
 use consensus_common::NoNetwork as DummyOracle;
+use consensus_common::import_queue::{
+	BoxBlockImport, BoxJustificationImport, BoxFinalityProofImport,
+};
 use network::test::*;
 use network::test::{Block as TestBlock, PeersClient};
+use network::config::BoxFinalityProofRequestBuilder;
 use sr_primitives::{generic::DigestItem, traits::{Block as BlockT, DigestFor}};
 use network::config::ProtocolConfig;
 use tokio::runtime::current_thread;
-use keyring::sr25519::Keyring;
 use client::BlockchainEvents;
 use test_client;
 use log::debug;
-use std::{time::Duration, borrow::Borrow, cell::RefCell};
+use std::{time::Duration, cell::RefCell};
 
 type Item = DigestItem<Hash>;
 
@@ -46,8 +49,28 @@ type TestClient = client::Client<
 	test_client::runtime::RuntimeApi,
 >;
 
-struct DummyFactory(Arc<TestClient>);
-struct DummyProposer(u64, Arc<TestClient>);
+#[derive(Copy, Clone, PartialEq)]
+enum Stage {
+	PreSeal,
+	PostSeal,
+}
+
+type Mutator = Arc<dyn Fn(&mut TestHeader, Stage) + Send + Sync>;
+
+#[derive(Clone)]
+struct DummyFactory {
+	client: Arc<TestClient>,
+	epoch_changes: crate::SharedEpochChanges<TestBlock>,
+	config: Config,
+	mutator: Mutator,
+}
+
+struct DummyProposer {
+	factory: DummyFactory,
+	parent_hash: Hash,
+	parent_number: u64,
+	parent_slot: SlotNumber,
+}
 
 impl Environment<TestBlock> for DummyFactory {
 	type Proposer = DummyProposer;
@@ -56,7 +79,69 @@ impl Environment<TestBlock> for DummyFactory {
 	fn init(&mut self, parent_header: &<TestBlock as BlockT>::Header)
 		-> Result<DummyProposer, Error>
 	{
-		Ok(DummyProposer(parent_header.number + 1, self.0.clone()))
+
+		let parent_slot = crate::find_pre_digest::<TestBlock>(parent_header)
+			.expect("parent header has a pre-digest")
+			.slot_number();
+
+		Ok(DummyProposer {
+			factory: self.clone(),
+			parent_hash: parent_header.hash(),
+			parent_number: *parent_header.number(),
+			parent_slot,
+		})
+	}
+}
+
+impl DummyProposer {
+	fn propose_with(&mut self, pre_digests: DigestFor<TestBlock>)
+		-> future::Ready<Result<TestBlock, Error>>
+	{
+		let block_builder = self.factory.client.new_block_at(
+			&BlockId::Hash(self.parent_hash),
+			pre_digests,
+		).unwrap();
+		let mut block = match block_builder.bake().map_err(|e| e.into()) {
+			Ok(b) => b,
+			Err(e) => return future::ready(Err(e)),
+		};
+
+		let this_slot = crate::find_pre_digest::<TestBlock>(block.header())
+			.expect("baked block has valid pre-digest")
+			.slot_number();
+
+		// figure out if we should add a consensus digest, since the test runtime
+		// doesn't.
+		let epoch_changes = self.factory.epoch_changes.lock();
+		let epoch = epoch_changes.epoch_for_child_of(
+			descendent_query(&*self.factory.client),
+			&self.parent_hash,
+			self.parent_number,
+			this_slot,
+			|slot| self.factory.config.genesis_epoch(slot),
+		)
+			.expect("client has data to find epoch")
+			.expect("can compute epoch for baked block")
+			.into_inner();
+
+		let first_in_epoch = self.parent_slot < epoch.start_slot;
+		if first_in_epoch {
+			// push a `Consensus` digest signalling next change.
+			// we just reuse the same randomness and authorities as the prior
+			// epoch. this will break when we add light client support, since
+			// that will re-check the randomness logic off-chain.
+			let digest_data = ConsensusLog::NextEpochData(NextEpochDescriptor {
+				authorities: epoch.authorities.clone(),
+				randomness: epoch.randomness.clone(),
+			}).encode();
+			let digest = DigestItem::Consensus(BABE_ENGINE_ID, digest_data);
+			block.header.digest_mut().push(digest)
+		}
+
+		// mutate the block header according to the mutator.
+		(self.factory.mutator)(&mut block.header, Stage::PreSeal);
+
+		future::ready(Ok(block))
 	}
 }
 
@@ -67,21 +152,42 @@ impl Proposer<TestBlock> for DummyProposer {
 	fn propose(
 		&mut self,
 		_: InherentData,
-		digests: DigestFor<TestBlock>,
+		pre_digests: DigestFor<TestBlock>,
 		_: Duration,
 	) -> Self::Create {
-		future::ready(self.1.new_block(digests).unwrap().bake().map_err(|e| e.into()))
+		self.propose_with(pre_digests)
 	}
 }
 
-type Mutator = Arc<dyn for<'r> Fn(&'r mut TestHeader) + Send + Sync>;
-
 thread_local! {
-	static MUTATOR: RefCell<Mutator> = RefCell::new(Arc::new(|_|()));
+	static MUTATOR: RefCell<Mutator> = RefCell::new(Arc::new(|_, _|()));
+}
+
+#[derive(Clone)]
+struct PanickingBlockImport<B>(B);
+
+impl<B: BlockImport<TestBlock>> BlockImport<TestBlock> for PanickingBlockImport<B> {
+	type Error = B::Error;
+
+	fn import_block(
+		&mut self,
+		block: BlockImportParams<TestBlock>,
+		new_cache: HashMap<CacheKeyId, Vec<u8>>,
+	) -> Result<ImportResult, Self::Error> {
+		Ok(self.0.import_block(block, new_cache).expect("importing block failed"))
+	}
+
+	fn check_block(
+		&mut self,
+		hash: Hash,
+		parent_hash: Hash,
+	) -> Result<ImportResult, Self::Error> {
+		Ok(self.0.check_block(hash, parent_hash).expect("checking block failed"))
+	}
 }
 
 pub struct BabeTestNet {
-	peers: Vec<Peer<(), DummySpecialization>>,
+	peers: Vec<Peer<Option<PeerData>, DummySpecialization>>,
 }
 
 type TestHeader = <TestBlock as BlockT>::Header;
@@ -94,7 +200,6 @@ pub struct TestVerifier {
 		TestBlock,
 		test_client::runtime::RuntimeApi,
 		PeersFullClient,
-		(),
 	>,
 	mutator: Mutator,
 }
@@ -110,16 +215,22 @@ impl Verifier<TestBlock> for TestVerifier {
 		justification: Option<Justification>,
 		body: Option<Vec<TestExtrinsic>>,
 	) -> Result<(BlockImportParams<TestBlock>, Option<Vec<(CacheKeyId, Vec<u8>)>>), String> {
-		let cb: &(dyn Fn(&mut TestHeader) + Send + Sync) = self.mutator.borrow();
-		cb(&mut header);
+		// apply post-sealing mutations (i.e. stripping seal, if desired).
+		(self.mutator)(&mut header, Stage::PostSeal);
 		Ok(self.inner.verify(origin, header, justification, body).expect("verification failed!"))
 	}
 }
 
+pub struct PeerData {
+	link: BabeLink<TestBlock>,
+	inherent_data_providers: InherentDataProviders,
+	block_import: Mutex<Option<BoxBlockImport<TestBlock>>>,
+}
+
 impl TestNetFactory for BabeTestNet {
 	type Specialization = DummySpecialization;
 	type Verifier = TestVerifier;
-	type PeerData = ();
+	type PeerData = Option<PeerData>;
 
 	/// Create new test network with peers and given config.
 	fn from_config(_config: &ProtocolConfig) -> Self {
@@ -129,31 +240,62 @@ impl TestNetFactory for BabeTestNet {
 		}
 	}
 
-	/// KLUDGE: this function gets the mutator from thread-local storage.
-	fn make_verifier(&self, client: PeersClient, _cfg: &ProtocolConfig)
+	fn make_block_import(&self, client: PeersClient)
+		-> (
+			BoxBlockImport<Block>,
+			Option<BoxJustificationImport<Block>>,
+			Option<BoxFinalityProofImport<Block>>,
+			Option<BoxFinalityProofRequestBuilder<Block>>,
+			Option<PeerData>,
+		)
+	{
+		let client = client.as_full().expect("only full clients are tested");
+		let inherent_data_providers = InherentDataProviders::new();
+
+		let config = Config::get_or_compute(&*client).expect("config available");
+		let (block_import, link) = crate::block_import(
+			config,
+			client.clone(),
+			client.clone(),
+			client.clone(),
+		).expect("can initialize block-import");
+
+		let block_import = PanickingBlockImport(block_import);
+
+		let data_block_import = Mutex::new(Some(Box::new(block_import.clone()) as BoxBlockImport<_>));
+		(
+			Box::new(block_import),
+			None,
+			None,
+			None,
+			Some(PeerData { link, inherent_data_providers, block_import: data_block_import }),
+		)
+	}
+
+	fn make_verifier(
+		&self,
+		client: PeersClient,
+		_cfg: &ProtocolConfig,
+		maybe_link: &Option<PeerData>,
+	)
 		-> Self::Verifier
 	{
 		let client = client.as_full().expect("only full clients are used in test");
 		trace!(target: "babe", "Creating a verifier");
-		let config = Config::get_or_compute(&*client)
-			.expect("slot duration available");
-		let inherent_data_providers = InherentDataProviders::new();
-		register_babe_inherent_data_provider(
-			&inherent_data_providers,
-			config.get()
-		).expect("Registers babe inherent data provider");
-		trace!(target: "babe", "Provider registered");
+
+		// ensure block import and verifier are linked correctly.
+		let data = maybe_link.as_ref().expect("babe link always provided to verifier instantiation");
 
 		TestVerifier {
 			inner: BabeVerifier {
 				client: client.clone(),
 				api: client,
-				inherent_data_providers,
-				config,
-				time_source: Default::default(),
-				transaction_pool : Default::default(),
+				inherent_data_providers: data.inherent_data_providers.clone(),
+				config: data.link.config.clone(),
+				epoch_changes: data.link.epoch_changes.clone(),
+				time_source: data.link.time_source.clone(),
 			},
-			mutator: MUTATOR.with(|s| s.borrow().clone()),
+			mutator: MUTATOR.with(|m| m.borrow().clone()),
 		}
 	}
 
@@ -188,8 +330,13 @@ fn rejects_empty_block() {
 	})
 }
 
-fn run_one_test() {
+fn run_one_test(
+	mutator: impl Fn(&mut TestHeader, Stage) + Send + Sync + 'static,
+) {
 	let _ = env_logger::try_init();
+	let mutator = Arc::new(mutator) as Mutator;
+
+	MUTATOR.with(|m| *m.borrow_mut() = mutator.clone());
 	let net = BabeTestNet::new(3);
 
 	let peers = &[
@@ -202,6 +349,7 @@ fn run_one_test() {
 	let mut import_notifications = Vec::new();
 	let mut runtime = current_thread::Runtime::new().unwrap();
 	let mut keystore_paths = Vec::new();
+
 	for (peer_id, seed) in peers {
 		let mut net = net.lock();
 		let peer = net.peer(*peer_id);
@@ -213,30 +361,46 @@ fn run_one_test() {
 		keystore.write().insert_ephemeral_from_seed::<AuthorityPair>(seed).expect("Generates authority key");
 		keystore_paths.push(keystore_path);
 
-		let environ = DummyFactory(client.clone());
+		let mut got_own = false;
+		let mut got_other = false;
+
+		let data = peer.data.as_ref().expect("babe link set up during initialization");
+
+		let environ = DummyFactory {
+			client: client.clone(),
+			config: data.link.config.clone(),
+			epoch_changes: data.link.epoch_changes.clone(),
+			mutator: mutator.clone(),
+		};
+
 		import_notifications.push(
+			// run each future until we get one of our own blocks with number higher than 5
+			// that was produced locally.
 			client.import_notification_stream()
-				.take_while(|n| future::ready(!(n.origin != BlockOrigin::Own && n.header.number() < &5)))
-				.for_each(move |_| future::ready(()))
+				.take_while(move |n| future::ready(n.header.number() < &5 || {
+					if n.origin == BlockOrigin::Own {
+						got_own = true;
+					} else {
+						got_other = true;
+					}
+
+					// continue until we have at least one block of our own
+					// and one of another peer.
+					!(got_own && got_other)
+				}))
+				.for_each(|_| future::ready(()) )
 		);
 
-		let config = Config::get_or_compute(&*client).expect("slot duration available");
-
-		let inherent_data_providers = InherentDataProviders::new();
-		register_babe_inherent_data_provider(
-			&inherent_data_providers, config.get()
-		).expect("Registers babe inherent data provider");
 
 		runtime.spawn(start_babe(BabeParams {
-			config,
-			block_import: client.clone(),
+			block_import: data.block_import.lock().take().expect("import set up during init"),
 			select_chain,
 			client,
 			env: environ,
 			sync_oracle: DummyOracle,
-			inherent_data_providers,
+			inherent_data_providers: data.inherent_data_providers.clone(),
 			force_authoring: false,
-			time_source: Default::default(),
+			babe_link: data.link.clone(),
 			keystore,
 		}).expect("Starts babe"));
 	}
@@ -251,45 +415,41 @@ fn run_one_test() {
 }
 
 #[test]
-fn authoring_blocks() { run_one_test() }
+fn authoring_blocks() {
+	run_one_test(|_, _| ())
+}
 
 #[test]
 #[should_panic]
 fn rejects_missing_inherent_digest() {
-	MUTATOR.with(|s| *s.borrow_mut() = Arc::new(move |header: &mut TestHeader| {
+	run_one_test(|header: &mut TestHeader, stage| {
 		let v = std::mem::replace(&mut header.digest_mut().logs, vec![]);
 		header.digest_mut().logs = v.into_iter()
-			.filter(|v| v.as_babe_pre_digest().is_none())
+			.filter(|v| stage == Stage::PostSeal || v.as_babe_pre_digest().is_none())
 			.collect()
-	}));
-	run_one_test()
+	})
 }
 
 #[test]
 #[should_panic]
 fn rejects_missing_seals() {
-	MUTATOR.with(|s| *s.borrow_mut() = Arc::new(move |header: &mut TestHeader| {
+	run_one_test(|header: &mut TestHeader, stage| {
 		let v = std::mem::replace(&mut header.digest_mut().logs, vec![]);
 		header.digest_mut().logs = v.into_iter()
-			.filter(|v| v.as_babe_seal().is_none())
+			.filter(|v| stage == Stage::PreSeal || v.as_babe_seal().is_none())
 			.collect()
-	}));
-	run_one_test()
+	})
 }
 
-// TODO: this test assumes that the test runtime will trigger epoch changes
-// which isn't the case since it doesn't include the session module.
 #[test]
 #[should_panic]
-#[ignore]
 fn rejects_missing_consensus_digests() {
-	MUTATOR.with(|s| *s.borrow_mut() = Arc::new(move |header: &mut TestHeader| {
+	run_one_test(|header: &mut TestHeader, stage| {
 		let v = std::mem::replace(&mut header.digest_mut().logs, vec![]);
 		header.digest_mut().logs = v.into_iter()
-			.filter(|v| v.as_babe_epoch().is_none())
+			.filter(|v| stage == Stage::PostSeal || v.as_next_epoch_descriptor().is_none())
 			.collect()
-	}));
-	run_one_test()
+	});
 }
 
 #[test]
@@ -326,28 +486,34 @@ fn can_author_block() {
 		.expect("Generates authority pair");
 
 	let mut i = 0;
-	let mut epoch = Epoch {
+	let epoch = Epoch {
 		start_slot: 0,
 		authorities: vec![(pair.public(), 1)],
 		randomness: [0; 32],
 		epoch_index: 1,
 		duration: 100,
+	};
+
+	let mut config = crate::BabeConfiguration {
+		slot_duration: 1000,
+		epoch_length: 100,
+		c: (3, 10),
+		genesis_authorities: Vec::new(),
+		randomness: [0; 32],
 		secondary_slots: true,
 	};
 
-	let parent_weight = 0;
-
 	// with secondary slots enabled it should never be empty
-	match claim_slot(i, parent_weight, &epoch, (3, 10), &keystore) {
+	match claim_slot(i, &epoch, &config, &keystore) {
 		None => i += 1,
 		Some(s) => debug!(target: "babe", "Authored block {:?}", s.0),
 	}
 
 	// otherwise with only vrf-based primary slots we might need to try a couple
 	// of times.
-	epoch.secondary_slots = false;
+	config.secondary_slots = false;
 	loop {
-		match claim_slot(i, parent_weight, &epoch, (3, 10), &keystore) {
+		match claim_slot(i, &epoch, &config, &keystore) {
 			None => i += 1,
 			Some(s) => {
 				debug!(target: "babe", "Authored block {:?}", s.0);
@@ -358,14 +524,75 @@ fn can_author_block() {
 }
 
 #[test]
-fn authorities_call_works() {
-	let _ = env_logger::try_init();
-	let client = test_client::new();
+fn importing_block_one_sets_genesis_epoch() {
+	let mut net = BabeTestNet::new(1);
 
-	assert_eq!(client.info().chain.best_number, 0);
-	assert_eq!(epoch(&client, &BlockId::Number(0)).unwrap().into_regular().unwrap().authorities, vec![
-		(Keyring::Alice.public().into(), 1),
-		(Keyring::Bob.public().into(), 1),
-		(Keyring::Charlie.public().into(), 1),
-	]);
+	let peer = net.peer(0);
+	let data = peer.data.as_ref().expect("babe link set up during initialization");
+	let client = peer.client().as_full().expect("Only full clients are used in tests").clone();
+
+	let mut environ = DummyFactory {
+		client: client.clone(),
+		config: data.link.config.clone(),
+		epoch_changes: data.link.epoch_changes.clone(),
+		mutator: Arc::new(|_, _| ()),
+	};
+
+	let genesis_header = client.header(&BlockId::Number(0)).unwrap().unwrap();
+
+	let mut proposer = environ.init(&genesis_header).unwrap();
+	let babe_claim = Item::babe_pre_digest(babe_primitives::BabePreDigest::Secondary {
+		authority_index: 0,
+		slot_number: 999,
+	});
+	let pre_digest = sr_primitives::generic::Digest { logs: vec![babe_claim] };
+
+	let genesis_epoch = data.link.config.genesis_epoch(999);
+
+	let mut block = futures::executor::block_on(proposer.propose_with(pre_digest)).unwrap();
+
+	// seal by alice.
+	let seal = {
+		// sign the pre-sealed hash of the block and then
+		// add it to a digest item.
+		let pair = AuthorityPair::from_seed(&[1; 32]);
+		let pre_hash = block.header.hash();
+		let signature = pair.sign(pre_hash.as_ref());
+		Item::babe_seal(signature)
+	};
+
+	let post_hash = {
+		block.header.digest_mut().push(seal.clone());
+		let h = block.header.hash();
+		block.header.digest_mut().pop();
+		h
+	};
+	assert_eq!(*block.header.number(), 1);
+	let (header, body) = block.deconstruct();
+
+	let post_digests = vec![seal];
+	let mut block_import = data.block_import.lock().take().expect("import set up during init");
+	block_import.import_block(
+		BlockImportParams {
+			origin: BlockOrigin::Own,
+			header,
+			justification: None,
+			post_digests,
+			body: Some(body),
+			finalized: false,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::LongestChain,
+		},
+		Default::default(),
+	).unwrap();
+
+	let epoch_changes = data.link.epoch_changes.lock();
+	let epoch_for_second_block = epoch_changes.epoch_for_child_of(
+		descendent_query(&*client),
+		&post_hash,
+		1,
+		1000,
+		|slot| data.link.config.genesis_epoch(slot),
+	).unwrap().unwrap().into_inner();
+	assert_eq!(epoch_for_second_block, genesis_epoch);
 }