Skip to main content

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