internal: stop toposort panicking when a dependency key is repeated

prqlc/prqlc/src/utils/toposort.rs

@@ -13,6 +13,24 @@ struct Node {
     done: bool,
 }
 
+/// Sorts `dependencies` so that each entry follows the entries it depends on.
+///
+/// Returns one element per element of `dependencies` — not per distinct key —
+/// or `None` if a cycle is reached. With `start`, only the reachable portion
+/// is visited, so a cycle elsewhere in `dependencies` isn't detected.
+///
+/// - A key may be declared more than once. Edges resolve to the *last*
+///   declaration of a key, and every declaration appears in the output, so a
+///   repeated key appears in the output more than once.
+/// - A dependency on a key that is never declared is ignored.
+/// - `start` limits the output to the entries reachable from that key,
+///   including the entry for the key itself.
+///
+/// # Panics
+///
+/// Panics if `start` is `Some` and names a key that isn't declared in
+/// `dependencies`. Note `None` is already used for "there's a cycle", so an
+/// unknown start key can't currently be reported through the return type.
 pub fn toposort<'a, Key: Eq + std::hash::Hash + Clone>(
     dependencies: &'a [(Key, Vec<Key>)],
     start: Option<&'_ Key>,
@@ -35,20 +53,25 @@ pub fn toposort<'a, Key: Eq + std::hash::Hash + Clone>(
         visiting: false,
         done: false,
     };
+    // Sized by `dependencies`, not `index`: a repeated key collapses to a single
+    // `index` entry, but `dag` still has a node per element of `dependencies`.
     let mut toposort = Toposort {
-        nodes: vec![empty; index.len()],
-        order: Vec::with_capacity(index.len()),
+        nodes: vec![empty; dependencies.len()],
+        order: Vec::with_capacity(dependencies.len()),
     };
 
-    if let Some(start) = start.map(|s| index.get(s).unwrap()) {
+    if let Some(start) = start.map(|s| {
+        index
+            .get(s)
+            .expect("`start` must name a key declared in `dependencies`")
+    }) {
         // use only the provided visit start
         toposort.visit(&dag, *start).ok()?;
     } else {
-        // start visits from all nodes
-        while toposort.order.len() < dependencies.len() {
-            for start_at in 0..index.len() {
-                toposort.visit(&dag, start_at).ok()?;
-            }
+        // start visits from all nodes; one pass reaches every node, since
+        // `visit` returns immediately for one that's already done
+        for start_at in 0..dependencies.len() {
+            toposort.visit(&dag, start_at).ok()?;
         }
     }
 
@@ -142,6 +165,34 @@ mod tests {
         assert_eq!(order, vec!["b", "a", "c", "d"]);
     }
 
+    /// A repeated key used to panic (or, for other shapes, loop forever),
+    /// because `nodes` was sized by the deduplicated key count while `dag` was
+    /// sized by the number of entries.
+    #[test]
+    fn repeated_key() {
+        let dependencies = vec![("c", vec!["a"]), ("a", vec![]), ("a", vec![])];
+
+        let order = toposort(&dependencies, None).unwrap();
+
+        // Every entry appears once, and the `a` that `c` resolves to — the
+        // last one declared — comes before it.
+        let order = order.into_iter().copied().collect_vec();
+        assert_eq!(order, vec!["a", "c", "a"]);
+    }
+
+    /// Cycle detection only covers the visited part of the graph, so a cycle
+    /// that `start` can't reach goes unreported.
+    #[test]
+    fn cycle_unreachable_from_root() {
+        let dependencies = vec![("a", vec![]), ("b", vec!["c"]), ("c", vec!["b"])];
+        let root = "a";
+
+        let order = toposort(&dependencies, Some(&root)).unwrap();
+
+        let order = order.into_iter().copied().collect_vec();
+        assert_eq!(order, vec!["a"]);
+    }
+
     #[test]
     fn with_root() {
         let dependencies = vec![