Move lexicographical_topological_sort() functionality to rustworkx-core

[mtreinish]

The lexicographical_topological_sort() function is only exposed via a Python interface currently:

rustworkx/src/dag_algo/mod.rs

Lines 347 to 485 in ab93329

	/// Get the lexicographical topological sorted nodes from the provided DAG
	///
	/// This function returns a list of nodes data in a graph lexicographically
	/// topologically sorted using the provided key function. A topological sort
	/// is a linear ordering of vertices such that for every directed edge from
	/// node :math:`u` to node :math:`v`, :math:`u` comes before :math:`v`
	/// in the ordering. If ``reverse`` is set to ``False``, the edges are treated
	/// as if they pointed in the opposite direction.
	///
	/// This function differs from :func:`~rustworkx.topological_sort` because
	/// when there are ties between nodes in the sort order this function will
	/// use the string returned by the ``key`` argument to determine the output
	/// order used. The ``reverse`` argument does not affect the ordering of keys
	/// from this function, only the edges of the graph.
	///
	/// :param PyDiGraph dag: The DAG to get the topological sorted nodes from
	/// :param callable key: key is a python function or other callable that
	/// gets passed a single argument the node data from the graph and is
	/// expected to return a string which will be used for resolving ties
	/// in the sorting order.
	/// :param bool reverse: If ``False`` (the default), perform a regular
	/// topological ordering. If ``True``, return the lexicographical
	/// topological order that would have been found if all the edges in the
	/// graph were reversed. This does not affect the comparisons from the
	/// ``key``.
	/// :param Iterable[int] initial: If given, the initial node indices to start the topological
	/// ordering from. If not given, the topological ordering will certainly contain every node in
	/// the graph. If given, only the ``initial`` nodes and nodes that are dominated by the
	/// ``initial`` set will be in the ordering. Notably, any node that has a natural in degree of
	/// zero will not be in the output ordering if ``initial`` is given and the zero-in-degree node
	/// is not in it. It is a :exc:`ValueError` to give an `initial` set where the nodes have even
	/// a partial topological order between themselves.
	///
	/// :returns: A list of node's data lexicographically topologically sorted.
	/// :rtype: list
	#[pyfunction]
	#[pyo3(signature = (dag, /, key, *, reverse=false, initial=None))]
	pub fn lexicographical_topological_sort(
	py: Python,
	dag: &digraph::PyDiGraph,
	key: PyObject,
	reverse: bool,
	initial: Option<&Bound<PyAny>>,
	) -> PyResult<PyObject> {
	let key_callable = |a: &PyObject| -> PyResult<PyObject> {
	let res = key.call1(py, (a,))?;
	Ok(res.to_object(py))
	};
	// HashMap of node_index indegree
	let node_count = dag.node_count();
	let (in_dir, out_dir) = traversal_directions(reverse);
	#[derive(Clone, Eq, PartialEq)]
	struct State {
	key: String,
	node: NodeIndex,
	}
	impl Ord for State {
	fn cmp(&self, other: &State) -> Ordering {
	// Notice that the we flip the ordering on costs.
	// In case of a tie we compare positions - this step is necessary
	// to make implementations of `PartialEq` and `Ord` consistent.
	other
	.key
	.cmp(&self.key)
	.then_with(|| other.node.index().cmp(&self.node.index()))
	}
	}
	// `PartialOrd` needs to be implemented as well.
	impl PartialOrd for State {
	fn partial_cmp(&self, other: &State) -> Option<Ordering> {
	Some(self.cmp(other))
	}
	}
	let mut in_degree_map: HashMap<NodeIndex, usize> = HashMap::with_capacity(node_count);
	if let Some(initial) = initial {
	// In this case, we don't iterate through all the nodes in the graph, and most nodes aren't
	// in `in_degree_map`; we'll fill in the relevant edge counts lazily.
	for maybe_index in initial.iter()? {
	let node = NodeIndex::new(maybe_index?.extract::<usize>()?);
	if dag.graph.contains_node(node) {
	// It's not necessarily actually zero, but we treat it as if it is. If the node is
	// reachable from another we visit during the iteration, then there was a defined
	// topological order between the `initial` set, and we'll throw an error.
	in_degree_map.insert(node, 0);
	} else {
	return Err(PyValueError::new_err(format!(
	"node index {} is not in this graph",
	node.index()
	)));
	}
	}
	} else {
	for node in dag.graph.node_indices() {
	in_degree_map.insert(node, dag.graph.edges_directed(node, in_dir).count());
	}
	}
	let mut zero_indegree = BinaryHeap::with_capacity(node_count);
	for (node, degree) in in_degree_map.iter() {
	if *degree == 0 {
	let map_key_raw = key_callable(&dag.graph[*node])?;
	let map_key: String = map_key_raw.extract(py)?;
	zero_indegree.push(State {
	key: map_key,
	node: *node,
	});
	}
	}
	let mut out_list: Vec<&PyObject> = Vec::with_capacity(node_count);
	while let Some(State { node, .. }) = zero_indegree.pop() {
	let neighbors = dag.graph.neighbors_directed(node, out_dir);
	for child in neighbors {
	let child_degree = in_degree_map
	.entry(child)
	.or_insert_with(|| dag.graph.edges_directed(child, in_dir).count());
	if *child_degree == 0 {
	return Err(PyValueError::new_err(
	"at least one initial node is reachable from another",
	));
	} else if *child_degree == 1 {
	let map_key_raw = key_callable(&dag.graph[child])?;
	let map_key: String = map_key_raw.extract(py)?;
	zero_indegree.push(State {
	key: map_key,
	node: child,
	});
	in_degree_map.remove(&child);
	} else {
	*child_degree -= 1;
	}
	}
	out_list.push(&dag.graph[node])
	}
	Ok(PyList::new_bound(py, out_list).into())
	}

We should port it to rustworkx-core, so that rust users can leverage the function.

One tweak that probably makes sense for the rustworkx-core version is that instead of returning a PyList (realistically a Vec<N>) of node weights we should have it return of an iterator of node ids. This would be more flexible and performant for rust space users and for the python side of rustworkx we can just consume the iterator to build the pylist return there (for api compatibility).

[hunterkemeny]

Self-assigning this.

[mtreinish]

There was a bit more urgency behind this as I had a few direct requests to have this implemented in rustworkx-core. So I hope you don't mind that I pushed up a PR to implement this here: #1197.