eqlog-runtime 0.9.0

use std::collections::{BTreeMap, VecDeque};

use crate::{PrefixTree1, PrefixTree2};

#[derive(Debug, Copy, Clone)]
pub enum ToposortError {
    CycleDetected,
}

#[derive(Debug)]
pub struct MorphismWithSignature {
    pub morph: u32,
    pub dom: u32,
    pub cod: u32,
}

/// Topologically sorts a directed category.
///
/// This is called from code generated by eqlog. The arguments represent the domain and codomain
/// functions Mor -> Obj of the category and the set of objects in the category. The output is a
/// list of morphisms in topological order, i.e. an order such that if `f: A -> B` is before `g: C
/// -> D`, then D != A. Only morphisms for whic both domain and codomain are defined are included
/// in the output.
///
/// To account for the data representation used for seminaive evaluation at the caller site, each
/// of the three input sets is passed in as a pair of disjoint sets of old and new tuples, but this
/// function does not distinguish between old and new tuples. The caller must pass in domain sets
/// as ordered set of (dom(f), f) pairs, whereas the codomain set must be passed in as an ordered
/// set of (f, cod(f)) pairs.
pub fn morphism_toposort(
    dom_new_order_1_0: &PrefixTree2,
    dom_old_order_1_0: &PrefixTree2,
    cod_new_order_0_1: &PrefixTree2,
    cod_old_order_0_1: &PrefixTree2,
    obj_old_order_0: &PrefixTree1,
    obj_new_order_0: &PrefixTree1,
) -> Result<Vec<MorphismWithSignature>, ToposortError> {
    let get_cod = |mor: u32| -> Option<u32> {
        let cods: &PrefixTree1 = cod_new_order_0_1
            .get(mor)
            .or_else(|| cod_old_order_0_1.get(mor))?;
        let [cod] = cods.iter().next().unwrap();
        Some(cod)
    };

    let mut in_degree: BTreeMap<u32, u32> = obj_old_order_0
        .iter()
        .chain(obj_new_order_0.iter())
        .map(|[obj]| (obj, 0))
        .collect();

    for [_, mor] in dom_new_order_1_0.iter().chain(dom_old_order_1_0.iter()) {
        if let Some(cod) = get_cod(mor) {
            *in_degree.get_mut(&cod).unwrap() += 1;
        }
    }

    let mut queue: VecDeque<u32> = in_degree
        .iter()
        .filter_map(|(&obj, &in_deg)| if in_deg == 0 { Some(obj) } else { None })
        .collect();
    let mut ordered_mors: Vec<MorphismWithSignature> = Vec::new();

    // Remove initial zero in-degree nodes from the map
    for &obj in &queue {
        in_degree.remove(&obj);
    }

    while let Some(obj) = queue.pop_front() {
        for dom_order_1_0 in [&dom_new_order_1_0, &dom_old_order_1_0] {
            let out_morphisms = match dom_order_1_0.get(obj) {
                Some(out_morphisms) => out_morphisms,
                None => continue,
            };

            for [out_morphism] in out_morphisms.iter() {
                let cod = match get_cod(out_morphism) {
                    Some(cod) => cod,
                    None => continue,
                };

                ordered_mors.push(MorphismWithSignature {
                    morph: out_morphism,
                    dom: obj,
                    cod,
                });
                let cod_in_degree = in_degree.get_mut(&cod).unwrap();
                *cod_in_degree -= 1;
                if *cod_in_degree == 0 {
                    in_degree.remove(&cod);
                    queue.push_back(cod);
                }
            }
        }
    }

    if !in_degree.is_empty() {
        return Err(ToposortError::CycleDetected);
    }

    Ok(ordered_mors)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_prefix_tree_1(values: &[u32]) -> PrefixTree1 {
        let mut tree = PrefixTree1::new();
        for &value in values {
            tree.insert([value]);
        }
        tree
    }

    fn make_prefix_tree_2(pairs: &[(u32, u32)]) -> PrefixTree2 {
        let mut tree = PrefixTree2::new();
        for &(a, b) in pairs {
            tree.insert([a, b]);
        }
        tree
    }

    #[test]
    fn test_single_morphism() {
        // Simplest case: single morphism 0 -> 1
        let mut dom_new = PrefixTree2::new();
        dom_new.insert([0, 100]); // object 0, morphism 100

        let mut cod_new = PrefixTree2::new();
        cod_new.insert([100, 1]); // morphism 100, object 1

        let mut obj_new = PrefixTree1::new();
        obj_new.insert([0]);
        obj_new.insert([1]);

        let dom_old = PrefixTree2::new();
        let cod_old = PrefixTree2::new();
        let obj_old = PrefixTree1::new();

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(result.is_ok(), "Expected Ok, got: {:?}", result);
        let ordered = result.unwrap();
        assert_eq!(ordered.len(), 1);
        assert_eq!(ordered[0].morph, 100);
        assert_eq!(ordered[0].dom, 0);
        assert_eq!(ordered[0].cod, 1);
    }

    #[test]
    fn test_simple_dag() {
        // Graph: 0 -> 1 -> 2
        // Morphisms: 0: 0 -> 1, 1: 1 -> 2
        let dom_new = make_prefix_tree_2(&[(0, 0), (1, 1)]); // (dom, morphism)
        let dom_old = make_prefix_tree_2(&[]);
        let cod_new = make_prefix_tree_2(&[(0, 1), (1, 2)]); // (morphism, cod)
        let cod_old = make_prefix_tree_2(&[]);
        let obj_new = make_prefix_tree_1(&[0, 1, 2]);
        let obj_old = make_prefix_tree_1(&[]);

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(result.is_ok());
        let ordered = result.unwrap();
        assert_eq!(ordered.len(), 2);
        assert_eq!(ordered[0].morph, 0); // morphism 0 comes before morphism 1
        assert_eq!(ordered[1].morph, 1);
    }

    #[test]
    fn test_cycle_detection() {
        // Graph: 0 -> 1 -> 0 (cycle)
        // Morphisms: 0: 0 -> 1, 1: 1 -> 0
        let dom_new = make_prefix_tree_2(&[(0, 0), (1, 1)]); // (dom, morphism)
        let dom_old = make_prefix_tree_2(&[]);
        let cod_new = make_prefix_tree_2(&[(0, 1), (1, 0)]); // (morphism, cod)
        let cod_old = make_prefix_tree_2(&[]);
        let obj_new = make_prefix_tree_1(&[0, 1]);
        let obj_old = make_prefix_tree_1(&[]);

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(matches!(result, Err(ToposortError::CycleDetected)));
    }

    #[test]
    fn test_disconnected_graph() {
        // Graph: 0 -> 1, 2 -> 3 (two disconnected components)
        // Morphisms: 0: 0 -> 1, 1: 2 -> 3
        let dom_new = make_prefix_tree_2(&[(0, 0), (2, 1)]); // (dom, morphism)
        let dom_old = make_prefix_tree_2(&[]);
        let cod_new = make_prefix_tree_2(&[(0, 1), (1, 3)]); // (morphism, cod)
        let cod_old = make_prefix_tree_2(&[]);
        let obj_new = make_prefix_tree_1(&[0, 1, 2, 3]);
        let obj_old = make_prefix_tree_1(&[]);

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(result.is_ok());
        let ordered = result.unwrap();
        assert_eq!(ordered.len(), 2);
        // Both morphisms should be included
        assert!(ordered.iter().any(|m| m.morph == 0));
        assert!(ordered.iter().any(|m| m.morph == 1));
    }

    #[test]
    fn test_empty_graph() {
        let dom_new = make_prefix_tree_2(&[]);
        let dom_old = make_prefix_tree_2(&[]);
        let cod_new = make_prefix_tree_2(&[]);
        let cod_old = make_prefix_tree_2(&[]);
        let obj_new = make_prefix_tree_1(&[]);
        let obj_old = make_prefix_tree_1(&[]);

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(result.is_ok());
        let ordered = result.unwrap();
        assert_eq!(ordered.len(), 0);
    }

    #[test]
    fn test_self_loop() {
        // Graph: 0 -> 0 (self-loop)
        // Morphism: 0: 0 -> 0
        let dom_new = make_prefix_tree_2(&[(0, 0)]); // (dom, morphism)
        let dom_old = make_prefix_tree_2(&[]);
        let cod_new = make_prefix_tree_2(&[(0, 0)]); // (morphism, cod)
        let cod_old = make_prefix_tree_2(&[]);
        let obj_new = make_prefix_tree_1(&[0]);
        let obj_old = make_prefix_tree_1(&[]);

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(matches!(result, Err(ToposortError::CycleDetected)));
    }

    #[test]
    fn test_complex_dag() {
        // Graph:
        //   0 -> 1 -> 3
        //   |    |
        //   v    v
        //   2 -> 4
        // Morphisms: 0: 0->1, 1: 0->2, 2: 1->3, 3: 1->4, 4: 2->4
        let dom_new = make_prefix_tree_2(&[(0, 0), (0, 1), (1, 2), (1, 3), (2, 4)]); // (dom, morphism)
        let dom_old = make_prefix_tree_2(&[]);
        let cod_new = make_prefix_tree_2(&[(0, 1), (1, 2), (2, 3), (3, 4), (4, 4)]); // (morphism, cod)
        let cod_old = make_prefix_tree_2(&[]);
        let obj_new = make_prefix_tree_1(&[0, 1, 2, 3, 4]);
        let obj_old = make_prefix_tree_1(&[]);

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(result.is_ok());
        let ordered = result.unwrap();
        assert_eq!(ordered.len(), 5);

        // Verify topological order constraints
        let positions: BTreeMap<u32, usize> = ordered
            .iter()
            .enumerate()
            .map(|(i, m)| (m.morph, i))
            .collect();

        // morphism 0 (0->1) should come before morphisms 2 (1->3) and 3 (1->4)
        assert!(positions[&0] < positions[&2]);
        assert!(positions[&0] < positions[&3]);

        // morphism 1 (0->2) should come before morphism 4 (2->4)
        assert!(positions[&1] < positions[&4]);
    }

    #[test]
    fn test_mixed_old_new_data() {
        // Test with data split between old and new prefix trees
        // Graph: 0 -> 1 -> 2
        // Morphisms: 0: 0->1, 1: 1->2
        let dom_new = make_prefix_tree_2(&[(0, 0)]); // (dom, morphism)
        let dom_old = make_prefix_tree_2(&[(1, 1)]); // (dom, morphism)
        let cod_new = make_prefix_tree_2(&[(0, 1)]); // (morphism, cod)
        let cod_old = make_prefix_tree_2(&[(1, 2)]); // (morphism, cod)
        let obj_new = make_prefix_tree_1(&[0, 2]);
        let obj_old = make_prefix_tree_1(&[1]);

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(result.is_ok());
        let ordered = result.unwrap();
        assert_eq!(ordered.len(), 2);
        assert_eq!(ordered[0].morph, 0);
        assert_eq!(ordered[1].morph, 1);
    }

    #[test]
    fn test_missing_codomain() {
        // Test morphism with missing codomain
        // Morphisms: 0: 0->1, 1: 1->? (missing codomain)
        let dom_new = make_prefix_tree_2(&[(0, 0), (1, 1)]); // (dom, morphism)
        let dom_old = make_prefix_tree_2(&[]);
        let cod_new = make_prefix_tree_2(&[(0, 1)]); // Missing codomain for morphism 1
        let cod_old = make_prefix_tree_2(&[]);
        let obj_new = make_prefix_tree_1(&[0, 1]);
        let obj_old = make_prefix_tree_1(&[]);

        let result = morphism_toposort(&dom_new, &dom_old, &cod_new, &cod_old, &obj_old, &obj_new);
        assert!(result.is_ok());
        let ordered = result.unwrap();
        assert_eq!(ordered.len(), 1);
        assert_eq!(ordered[0].morph, 0);
    }
}