1use std::collections::{BTreeMap, BTreeSet};
12
13use nalgebra::DMatrix;
14use snafu::{ResultExt, Snafu};
15use uuid::Uuid;
16
17use crate::{
18 Aggregate, Aggregator, Meta, Metadata, MetadataFilter, Node, NodeNotInStoreSnafu, NodeStore,
19};
20
21#[derive(Clone, Debug, Snafu)]
22#[snafu(visibility(pub(crate)))]
23pub enum EdgeStoreError {
24 NodeStore { source: crate::node::NodeStoreError },
26 MissingNodeStore,
28}
29
30#[derive(Clone, Debug, PartialEq)]
32#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
33pub struct Edge {
34 pub metadata: Metadata,
36 pub source: Uuid,
38 pub target: Uuid,
40}
41impl Edge {
42 pub fn new(metadata: Option<Metadata>, source: Uuid, target: Uuid) -> Self {
44 Edge {
45 metadata: metadata.unwrap_or_default(),
46 source: source.to_owned(),
47 target: target.to_owned(),
48 }
49 }
50 pub fn new_and_nodes(metadata: Option<Metadata>) -> (Self, Node, Node) {
52 let source = Node::default();
53 let target = Node::default();
54 (
55 Edge::new(metadata, source.id().to_owned(), target.id().to_owned()),
56 source,
57 target,
58 )
59 }
60}
61
62impl Meta for Edge {
63 fn get_meta(&self) -> &Metadata {
65 &self.metadata
66 }
67}
68
69#[derive(Clone, Debug, Default, PartialEq)]
70#[cfg_attr(
71 feature = "serde",
72 derive(serde::Serialize, serde::Deserialize),
73 serde(default)
74)]
75pub struct EdgeStoreData {
76 edges: BTreeMap<Uuid, Edge>,
77 forward: BTreeMap<Uuid, BTreeMap<Uuid, BTreeSet<Uuid>>>,
78 reverse: BTreeMap<Uuid, BTreeMap<Uuid, BTreeSet<Uuid>>>,
79 aggregate: Aggregate,
80}
81impl HasEdgeStore for EdgeStoreData {
82 fn edge_store(&self) -> &EdgeStoreData {
83 self
84 }
85 fn edge_store_mut(&mut self) -> &mut EdgeStoreData {
86 self
87 }
88}
89impl EdgeStore for EdgeStoreData {}
90
91pub trait HasEdgeStore {
92 fn edge_store(&self) -> &EdgeStoreData;
94 fn edge_store_mut(&mut self) -> &mut EdgeStoreData;
96}
97
98pub trait EdgeStore: HasEdgeStore {
99 fn edges_len(&self) -> usize {
101 self.edge_store().edges.len()
102 }
103
104 fn is_edges_empty(&self) -> bool {
106 self.edges_len() == 0
107 }
108
109 fn add_edge<N: NodeStore>(
111 &mut self,
112 edge: Edge,
113 safe: bool,
114 nodes: Option<&N>,
115 ) -> Result<Option<Edge>, EdgeStoreError> {
116 if safe {
117 match nodes {
118 None => MissingNodeStoreSnafu.fail(),
119 Some(nodes) => self.check_edge(&edge, nodes),
120 }?;
121 }
122
123 let id = edge.id().to_owned();
124 let replaced = self.del_edge(&id);
125 let source = edge.source.to_owned();
126 let target = edge.target.to_owned();
127
128 let store = self.edge_store_mut();
129
130 store.aggregate.add(edge.get_meta());
131 store.edges.insert(id.to_owned(), edge);
132
133 store
134 .forward
135 .entry(source.to_owned())
136 .or_default()
137 .entry(target.to_owned())
138 .or_default()
139 .insert(id.to_owned());
140
141 store
142 .reverse
143 .entry(target)
144 .or_default()
145 .entry(source)
146 .or_default()
147 .insert(id);
148
149 Ok(replaced)
150 }
151
152 fn extend_edges<N: NodeStore>(
154 &mut self,
155 edges: Vec<Edge>,
156 safe: bool,
157 nodes: Option<&N>,
158 ) -> Result<Vec<Edge>, EdgeStoreError> {
159 let mut replaced = vec![];
160 let store = self.edge_store_mut();
161 for edge in edges.into_iter() {
162 if let Some(edge) = store.add_edge(edge, safe, nodes)? {
163 replaced.push(edge);
164 }
165 }
166 Ok(replaced)
167 }
168
169 fn check_edge<N: NodeStore>(&self, edge: &Edge, nodes: &N) -> Result<(), EdgeStoreError> {
171 if !nodes.has_node(&edge.source) {
172 return NodeNotInStoreSnafu { id: edge.source }
173 .fail()
174 .with_context(|_| NodeStoreSnafu);
175 }
176 if !nodes.has_node(&edge.target) {
177 return NodeNotInStoreSnafu { id: edge.target }
178 .fail()
179 .with_context(|_| NodeStoreSnafu);
180 }
181 Ok(())
182 }
183
184 fn del_edge(&mut self, id: &Uuid) -> Option<Edge> {
186 let store = self.edge_store_mut();
187 if let Some(edge) = store.edges.remove(id) {
188 store.aggregate.subtract(edge.get_meta());
189 let id = edge.id();
190 let source = edge.source;
191 let target = edge.target;
192
193 if let Some(tgt_map) = store.forward.get_mut(&source) {
194 if let Some(edge_ids) = tgt_map.get_mut(&target) {
195 edge_ids.remove(id);
196 if edge_ids.is_empty() {
197 tgt_map.remove(&target);
198 }
199 }
200 if tgt_map.is_empty() {
201 store.forward.remove(&source);
202 }
203 }
204
205 if let Some(src_map) = store.reverse.get_mut(&target) {
206 if let Some(edge_ids) = src_map.get_mut(&source) {
207 edge_ids.remove(id);
208 if edge_ids.is_empty() {
209 src_map.remove(&source);
210 }
211 }
212 if src_map.is_empty() {
213 store.reverse.remove(&target);
214 }
215 }
216
217 return Some(edge);
218 }
219 None
220 }
221
222 fn get_edge(&self, id: &Uuid) -> Option<&Edge> {
224 self.edge_store().edges.get(id)
225 }
226
227 fn all_edges(&self) -> impl Iterator<Item = &Edge> {
229 self.edge_store().edges.values()
230 }
231
232 fn edge_ids_between(&self, source: &Uuid, target: &Uuid) -> impl Iterator<Item = Uuid> {
234 self.edge_store()
235 .forward
236 .get(source)
237 .and_then(|tgt_map| tgt_map.get(target))
238 .map(|ids| ids.iter().copied())
239 .unwrap_or_default()
240 }
241
242 fn edges_between(&self, source: &Uuid, target: &Uuid) -> impl Iterator<Item = &Edge> {
244 self.edge_ids_between(source, target)
245 .filter_map(|id| self.get_edge(&id))
246 }
247
248 fn edge_ids_between_all(
250 &self,
251 sources: &BTreeSet<Uuid>,
252 targets: &BTreeSet<Uuid>,
253 ) -> impl Iterator<Item = Uuid> {
254 sources
255 .iter()
256 .flat_map(|s| targets.iter().map(|t| self.edge_ids_between(s, t)))
257 .flatten()
258 }
259
260 fn edges_between_all(
262 &self,
263 sources: &BTreeSet<Uuid>,
264 targets: &BTreeSet<Uuid>,
265 ) -> impl Iterator<Item = &Edge> {
266 self.edge_ids_between_all(sources, targets)
267 .filter_map(|id| self.get_edge(&id))
268 }
269
270 fn aggregate_between(
273 &self,
274 source: &Uuid,
275 target: &Uuid,
276 filter: Option<&MetadataFilter>,
277 ) -> Aggregate {
278 let mut agg = Aggregate::default();
279 for edge in self.edges_between(source, target) {
280 if let Some(filter) = filter {
281 if !filter.satisfies(edge) {
282 continue;
283 }
284 }
285 agg.add(edge.get_meta());
286 }
287 agg
288 }
289
290 fn aggregate_between_all(
293 &self,
294 sources: &BTreeSet<Uuid>,
295 targets: &BTreeSet<Uuid>,
296 filter: Option<&MetadataFilter>,
297 ) -> Aggregate {
298 let mut agg = Aggregate::default();
299 for source in sources.iter() {
300 for target in targets.iter() {
301 agg.extend(self.aggregate_between(source, target, filter));
302 }
303 }
304 agg
305 }
306
307 fn aggregate_value_between(
310 &self,
311 source: &Uuid,
312 target: &Uuid,
313 aggregator: &Aggregator,
314 filter: Option<&MetadataFilter>,
315 fields: Option<&BTreeSet<String>>,
316 absolute: bool,
317 ) -> f64 {
318 let agg = self.aggregate_between(source, target, filter);
319 agg.sum(aggregator, fields, absolute)
320 }
321
322 fn aggregate_map(
325 &self,
326 sources: &BTreeSet<Uuid>,
327 targets: &BTreeSet<Uuid>,
328 filter: Option<&MetadataFilter>,
329 ) -> BTreeMap<Uuid, BTreeMap<Uuid, Aggregate>> {
330 let mut map = BTreeMap::new();
331 for source in sources.iter() {
332 for target in targets.iter() {
333 let agg = self.aggregate_between(source, target, filter);
334 map.entry(source.to_owned())
335 .or_insert(BTreeMap::<Uuid, Aggregate>::new())
336 .insert(target.to_owned(), agg);
337 }
338 }
339 map
340 }
341
342 fn aggregate_matrix(
346 &self,
347 sources: &[&Uuid],
348 targets: &[&Uuid],
349 filter: Option<&MetadataFilter>,
350 ) -> Vec<Vec<Aggregate>> {
351 targets
352 .iter()
353 .map(|&target| {
354 sources
355 .iter()
356 .map(|&source| self.aggregate_between(source, target, filter))
357 .collect()
358 })
359 .collect()
360 }
361
362 fn outgoing_ids_from(&self, source: &Uuid) -> impl Iterator<Item = Uuid> {
364 self.edge_store()
365 .forward
366 .get(source)
367 .into_iter()
368 .flat_map(|tgt_map| tgt_map.values())
369 .flatten()
370 .copied()
371 }
372
373 fn outgoing_edges_from(&self, source: &Uuid) -> impl Iterator<Item = &Edge> {
375 self.outgoing_ids_from(source)
376 .filter_map(|id| self.get_edge(&id))
377 }
378
379 fn incoming_ids_to(&self, target: &Uuid) -> impl Iterator<Item = Uuid> {
381 self.edge_store()
382 .reverse
383 .get(target)
384 .into_iter()
385 .flat_map(|src_map| src_map.values())
386 .flatten()
387 .copied()
388 }
389
390 fn incoming_edges_to(&self, target: &Uuid) -> impl Iterator<Item = &Edge> {
392 self.incoming_ids_to(target)
393 .filter_map(|id| self.get_edge(&id))
394 }
395
396 fn targets_of(&self, source: &Uuid) -> impl Iterator<Item = Uuid> {
399 self.edge_store()
400 .forward
401 .get(source)
402 .into_iter()
403 .flat_map(|tgt_map| {
404 tgt_map.iter().filter_map(|(target, edges)| {
405 if edges.is_empty() {
406 None
407 } else {
408 Some(*target)
409 }
410 })
411 })
412 }
413
414 fn sources_to(&self, target: &Uuid) -> impl Iterator<Item = Uuid> {
416 self.edge_store()
417 .reverse
418 .get(target)
419 .into_iter()
420 .flat_map(|src_map| {
421 src_map.iter().filter_map(|(source, edges)| {
422 if edges.is_empty() {
423 None
424 } else {
425 Some(*source)
426 }
427 })
428 })
429 }
430
431 fn is_connected_to(
434 &self,
435 node: &Uuid,
436 others: &BTreeSet<Uuid>,
437 edge_ids: Option<&BTreeSet<Uuid>>,
438 ) -> bool {
439 match edge_ids {
440 None => {
441 self.targets_of(node).any(|id| others.contains(&id))
442 || self.sources_to(node).any(|id| others.contains(&id))
443 }
444 Some(edge_ids) => self
445 .edge_ids_between_all(&BTreeSet::from([node.to_owned()]), others)
446 .any(|edge_id| edge_ids.contains(&edge_id)),
447 }
448 }
449
450 fn adjacency_matrix(
453 &self,
454 sources: &Vec<&Uuid>,
455 targets: &Vec<&Uuid>,
456 aggregator: &Aggregator,
457 filter: Option<&MetadataFilter>,
458 fields: Option<&BTreeSet<String>>,
459 absolute: bool,
460 ) -> DMatrix<f64> {
461 let mut matrix = DMatrix::<f64>::zeros(targets.len(), sources.len());
462 for (col, &source) in sources.iter().enumerate() {
463 for (row, &target) in targets.iter().enumerate() {
464 matrix[(row, col)] = self
465 .aggregate_value_between(source, target, aggregator, filter, fields, absolute)
466 }
467 }
468 matrix
469 }
470
471 fn edge_aggregate(&self) -> &Aggregate {
473 &self.edge_store().aggregate
474 }
475}