1use std::collections::HashMap;
8
9use nulid::Nulid;
10use serde::{Deserialize, Serialize};
11
12use crate::{Edge, Node};
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct Subgraph {
21 pub nodes: Vec<Node>,
22 pub edges: Vec<Edge>,
23}
24
25#[derive(Debug)]
30pub struct IndexedSubgraph {
31 pub nodes: HashMap<Nulid, Node>,
33 pub edges: HashMap<Nulid, Edge>,
35 pub outgoing: HashMap<Nulid, Vec<usize>>,
38 pub edge_list: Vec<Edge>,
40}
41
42impl IndexedSubgraph {
43 #[must_use]
45 pub fn from_subgraph(sg: &Subgraph) -> Self {
46 let nodes: HashMap<Nulid, Node> = sg.nodes.iter().map(|n| (n.id, n.clone())).collect();
47
48 let edges: HashMap<Nulid, Edge> = sg.edges.iter().map(|e| (e.id, e.clone())).collect();
49
50 let mut outgoing: HashMap<Nulid, Vec<usize>> = HashMap::new();
51 let edge_list: Vec<Edge> = sg.edges.clone();
52
53 for (idx, edge) in edge_list.iter().enumerate() {
54 outgoing.entry(edge.source_id).or_default().push(idx);
55 }
56
57 Self {
58 nodes,
59 edges,
60 outgoing,
61 edge_list,
62 }
63 }
64
65 #[inline]
67 pub fn outgoing_edges(&self, node_id: Nulid) -> &[usize] {
68 self.outgoing.get(&node_id).map_or(&[], Vec::as_slice)
69 }
70
71 #[inline]
73 #[must_use]
74 pub fn node_label(&self, node_id: Nulid) -> Option<&str> {
75 self.nodes.get(&node_id).map(|n| n.label.as_str())
76 }
77}
78
79#[cfg(test)]
80mod tests {
81 use super::*;
82
83 fn make_nulid(n: u8) -> Nulid {
84 Nulid::from_bytes([n; 16])
85 }
86
87 fn make_node(id: Nulid, label: &str, name: &str) -> Node {
88 Node {
89 id,
90 label: label.to_string(),
91 name: name.to_string(),
92 }
93 }
94
95 fn make_edge(id: Nulid, src: Nulid, tgt: Nulid, rel: &str) -> Edge {
96 Edge {
97 id,
98 source_id: src,
99 target_id: tgt,
100 rel_type: rel.to_string(),
101 source_urls: vec!["https://example.com".to_string()],
102 }
103 }
104
105 #[test]
106 fn indexed_subgraph_builds_adjacency_list() {
107 let a = make_nulid(1);
108 let b = make_nulid(2);
109 let e1 = make_nulid(10);
110 let e2 = make_nulid(11);
111
112 let sg = Subgraph {
113 nodes: vec![
114 make_node(a, "Person", "Alice"),
115 make_node(b, "Person", "Bob"),
116 ],
117 edges: vec![
118 make_edge(e1, a, b, "PAID_TO"),
119 make_edge(e2, b, a, "APPOINTED_BY"),
120 ],
121 };
122
123 let idx = IndexedSubgraph::from_subgraph(&sg);
124
125 assert_eq!(idx.nodes.len(), 2);
126 assert_eq!(idx.edges.len(), 2);
127 assert_eq!(idx.outgoing_edges(a).len(), 1);
128 assert_eq!(idx.outgoing_edges(b).len(), 1);
129 assert_eq!(idx.edge_list[idx.outgoing_edges(a)[0]].rel_type, "PAID_TO");
130 assert_eq!(
131 idx.edge_list[idx.outgoing_edges(b)[0]].rel_type,
132 "APPOINTED_BY"
133 );
134 }
135
136 #[test]
137 fn node_label_lookup() {
138 let a = make_nulid(1);
139 let sg = Subgraph {
140 nodes: vec![make_node(a, "Person", "Alice")],
141 edges: vec![],
142 };
143 let idx = IndexedSubgraph::from_subgraph(&sg);
144
145 assert_eq!(idx.node_label(a), Some("Person"));
146 assert_eq!(idx.node_label(make_nulid(99)), None);
147 }
148}