prospective-parachains rework: take II

polkadot/node/core/prospective-parachains/src/fragment_chain/mod.rs

@@ -216,7 +216,7 @@ impl CandidateStorage {
 	}
 
 	/// Introduce a new candidate entry.
-	pub fn add_candidate_entry(&mut self, candidate: CandidateEntry) -> Result<(), Error> {
+	fn add_candidate_entry(&mut self, candidate: CandidateEntry) -> Result<(), Error> {
 		let candidate_hash = candidate.candidate_hash;
 		if self.by_candidate_hash.contains_key(&candidate_hash) {
 			return Err(Error::CandidateAlreadyKnown)
@@ -274,11 +274,6 @@ impl CandidateStorage {
 		self.by_candidate_hash.values()
 	}
 
-	/// Consume self into an iterator over the stored candidates, in arbitrary order.
-	pub fn into_candidates(self) -> impl Iterator<Item = CandidateEntry> {
-		self.by_candidate_hash.into_values()
-	}
-
 	/// Try getting head-data by hash.
 	fn head_data_by_hash(&self, hash: &Hash) -> Option<&HeadData> {
 		// First, search for candidates outputting this head data and extract the head data
@@ -363,6 +358,10 @@ impl CandidateEntry {
 		Self::new(candidate_hash, candidate, persisted_validation_data, CandidateState::Seconded)
 	}
 
+	pub fn hash(&self) -> CandidateHash {
+		self.candidate_hash
+	}
+
 	fn new(
 		candidate_hash: CandidateHash,
 		candidate: CommittedCandidateReceipt,
@@ -394,10 +393,6 @@ impl CandidateEntry {
 			}),
 		})
 	}
-
-	pub fn hash(&self) -> CandidateHash {
-		self.candidate_hash
-	}
 }
 
 impl HypotheticalOrConcreteCandidate for CandidateEntry {
@@ -542,22 +537,24 @@ impl Scope {
 		self.ancestors_by_hash.get(hash).map(|info| info.clone())
 	}
 
+	/// Get the base constraints of the scope
+	pub fn base_constraints(&self) -> &Constraints {
+		&self.base_constraints
+	}
+
 	/// Whether the candidate in question is one pending availability in this scope.
-	pub fn get_pending_availability(
+	fn get_pending_availability(
 		&self,
 		candidate_hash: &CandidateHash,
 	) -> Option<&PendingAvailability> {
 		self.pending_availability.iter().find(|c| &c.candidate_hash == candidate_hash)
 	}
-
-	/// Get the base constraints of the scope
-	pub fn base_constraints(&self) -> &Constraints {
-		&self.base_constraints
-	}
 }
 
 #[cfg_attr(test, derive(Clone))]
-pub struct FragmentNode {
+/// A node that is part of a `BackedChain`. It holds constraints based on the ancestors in the
+/// chain.
+struct FragmentNode {
 	fragment: Fragment,
 	candidate_hash: CandidateHash,
 	cumulative_modifications: ConstraintModifications,
@@ -571,6 +568,22 @@ impl FragmentNode {
 	}
 }
 
+impl From<&FragmentNode> for CandidateEntry {
+	fn from(node: &FragmentNode) -> Self {
+		// We don't need to perform the checks done in `CandidateEntry::new()`, since a
+		// `FragmentNode` always comes from a `CandidateEntry`
+		Self {
+			candidate_hash: node.candidate_hash,
+			parent_head_data_hash: node.parent_head_data_hash,
+			output_head_data_hash: node.output_head_data_hash,
+			candidate: node.fragment.candidate_clone(),
+			relay_parent: node.relay_parent(),
+			// A fragment node is always backed.
+			state: CandidateState::Backed,
+		}
+	}
+}
+
 /// A candidate chain of backed/backable candidates.
 /// Includes the candidates pending availability and candidates which may be backed on-chain.
 #[derive(Default)]
@@ -659,28 +672,56 @@ pub(crate) struct FragmentChain {
 }
 
 impl FragmentChain {
-	/// Create a new [`FragmentChain`] with given scope and populated from the given storage.
-	/// The `prev_storage` should contain the candidates of the `FragmentChain` at the previous
-	/// relay parent, as well as the candidates pending availability at this relay parent.
-	pub fn populate(scope: Scope, mut prev_storage: CandidateStorage) -> Self {
-		// Initialize as empty
+	/// Create a new [`FragmentChain`] with the given scope and populate it with the candidates
+	/// pending availability.
+	pub fn init(scope: Scope, mut candidates_pending_availability: CandidateStorage) -> Self {
 		let mut fragment_chain = Self {
 			scope,
 			best_chain: BackedChain::default(),
 			unconnected: CandidateStorage::default(),
 		};
 
+		// We only need to populate the best backable chain. Candidates pending availability must
+		// form a chain with the latest included head.
+		fragment_chain.populate_chain(&mut candidates_pending_availability);
+
+		// TODO: return error if not all candidates were introduced successfully.
+
+		fragment_chain
+	}
+
+	/// Populate the [`FragmentChain`] given the new candidates pending availability and the
+	/// optional previous fragment chain (of the previous relay parent).
+	pub fn populate_from_previous(&mut self, prev_fragment_chain: &FragmentChain) {
+		let mut prev_storage = prev_fragment_chain.unconnected.clone();
+
+		for candidate in prev_fragment_chain.best_chain.chain.iter() {
+			// If they used to be pending availability, don't add them. This is fine
+			// because:
+			// - if they still are pending availability, they have already been added to the new
+			//   storage.
+			// - if they were included, no point in keeping them.
+			//
+			// This cannot happen for the candidates in the unconnected storage. The pending
+			// availability candidates will always be part of the best chain.
+			if prev_fragment_chain
+				.scope
+				.get_pending_availability(&candidate.candidate_hash)
+				.is_none()
+			{
+				let _ = prev_storage.add_candidate_entry(candidate.into());
+			}
+		}
+
 		// First populate the best backable chain.
-		fragment_chain.populate_chain(&mut prev_storage);
+		self.populate_chain(&mut prev_storage);
 
 		// Now that we picked the best backable chain, trim the forks generated by candidates which
 		// are not present in the best chain.
-		fragment_chain.trim_uneligible_forks(&mut prev_storage, None);
+		self.trim_uneligible_forks(&mut prev_storage, None);
 
 		// Finally, keep any candidates which haven't been trimmed but still have potential.
-		fragment_chain.populate_unconnected_potential_candidates(prev_storage);
-
-		fragment_chain
+		self.populate_unconnected_potential_candidates(prev_storage);
 	}
 
 	/// Get the scope of the [`FragmentChain`].
@@ -722,34 +763,78 @@ impl FragmentChain {
 			)
 	}
 
-	/// Return a new [`CandidateStorage`] containing all the candidates from this `FragmentChain`,
-	/// as well as the unconnected ones. This does not contain the candidates that used to be
-	/// pending availability.
-	pub fn advance_scope(&self) -> CandidateStorage {
-		let mut storage = self.unconnected.clone();
+	/// Mark a candidate as backed. This can trigger a recreation of the best backable chain.
+	pub fn candidate_backed(&mut self, newly_backed_candidate: &CandidateHash) {
+		// Already backed.
+		if self.best_chain.candidates.contains(newly_backed_candidate) {
+			return
+		}
+		let Some(parent_head_hash) = self
+			.unconnected
+			.by_candidate_hash
+			.get(newly_backed_candidate)
+			.map(|entry| entry.parent_head_data_hash)
+		else {
+			// Candidate is not in unconnected storage.
+			return
+		};
 
-		for candidate in self.best_chain.chain.iter() {
-			// If they used to be pending availability, don't add them. This is fine
-			// because:
-			// - if they still are pending availability, they have already been added to the new
-			//   storage.
-			// - if they were included, no point in keeping them.
-			//
-			// This cannot happen for the candidates in the unconnected storage. The pending
-			// availability candidates will always be part of the best chain.
-			if self.scope.get_pending_availability(&candidate.candidate_hash).is_none() {
-				let _ = storage.add_candidate_entry(CandidateEntry {
-					candidate_hash: candidate.candidate_hash,
-					parent_head_data_hash: candidate.parent_head_data_hash,
-					output_head_data_hash: candidate.output_head_data_hash,
-					relay_parent: candidate.relay_parent(),
-					candidate: candidate.fragment.candidate_clone(), // This clone is very cheap.
-					state: CandidateState::Backed,
-				});
-			}
+		// Mark the candidate hash.
+		self.unconnected.mark_backed(newly_backed_candidate);
+
+		// Revert to parent_head_hash
+		if !self.revert_to(&parent_head_hash) {
+			// If nothing was reverted, there is nothing we can do for now.
+			return
+		}
+
+		let mut prev_storage = std::mem::take(&mut self.unconnected);
+
+		// Populate the chain.
+		self.populate_chain(&mut prev_storage);
+
+		// Now that we picked the best backable chain, trim the forks generated by candidates
+		// which are not present in the best chain. We can start trimming from this candidate
+		// onwards.
+		self.trim_uneligible_forks(&mut prev_storage, Some(parent_head_hash));
+
+		// Finally, keep any candidates which haven't been trimmed but still have potential.
+		self.populate_unconnected_potential_candidates(prev_storage);
+	}
+
+	/// Checks if this candidate could be added in the future to this chain.
+	/// This will return `Error::CandidateAlreadyKnown` if the candidate is already in the chain or
+	/// the unconnected candidate storage.
+	pub fn can_add_candidate_as_potential(
+		&self,
+		candidate: &impl HypotheticalOrConcreteCandidate,
+	) -> Result<(), Error> {
+		let candidate_hash = candidate.candidate_hash();
+
+		if self.best_chain.contains(&candidate_hash) || self.unconnected.contains(&candidate_hash) {
+			return Err(Error::CandidateAlreadyKnown)
+		}
+
+		self.check_potential(candidate)
+	}
+
+	/// Try adding a seconded candidate, if the candidate has potential. It will never be added to
+	/// the chain directly in the seconded state, it will only be part of the unconnected storage.
+	pub fn try_adding_seconded_candidate(
+		&mut self,
+		candidate: &CandidateEntry,
+	) -> Result<(), Error> {
+		if candidate.state == CandidateState::Backed {
+			return Err(Error::IntroduceBackedCandidate);
 		}
 
-		storage
+		self.can_add_candidate_as_potential(candidate)?;
+
+		// This clone is cheap, as it uses an Arc for the expensive stuff.
+		// We can't consume the candidate because other fragment chains may use it also.
+		self.unconnected.add_candidate_entry(candidate.clone())?;
+
+		Ok(())
 	}
 
 	/// Try getting the full head data associated with this hash.
@@ -877,44 +962,9 @@ impl FragmentChain {
 			.unwrap_or_else(|| self.scope.earliest_relay_parent())
 	}
 
-	/// Checks if this candidate could be added in the future to this chain.
-	/// This will return `Error::CandidateAlreadyKnown` if the candidate is already in the chain or
-	/// the unconnected candidate storage.
-	pub fn can_add_candidate_as_potential(
-		&self,
-		candidate: &impl HypotheticalOrConcreteCandidate,
-	) -> Result<(), Error> {
-		let candidate_hash = candidate.candidate_hash();
-
-		if self.best_chain.contains(&candidate_hash) || self.unconnected.contains(&candidate_hash) {
-			return Err(Error::CandidateAlreadyKnown)
-		}
-
-		self.check_potential(candidate)
-	}
-
-	/// Try adding a seconded candidate, if the candidate has potential. It will never be added to
-	/// the chain directly in the seconded state, it will only be part of the unconnected storage.
-	pub fn try_adding_seconded_candidate(
-		&mut self,
-		candidate: &CandidateEntry,
-	) -> Result<(), Error> {
-		if candidate.state == CandidateState::Backed {
-			return Err(Error::IntroduceBackedCandidate);
-		}
-
-		self.can_add_candidate_as_potential(candidate)?;
-
-		// This clone is cheap, as it uses an Arc for the expensive stuff.
-		// We can't consume the candidate because other fragment chains may use it also.
-		self.unconnected.add_candidate_entry(candidate.clone())?;
-
-		Ok(())
-	}
-
 	// Populate the unconnected potential candidate storage starting from a previous storage.
 	fn populate_unconnected_potential_candidates(&mut self, old_storage: CandidateStorage) {
-		for candidate in old_storage.into_candidates() {
+		for candidate in old_storage.by_candidate_hash.into_values() {
 			// Sanity check, all pending availability candidates should be already present in the
 			// chain.
 			if self.scope.get_pending_availability(&candidate.candidate_hash).is_some() {
@@ -1060,6 +1110,8 @@ impl FragmentChain {
 
 	// Once the backable chain was populated, trim the forks generated by candidates which
 	// are not present in the best chain. Fan this out into a full breadth-first search.
+	// If `starting_point` is `Some()`, start the search from the candidates having this parent head
+	// hash.
 	fn trim_uneligible_forks(&self, storage: &mut CandidateStorage, starting_point: Option<Hash>) {
 		// Start out with the candidates in the chain. They are all valid candidates.
 		let mut queue: VecDeque<_> = if let Some(starting_point) = starting_point {
@@ -1275,45 +1327,6 @@ impl FragmentChain {
 		}
 	}
 
-	/// Mark a candidate as backed. This can trigger a recreation of the best backable chain.
-	pub fn candidate_backed(&mut self, newly_backed_candidate: &CandidateHash) {
-		// Already backed.
-		if self.best_chain.candidates.contains(newly_backed_candidate) {
-			return
-		}
-		let Some(parent_head_hash) = self
-			.unconnected
-			.by_candidate_hash
-			.get(newly_backed_candidate)
-			.map(|entry| entry.parent_head_data_hash)
-		else {
-			// Candidate is not in unconnected storage.
-			return
-		};
-
-		// Mark the candidate hash.
-		self.unconnected.mark_backed(newly_backed_candidate);
-
-		// Revert to parent_head_hash
-		if !self.revert_to(&parent_head_hash) {
-			// If nothing was reverted, there is nothing we can do for now.
-			return
-		}
-
-		let mut prev_storage = std::mem::take(&mut self.unconnected);
-
-		// Populate the chain.
-		self.populate_chain(&mut prev_storage);
-
-		// Now that we picked the best backable chain, trim the forks generated by candidates
-		// which are not present in the best chain. We can start trimming from this candidate
-		// onwards.
-		self.trim_uneligible_forks(&mut prev_storage, Some(parent_head_hash));
-
-		// Finally, keep any candidates which haven't been trimmed but still have potential.
-		self.populate_unconnected_potential_candidates(prev_storage);
-	}
-
 	// Revert the best backable chain so that the last candidate will be one outputting the given
 	// `parent_head_hash`. If the `parent_head_hash` is exactly the required parent of the base
 	// constraints (builds on the latest included candidate), revert the entire chain.
@@ -1333,15 +1346,8 @@ impl FragmentChain {
 
 		// Even if it's empty, we need to return true, because we'll be able to add a new candidate
 		// to the chain.
-		for node in removed_items {
-			let _ = self.unconnected.add_candidate_entry(CandidateEntry {
-				candidate_hash: node.candidate_hash,
-				parent_head_data_hash: node.parent_head_data_hash,
-				output_head_data_hash: node.output_head_data_hash,
-				candidate: node.fragment.candidate_clone(),
-				relay_parent: node.relay_parent(),
-				state: CandidateState::Backed,
-			});
+		for node in &removed_items {
+			let _ = self.unconnected.add_candidate_entry(node.into());
 		}
 		true
 	}