use std::collections::HashMap;
use crate::id::AtomId;
#[derive(Debug, Clone, PartialEq)]
pub struct AdjacencyRow {
id: AtomId,
upstream_ids: Vec<AtomId>,
}
impl AdjacencyRow {
pub fn new(id: AtomId, upstream_ids: Vec<AtomId>) -> Self {
Self { id, upstream_ids }
}
pub fn id(&self) -> AtomId {
self.id
}
pub fn upstream_ids(&self) -> &[AtomId] {
&self.upstream_ids
}
pub fn is_headwater(&self) -> bool {
self.upstream_ids.is_empty()
}
}
#[derive(Debug, thiserror::Error)]
pub enum GraphError {
#[error("drainage graph must contain at least one atom")]
EmptyGraph,
#[error("duplicate atom id {id} at row indices {first} and {second}")]
DuplicateAtomId {
id: i64,
first: usize,
second: usize,
},
}
#[derive(Debug, Clone, PartialEq)]
pub struct DrainageGraph {
rows: Vec<AdjacencyRow>,
index: HashMap<AtomId, usize>,
}
impl DrainageGraph {
pub fn new(rows: Vec<AdjacencyRow>) -> Result<Self, GraphError> {
if rows.is_empty() {
return Err(GraphError::EmptyGraph);
}
let mut index = HashMap::with_capacity(rows.len());
for (i, row) in rows.iter().enumerate() {
if let Some(&first) = index.get(&row.id) {
return Err(GraphError::DuplicateAtomId {
id: row.id.get(),
first,
second: i,
});
}
index.insert(row.id, i);
}
Ok(Self { rows, index })
}
pub fn rows(&self) -> &[AdjacencyRow] {
&self.rows
}
pub fn len(&self) -> usize {
self.rows.len()
}
pub fn is_empty(&self) -> bool {
self.rows.is_empty()
}
pub fn get(&self, id: AtomId) -> Option<&AdjacencyRow> {
self.index.get(&id).map(|&i| &self.rows[i])
}
}
#[cfg(test)]
mod tests {
use super::*;
fn test_atom_id(raw: i64) -> AtomId {
AtomId::new(raw).unwrap()
}
fn headwater_row(id: i64) -> AdjacencyRow {
AdjacencyRow::new(test_atom_id(id), vec![])
}
fn interior_row(id: i64, upstream: &[i64]) -> AdjacencyRow {
let upstream_ids = upstream.iter().map(|&r| test_atom_id(r)).collect();
AdjacencyRow::new(test_atom_id(id), upstream_ids)
}
#[test]
fn single_headwater_row_has_len_one_and_is_not_empty() {
let graph = DrainageGraph::new(vec![headwater_row(1)]).unwrap();
assert_eq!(graph.len(), 1);
assert!(!graph.is_empty());
}
#[test]
fn empty_rows_fails_with_empty_graph_error() {
let err = DrainageGraph::new(vec![]).unwrap_err();
assert!(matches!(err, GraphError::EmptyGraph));
}
#[test]
fn duplicate_atom_id_fails_with_duplicate_error() {
let rows = vec![headwater_row(5), headwater_row(5)];
let err = DrainageGraph::new(rows).unwrap_err();
assert!(matches!(
err,
GraphError::DuplicateAtomId { id: 5, first: 0, second: 1 }
));
}
#[test]
fn get_returns_correct_row() {
let row = interior_row(10, &[11, 12]);
let graph = DrainageGraph::new(vec![
headwater_row(11),
headwater_row(12),
row.clone(),
])
.unwrap();
let found = graph.get(test_atom_id(10)).unwrap();
assert_eq!(found.id(), test_atom_id(10));
assert_eq!(found.upstream_ids().len(), 2);
}
#[test]
fn get_with_nonexistent_id_returns_none() {
let graph = DrainageGraph::new(vec![headwater_row(1)]).unwrap();
assert!(graph.get(test_atom_id(999)).is_none());
}
#[test]
fn is_headwater_true_for_empty_upstream_ids() {
let row = headwater_row(1);
assert!(row.is_headwater());
}
#[test]
fn is_headwater_false_for_nonempty_upstream_ids() {
let row = interior_row(10, &[1, 2]);
assert!(!row.is_headwater());
}
}