1use std::collections::HashMap;
4
5use crate::id::AtomId;
6
7#[derive(Debug, Clone, PartialEq)]
11pub struct AdjacencyRow {
12 id: AtomId,
13 upstream_ids: Vec<AtomId>,
14}
15
16impl AdjacencyRow {
17 pub fn new(id: AtomId, upstream_ids: Vec<AtomId>) -> Self {
19 Self { id, upstream_ids }
20 }
21
22 pub fn id(&self) -> AtomId {
24 self.id
25 }
26
27 pub fn upstream_ids(&self) -> &[AtomId] {
29 &self.upstream_ids
30 }
31
32 pub fn is_headwater(&self) -> bool {
34 self.upstream_ids.is_empty()
35 }
36}
37
38#[derive(Debug, thiserror::Error)]
40pub enum GraphError {
41 #[error("drainage graph must contain at least one atom")]
43 EmptyGraph,
44
45 #[error("duplicate atom id {id} at row indices {first} and {second}")]
47 DuplicateAtomId {
48 id: i64,
50 first: usize,
52 second: usize,
54 },
55}
56
57#[derive(Debug, Clone, PartialEq)]
68pub struct DrainageGraph {
69 rows: Vec<AdjacencyRow>,
70 index: HashMap<AtomId, usize>,
71}
72
73impl DrainageGraph {
74 pub fn new(rows: Vec<AdjacencyRow>) -> Result<Self, GraphError> {
85 if rows.is_empty() {
86 return Err(GraphError::EmptyGraph);
87 }
88
89 let mut index = HashMap::with_capacity(rows.len());
90 for (i, row) in rows.iter().enumerate() {
91 if let Some(&first) = index.get(&row.id) {
92 return Err(GraphError::DuplicateAtomId {
93 id: row.id.get(),
94 first,
95 second: i,
96 });
97 }
98 index.insert(row.id, i);
99 }
100
101 Ok(Self { rows, index })
102 }
103
104 pub fn rows(&self) -> &[AdjacencyRow] {
106 &self.rows
107 }
108
109 pub fn len(&self) -> usize {
111 self.rows.len()
112 }
113
114 pub fn is_empty(&self) -> bool {
116 self.rows.is_empty()
117 }
118
119 pub fn get(&self, id: AtomId) -> Option<&AdjacencyRow> {
123 self.index.get(&id).map(|&i| &self.rows[i])
124 }
125}
126
127#[cfg(test)]
128mod tests {
129 use super::*;
130
131 fn test_atom_id(raw: i64) -> AtomId {
132 AtomId::new(raw).unwrap()
133 }
134
135 fn headwater_row(id: i64) -> AdjacencyRow {
136 AdjacencyRow::new(test_atom_id(id), vec![])
137 }
138
139 fn interior_row(id: i64, upstream: &[i64]) -> AdjacencyRow {
140 let upstream_ids = upstream.iter().map(|&r| test_atom_id(r)).collect();
141 AdjacencyRow::new(test_atom_id(id), upstream_ids)
142 }
143
144 #[test]
145 fn single_headwater_row_has_len_one_and_is_not_empty() {
146 let graph = DrainageGraph::new(vec![headwater_row(1)]).unwrap();
147 assert_eq!(graph.len(), 1);
148 assert!(!graph.is_empty());
149 }
150
151 #[test]
152 fn empty_rows_fails_with_empty_graph_error() {
153 let err = DrainageGraph::new(vec![]).unwrap_err();
154 assert!(matches!(err, GraphError::EmptyGraph));
155 }
156
157 #[test]
158 fn duplicate_atom_id_fails_with_duplicate_error() {
159 let rows = vec![headwater_row(5), headwater_row(5)];
160 let err = DrainageGraph::new(rows).unwrap_err();
161 assert!(matches!(
162 err,
163 GraphError::DuplicateAtomId { id: 5, first: 0, second: 1 }
164 ));
165 }
166
167 #[test]
168 fn get_returns_correct_row() {
169 let row = interior_row(10, &[11, 12]);
170 let graph = DrainageGraph::new(vec![
171 headwater_row(11),
172 headwater_row(12),
173 row.clone(),
174 ])
175 .unwrap();
176
177 let found = graph.get(test_atom_id(10)).unwrap();
178 assert_eq!(found.id(), test_atom_id(10));
179 assert_eq!(found.upstream_ids().len(), 2);
180 }
181
182 #[test]
183 fn get_with_nonexistent_id_returns_none() {
184 let graph = DrainageGraph::new(vec![headwater_row(1)]).unwrap();
185 assert!(graph.get(test_atom_id(999)).is_none());
186 }
187
188 #[test]
189 fn is_headwater_true_for_empty_upstream_ids() {
190 let row = headwater_row(1);
191 assert!(row.is_headwater());
192 }
193
194 #[test]
195 fn is_headwater_false_for_nonempty_upstream_ids() {
196 let row = interior_row(10, &[1, 2]);
197 assert!(!row.is_headwater());
198 }
199}