use itertools::Itertools;
use rayon::prelude::*;
use crate::{
HyperedgeIndex,
HyperedgeTrait,
Hypergraph,
VertexIndex,
VertexTrait,
errors::HypergraphError,
};
pub(crate) enum Connection<Index = VertexIndex> {
In(Index),
Out(Index),
InAndOut(Index, Index),
}
type Connections = Vec<(HyperedgeIndex, Option<VertexIndex>)>;
impl<V, HE> Hypergraph<V, HE>
where
V: VertexTrait,
HE: HyperedgeTrait,
{
#[allow(clippy::type_complexity)]
pub(crate) fn get_connections(
&self,
connections: &Connection,
) -> Result<Connections, HypergraphError<V, HE>> {
let internal_index = self.get_internal_vertex(match connections {
Connection::InAndOut(vertex_index, _)
| Connection::In(vertex_index)
| Connection::Out(vertex_index) => *vertex_index,
})?;
let (_, hyperedges_index_set) = self
.vertices
.get(internal_index)
.ok_or(HypergraphError::InternalVertexIndexNotFound(internal_index))?;
let hyperedges =
self.get_hyperedges(&hyperedges_index_set.iter().copied().collect_vec())?;
let hyperedges_with_vertices = hyperedges
.into_par_iter()
.map(|hyperedge_index| {
self.get_hyperedge_vertices(hyperedge_index)
.map(|vertices| (hyperedge_index, vertices))
})
.collect::<Result<Vec<(HyperedgeIndex, Vec<VertexIndex>)>, HypergraphError<V, HE>>>()?;
let capacity = hyperedges_with_vertices.len();
let results = hyperedges_with_vertices
.into_par_iter()
.fold_with(
Vec::with_capacity(capacity),
|acc, (hyperedge_index, vertices)| {
vertices.iter().tuple_windows::<(_, _)>().fold(
acc,
|mut acc, (window_from, window_to)| {
match connections {
Connection::In(from) => {
if *window_from == *from {
acc.push((hyperedge_index, Some(*window_to)));
}
}
Connection::Out(to) => {
if *window_to == *to {
acc.push((hyperedge_index, Some(*window_from)));
}
}
Connection::InAndOut(from, to) => {
if *window_from == *from && *window_to == *to {
acc.push((hyperedge_index, None));
}
}
}
acc
},
)
},
)
.flatten()
.collect::<Connections>();
Ok(results)
}
}