provisioner: allow multiple cores assigned to the same para

polkadot/node/core/prospective-parachains/src/fragment_tree.rs

@@ -756,37 +756,31 @@ impl FragmentTree {
 		depths.iter_ones().collect()
 	}
 
-	/// Select `count` candidates after the given `required_path` which pass
+	/// Select `count` candidates after the given `required_ancestors` which pass
 	/// the predicate and have not already been backed on chain.
 	///
-	/// Does an exhaustive search into the tree starting after `required_path`.
-	/// If there are multiple possibilities of size `count`, this will select the first one.
-	/// If there is no chain of size `count` that matches the criteria, this will return the largest
-	/// chain it could find with the criteria.
+	/// Does an exhaustive search into the tree starting after `required_ancestors`.
+	//  If there are multiple possibilities of size `count`, this
+	/// will select the first one. If there is no chain of size `count` that matches the criteria,
+	/// this will return the largest chain it could find with the criteria.
 	/// If there are no candidates meeting those criteria, returns an empty `Vec`.
 	/// Cycles are accepted, see module docs for the `Cycles` section.
 	///
-	/// The intention of the `required_path` is to allow queries on the basis of
+	/// The intention of the `required_ancestors` is to allow queries on the basis of
 	/// one or more candidates which were previously pending availability becoming
-	/// available and opening up more room on the core.
+	/// available and opening up more room on the cores.
+	///
+	/// Assumes `required_ancestors` is not ordered and takes care of finding a valid path from the
+	/// root that includes all ancestors.
 	pub(crate) fn select_children(
 		&self,
-		required_path: &[CandidateHash],
+		required_ancestors: &[CandidateHash],
 		count: u32,
 		pred: impl Fn(&CandidateHash) -> bool,
 	) -> Vec<CandidateHash> {
-		let base_node = {
-			// traverse the required path.
-			let mut node = NodePointer::Root;
-			for required_step in required_path {
-				if let Some(next_node) = self.node_candidate_child(node, &required_step) {
-					node = next_node;
-				} else {
-					return vec![]
-				};
-			}
-
-			node
+		// First, we need to order the required_ancestors.
+		let Some(base_node) = self.find_path_from_candidates(required_ancestors) else {
+			return vec![]
 		};
 
 		// TODO: taking the first best selection might introduce bias
@@ -879,6 +873,66 @@ impl FragmentTree {
 		best_result
 	}
 
+	// Orders the ancestors into a viable path from root to the last one.
+	// Returns the pointer to the last node in the path.
+	fn find_path_from_candidates(
+		&self,
+		required_ancestors: &[CandidateHash],
+	) -> Option<NodePointer> {
+		let mut n_required_ancestors = required_ancestors.len();
+		// We may have duplicates, which are permitted. Use a HashMap to keep a record of ancestors
+		// and their occurrence count.
+		let mut ancestors_map: HashMap<CandidateHash, usize> = required_ancestors.iter().fold(
+			HashMap::with_capacity(n_required_ancestors),
+			|mut acc, candidate| {
+				acc.entry(*candidate).and_modify(|c| *c += 1).or_insert(1);
+				acc
+			},
+		);
+
+		let mut last_node = NodePointer::Root;
+		let mut next_node: Option<NodePointer> = Some(NodePointer::Root);
+
+		while let Some(node) = next_node {
+			last_node = node;
+
+			next_node = match node {
+				NodePointer::Root => self
+					.nodes
+					.iter()
+					.enumerate()
+					.take_while(|n| n.1.parent == NodePointer::Root)
+					.find_map(|(i, n)| match ancestors_map.get_mut(&n.candidate_hash) {
+						Some(count) if *count > 0 => {
+							*count -= 1;
+							n_required_ancestors -= 1;
+							Some(NodePointer::Storage(i))
+						},
+						_ => None,
+					}),
+				NodePointer::Storage(ptr) => self.nodes.get(ptr).and_then(|n| {
+					n.children.iter().find_map(|(node_ptr, hash)| {
+						match ancestors_map.get_mut(hash) {
+							Some(count) if *count > 0 => {
+								*count -= 1;
+								n_required_ancestors -= 1;
+								Some(*node_ptr)
+							},
+							_ => None,
+						}
+					})
+				}),
+			};
+		}
+
+		if n_required_ancestors != 0 {
+			// Could not traverse the path. We should have used every ancestor.
+			None
+		} else {
+			Some(last_node)
+		}
+	}
+
 	fn populate_from_bases(&mut self, storage: &CandidateStorage, initial_bases: Vec<NodePointer>) {
 		// Populate the tree breadth-first.
 		let mut last_sweep_start = None;
@@ -1614,6 +1668,30 @@ mod tests {
 					.collect::<Vec<_>>()
 			);
 		}
+		assert_eq!(
+			tree.select_children(
+				&iter::repeat(candidate_a_hash).take(2).collect::<Vec<_>>(),
+				2,
+				|_| true
+			),
+			vec![candidate_a_hash, candidate_a_hash]
+		);
+		assert_eq!(
+			tree.select_children(
+				&iter::repeat(candidate_a_hash).take(3).collect::<Vec<_>>(),
+				3,
+				|_| true
+			),
+			vec![candidate_a_hash, candidate_a_hash]
+		);
+		assert_eq!(
+			tree.select_children(
+				&iter::repeat(candidate_a_hash).take(max_depth + 1).collect::<Vec<_>>(),
+				5,
+				|_| true
+			),
+			vec![]
+		);
 	}
 
 	#[test]
@@ -1710,6 +1788,38 @@ mod tests {
 			tree.select_children(&[candidate_a_hash, candidate_b_hash], 6, |_| true),
 			vec![candidate_a_hash, candidate_b_hash, candidate_a_hash,],
 		);
+		assert_eq!(
+			tree.select_children(&[candidate_b_hash, candidate_a_hash], 6, |_| true),
+			vec![candidate_a_hash, candidate_b_hash, candidate_a_hash,],
+		);
+
+		// Unordered ancestors.
+		for count in 1..5 {
+			assert_eq!(
+				tree.select_children(
+					&[candidate_b_hash, candidate_a_hash, candidate_b_hash, candidate_a_hash],
+					count,
+					|_| true
+				),
+				vec![candidate_a_hash]
+			);
+		}
+
+		// Invalid ancestors. a->b->a->b->b.
+		assert_eq!(
+			tree.select_children(
+				&[
+					candidate_b_hash,
+					candidate_a_hash,
+					candidate_b_hash,
+					candidate_a_hash,
+					candidate_b_hash
+				],
+				1,
+				|_| true
+			),
+			vec![]
+		);
 	}
 
 	#[test]

polkadot/node/core/prospective-parachains/src/lib.rs

@@ -150,14 +150,14 @@ async fn run_iteration<Context>(
 					relay_parent,
 					para,
 					count,
-					required_path,
+					required_ancestors,
 					tx,
 				) => answer_get_backable_candidates(
 					&view,
 					relay_parent,
 					para,
 					count,
-					required_path,
+					required_ancestors,
 					tx,
 				),
 				ProspectiveParachainsMessage::GetHypotheticalFrontier(request, tx) =>
@@ -565,7 +565,7 @@ fn answer_get_backable_candidates(
 	relay_parent: Hash,
 	para: ParaId,
 	count: u32,
-	required_path: Vec<CandidateHash>,
+	required_ancestors: Vec<CandidateHash>,
 	tx: oneshot::Sender<Vec<(CandidateHash, Hash)>>,
 ) {
 	let data = match view.active_leaves.get(&relay_parent) {
@@ -614,7 +614,7 @@ fn answer_get_backable_candidates(
 	};
 
 	let backable_candidates: Vec<_> = tree
-		.select_children(&required_path, count, |candidate| storage.is_backed(candidate))
+		.select_children(&required_ancestors, count, |candidate| storage.is_backed(candidate))
 		.into_iter()
 		.filter_map(|child_hash| {
 			storage.relay_parent_by_candidate_hash(&child_hash).map_or_else(
@@ -635,7 +635,7 @@ fn answer_get_backable_candidates(
 	if backable_candidates.is_empty() {
 		gum::trace!(
 			target: LOG_TARGET,
-			?required_path,
+			?required_ancestors,
 			para_id = ?para,
 			%relay_parent,
 			"Could not find any backable candidate",