kmp_adapter_embedded/adapter/
graph_read.rs

1use std::collections::{BTreeMap, BTreeSet, VecDeque};
2
3use kmp_domain::{
4    ContextPathNeighborhood, GraphNeighborhoodReader, MemoryAboutIndexReader, NodeNeighborhood,
5    NodeProjection, NodeRelationProjection, NodeRelationshipReader, NodeRelationships, PortError,
6};
7use redb::{ReadOnlyTable, ReadableTable};
8
9use super::projection_write::MEMORY_ANCHOR_KIND;
10use super::serdes::{NodeRecord, decode, decode_explanation};
11use super::store::{
12    ANCHORS, EmbeddedKernelStore, NODES, RELATIONS, RELATIONS_BY_TARGET, range_error,
13    storage_error, table_error,
14};
15
16type NodeReadTable = ReadOnlyTable<&'static str, &'static [u8]>;
17type RelationReadTable = ReadOnlyTable<(&'static str, &'static str, &'static str), &'static [u8]>;
18
19fn load_node(nodes: &NodeReadTable, node_id: &str) -> Result<Option<NodeProjection>, PortError> {
20    match nodes.get(node_id).map_err(storage_error)? {
21        Some(guard) => Ok(Some(
22            decode::<NodeRecord>("graph node", guard.value())?.into_projection()?,
23        )),
24        None => Ok(None),
25    }
26}
27
28fn outgoing_rows(
29    relations: &RelationReadTable,
30    source: &str,
31) -> Result<Vec<NodeRelationProjection>, PortError> {
32    let upper = format!("{source}\u{0}");
33    let mut rows = Vec::new();
34    for row in relations
35        .range((source, "", "")..(upper.as_str(), "", ""))
36        .map_err(range_error)?
37    {
38        let (key, value) = row.map_err(range_error)?;
39        let (source_node_id, target_node_id, relation_type) = key.value();
40        rows.push(NodeRelationProjection {
41            source_node_id: source_node_id.to_string(),
42            target_node_id: target_node_id.to_string(),
43            relation_type: relation_type.to_string(),
44            explanation: decode_explanation(value.value())?,
45        });
46    }
47    Ok(rows)
48}
49
50fn outgoing_targets(relations: &RelationReadTable, source: &str) -> Result<Vec<String>, PortError> {
51    let upper = format!("{source}\u{0}");
52    let mut targets = Vec::new();
53    for row in relations
54        .range((source, "", "")..(upper.as_str(), "", ""))
55        .map_err(range_error)?
56    {
57        let (key, _) = row.map_err(range_error)?;
58        targets.push(key.value().1.to_string());
59    }
60    Ok(targets)
61}
62
63fn reachable_outward(
64    relations: &RelationReadTable,
65    root_node_id: &str,
66    depth: u32,
67) -> Result<BTreeSet<String>, PortError> {
68    let mut visited = BTreeSet::from([root_node_id.to_string()]);
69    let mut reachable = BTreeSet::new();
70    let mut frontier = VecDeque::from([(root_node_id.to_string(), 0u32)]);
71
72    while let Some((node_id, hops)) = frontier.pop_front() {
73        if hops == depth {
74            continue;
75        }
76        for target in outgoing_targets(relations, &node_id)? {
77            if visited.insert(target.clone()) {
78                reachable.insert(target.clone());
79                frontier.push_back((target, hops + 1));
80            }
81        }
82    }
83
84    reachable.remove(root_node_id);
85    Ok(reachable)
86}
87
88fn relations_among(
89    relations: &RelationReadTable,
90    selected: &BTreeSet<String>,
91) -> Result<Vec<NodeRelationProjection>, PortError> {
92    let mut rows = Vec::new();
93    for source in selected {
94        for relation in outgoing_rows(relations, source)? {
95            if selected.contains(&relation.target_node_id) {
96                rows.push(relation);
97            }
98        }
99    }
100    Ok(rows)
101}
102
103fn selected_projections(
104    nodes: &NodeReadTable,
105    selected: &BTreeSet<String>,
106    root_node_id: &str,
107) -> Result<Vec<NodeProjection>, PortError> {
108    let mut projections = Vec::new();
109    for node_id in selected {
110        if node_id == root_node_id {
111            continue;
112        }
113        if let Some(projection) = load_node(nodes, node_id)? {
114            projections.push(projection);
115        }
116    }
117    Ok(projections)
118}
119
120fn shortest_outward_path(
121    relations: &RelationReadTable,
122    root_node_id: &str,
123    target_node_id: &str,
124) -> Result<Option<Vec<String>>, PortError> {
125    let mut predecessors = BTreeMap::<String, String>::new();
126    let mut visited = BTreeSet::from([root_node_id.to_string()]);
127    let mut frontier = VecDeque::from([root_node_id.to_string()]);
128
129    while let Some(node_id) = frontier.pop_front() {
130        for target in outgoing_targets(relations, &node_id)? {
131            if !visited.insert(target.clone()) {
132                continue;
133            }
134            predecessors.insert(target.clone(), node_id.clone());
135            if target == target_node_id {
136                let mut path = vec![target.clone()];
137                let mut current = target.as_str();
138                while let Some(previous) = predecessors.get(current) {
139                    path.push(previous.clone());
140                    current = previous;
141                }
142                path.reverse();
143                return Ok(Some(path));
144            }
145            frontier.push_back(target);
146        }
147    }
148
149    Ok(None)
150}
151
152impl GraphNeighborhoodReader for EmbeddedKernelStore {
153    async fn load_neighborhood(
154        &self,
155        root_node_id: &str,
156        depth: u32,
157    ) -> Result<Option<NodeNeighborhood>, PortError> {
158        let root_node_id = root_node_id.to_string();
159        self.run(move |store| {
160            let tx = store.begin_read()?;
161            let nodes = tx.open_table(NODES).map_err(table_error)?;
162            let relations = tx.open_table(RELATIONS).map_err(table_error)?;
163
164            let Some(root) = load_node(&nodes, &root_node_id)? else {
165                return Ok(None);
166            };
167
168            let reachable = reachable_outward(&relations, &root_node_id, depth)?;
169            // Mirrors the Neo4j neighborhood query: an empty neighborhood
170            // reports no relations, even for self-referential root edges.
171            let relation_rows = if reachable.is_empty() {
172                Vec::new()
173            } else {
174                let mut selected = reachable.clone();
175                selected.insert(root_node_id.clone());
176                relations_among(&relations, &selected)?
177            };
178
179            Ok(Some(NodeNeighborhood {
180                neighbors: selected_projections(&nodes, &reachable, &root_node_id)?,
181                relations: relation_rows,
182                root,
183            }))
184        })
185        .await
186    }
187
188    async fn load_context_path(
189        &self,
190        root_node_id: &str,
191        target_node_id: &str,
192        subtree_depth: u32,
193    ) -> Result<Option<ContextPathNeighborhood>, PortError> {
194        let root_node_id = root_node_id.to_string();
195        let target_node_id = target_node_id.to_string();
196        self.run(move |store| {
197            let tx = store.begin_read()?;
198            let nodes = tx.open_table(NODES).map_err(table_error)?;
199            let relations = tx.open_table(RELATIONS).map_err(table_error)?;
200
201            let Some(root) = load_node(&nodes, &root_node_id)? else {
202                return Ok(None);
203            };
204            if load_node(&nodes, &target_node_id)?.is_none() {
205                return Ok(None);
206            }
207            let Some(path_node_ids) =
208                shortest_outward_path(&relations, &root_node_id, &target_node_id)?
209            else {
210                return Ok(None);
211            };
212
213            let mut selected = path_node_ids.iter().cloned().collect::<BTreeSet<_>>();
214            selected.insert(target_node_id.clone());
215            selected.extend(reachable_outward(
216                &relations,
217                &target_node_id,
218                subtree_depth,
219            )?);
220
221            Ok(Some(ContextPathNeighborhood {
222                neighbors: selected_projections(&nodes, &selected, &root_node_id)?,
223                relations: relations_among(&relations, &selected)?,
224                path_node_ids,
225                root,
226            }))
227        })
228        .await
229    }
230}
231
232impl NodeRelationshipReader for EmbeddedKernelStore {
233    async fn load_node_relationships(
234        &self,
235        node_id: &str,
236    ) -> Result<Option<NodeRelationships>, PortError> {
237        let node_id = node_id.to_string();
238        self.run(move |store| {
239            let tx = store.begin_read()?;
240            let nodes = tx.open_table(NODES).map_err(table_error)?;
241            if load_node(&nodes, &node_id)?.is_none() {
242                return Ok(None);
243            }
244            let relations = tx.open_table(RELATIONS).map_err(table_error)?;
245            let by_target = tx.open_table(RELATIONS_BY_TARGET).map_err(table_error)?;
246
247            let mut incoming = Vec::new();
248            let upper = format!("{node_id}\u{0}");
249            for row in by_target
250                .range((node_id.as_str(), "", "")..(upper.as_str(), "", ""))
251                .map_err(range_error)?
252            {
253                let (key, _) = row.map_err(range_error)?;
254                let (target, source, relation_type) = key.value();
255                let Some(value) = relations
256                    .get((source, target, relation_type))
257                    .map_err(storage_error)?
258                else {
259                    return Err(PortError::InvalidState(format!(
260                        "embedded store adjacency index points at missing relation \
261                         `{source}` -> `{target}` ({relation_type})"
262                    )));
263                };
264                incoming.push(NodeRelationProjection {
265                    source_node_id: source.to_string(),
266                    target_node_id: target.to_string(),
267                    relation_type: relation_type.to_string(),
268                    explanation: decode_explanation(value.value())?,
269                });
270            }
271
272            Ok(Some(NodeRelationships {
273                incoming,
274                outgoing: outgoing_rows(&relations, &node_id)?,
275            }))
276        })
277        .await
278    }
279}
280
281impl MemoryAboutIndexReader for EmbeddedKernelStore {
282    async fn list_memory_abouts(&self) -> Result<Vec<String>, PortError> {
283        self.run(|store| {
284            let tx = store.begin_read()?;
285            let anchors = tx.open_table(ANCHORS).map_err(table_error)?;
286            let mut abouts = Vec::new();
287            for row in anchors.iter().map_err(range_error)? {
288                let (key, _) = row.map_err(range_error)?;
289                abouts.push(key.value().to_string());
290            }
291            Ok(abouts)
292        })
293        .await
294    }
295
296    async fn list_memory_abouts_by_dimensions(
297        &self,
298        dimension_ids: &[String],
299    ) -> Result<Vec<String>, PortError> {
300        let dimension_ids = dimension_ids.to_vec();
301        self.run(move |store| {
302            let tx = store.begin_read()?;
303            let anchors = tx.open_table(ANCHORS).map_err(table_error)?;
304            let nodes = tx.open_table(NODES).map_err(table_error)?;
305            let relations = tx.open_table(RELATIONS).map_err(table_error)?;
306
307            let mut abouts = BTreeSet::new();
308            for row in anchors.iter().map_err(range_error)? {
309                let (key, _) = row.map_err(range_error)?;
310                let anchor = key.value().to_string();
311                let is_anchor = load_node(&nodes, &anchor)?
312                    .is_some_and(|node| node.node_kind == MEMORY_ANCHOR_KIND);
313                if !is_anchor {
314                    continue;
315                }
316                for relation in outgoing_rows(&relations, &anchor)? {
317                    if relation.relation_type != "has_dimension" {
318                        continue;
319                    }
320                    let matches =
321                        load_node(&nodes, &relation.target_node_id)?.is_some_and(|dimension| {
322                            dimension.node_kind == "memory_dimension"
323                                && dimension_ids.iter().any(|dimension_id| {
324                                    dimension.node_id == *dimension_id
325                                        || dimension
326                                            .node_id
327                                            .ends_with(&format!(":dimension:{dimension_id}"))
328                                })
329                        });
330                    if matches {
331                        abouts.insert(anchor.clone());
332                        break;
333                    }
334                }
335            }
336            Ok(abouts.into_iter().collect())
337        })
338        .await
339    }
340}
kmp_adapter_embedded/adapter/graph_read.rs

kmp_adapter_embedded/adapter/
graph_read.rs
