1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! The RPSL `as-set` and `as-num` directed graph.
use hashbrown::HashMap;
use petgraph::{
    algo::{dijkstra, is_cyclic_directed},
    dot::Dot,
    prelude::{DiGraph, NodeIndex},
};

#[cfg(test)]
mod tests;

/// AS Sets and AS Num membership graph.
/// The display format can be used with Graphviz to visualize the graph.
/// (Try print one out and paste it into Graphviz online.)
#[derive(Clone, Debug, Default)]
pub struct ASSetGraph {
    pub as_num_and_sets: HashMap<ASNumOrSet, NodeIndex>,
    /// Membership graph of AS Nums and AS sets.
    /// Edge weights are 1 except for pseudo sets.
    pub graph: DiGraph<ASNumOrSet, u32>,
}

impl std::fmt::Display for ASSetGraph {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.to_dot())
    }
}

impl ASSetGraph {
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            as_num_and_sets: HashMap::with_capacity(capacity),
            graph: DiGraph::with_capacity(capacity, capacity << 1),
        }
    }

    pub fn get_or_insert(&mut self, as_num_or_set: ASNumOrSet) -> NodeIndex {
        if let Some(index) = self.as_num_and_sets.get(&as_num_or_set) {
            *index
        } else {
            let index = self.graph.add_node(as_num_or_set.clone());
            _ = self.as_num_and_sets.insert(as_num_or_set, index);
            index
        }
    }

    /// Add each member in `members` to `set`.
    pub fn add_members<I>(&mut self, members: I, set: ASNumOrSet) -> (Vec<NodeIndex>, NodeIndex)
    where
        I: IntoIterator<Item = ASNumOrSet>,
    {
        let set_index = self.get_or_insert(set);
        let member_indexes = members
            .into_iter()
            .map(|member| {
                let edge_weight = if member.is_pseudo_set() { 0 } else { 1 };
                let member_index = self.get_or_insert(member);
                self.graph.add_edge(set_index, member_index, edge_weight);
                member_index
            })
            .collect();
        (member_indexes, set_index)
    }

    pub fn count_stats(&self, set_index: NodeIndex) -> ASSetGraphStats {
        let mut stats = ASSetGraphStats::default();

        let shortest_distances = dijkstra(&self.graph, set_index, None, |e| *e.weight());
        for (as_num_or_set, node_index) in &self.as_num_and_sets {
            if *node_index == set_index {
                continue;
            }

            match as_num_or_set {
                ASNumOrSet::Num(_) => {
                    stats.n_nums += 1;
                    if let Some(distance) = shortest_distances.get(node_index) {
                        stats.depth = stats.depth.max(*distance);
                    }
                }
                ASNumOrSet::Set(_) => {
                    if !as_num_or_set.is_pseudo_set() {
                        stats.n_sets += 1
                    }
                }
            }
        }

        stats
    }

    pub fn has_cycle(&self) -> bool {
        is_cyclic_directed(&self.graph)
    }

    pub fn to_dot(&self) -> Dot<&DiGraph<ASNumOrSet, u32>> {
        Dot::new(&self.graph)
    }
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct ASSetGraphStats {
    pub n_sets: u32,
    pub n_nums: u32,
    pub depth: u32,
}

#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum ASNumOrSet {
    Num(u32),
    Set(String),
}

impl ASNumOrSet {
    pub fn set(s: &str) -> Self {
        Self::Set(s.into())
    }

    pub fn is_pseudo_set(&self) -> bool {
        match self {
            ASNumOrSet::Set(set) => set.contains('#'),
            _ => false,
        }
    }
}

impl std::fmt::Display for ASNumOrSet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ASNumOrSet::Num(num) => write!(f, "AS{num}"),
            ASNumOrSet::Set(set) => f.write_str(set),
        }
    }
}