kmp_adapter_embedded/adapter/

graph_read.rs

use std::collections::{BTreeMap, BTreeSet, VecDeque};

use kmp_domain::{
    ContextPathNeighborhood, GraphNeighborhoodReader, NeighborhoodRequest, NodeNeighborhood,
    NodeProjection, NodeRelationProjection, NodeRelationshipReader, NodeRelationships, PortError,
};

use super::engine::{Key, ReadTx, Table};
use super::outward_neighborhood::OutwardNeighborhoodRead;
use super::serdes::{NodeRecord, decode, decode_explanation};
use super::store::EmbeddedKernelStore;
#[path = "node_admission_header.rs"]
mod node_admission_header;

fn load_node(tx: &dyn ReadTx, node_id: &str) -> Result<Option<NodeProjection>, PortError> {
    match tx.get(Table::Nodes, Key::Str(node_id))? {
        Some(raw) => Ok(Some(
            decode::<NodeRecord>("graph node", &raw)?.into_projection()?,
        )),
        None => Ok(None),
    }
}

fn outgoing_rows(tx: &dyn ReadTx, source: &str) -> Result<Vec<NodeRelationProjection>, PortError> {
    tx.scan_str3_by_first(Table::Relations, source)?
        .into_iter()
        .map(|((source_node_id, target_node_id, relation_type), raw)| {
            Ok(NodeRelationProjection {
                source_node_id,
                target_node_id,
                relation_type,
                explanation: decode_explanation(&raw)?,
            })
        })
        .collect()
}

fn outgoing_targets(tx: &dyn ReadTx, source: &str) -> Result<Vec<String>, PortError> {
    Ok(tx
        .scan_str3_by_first(Table::Relations, source)?
        .into_iter()
        .map(|((_, target, _), _)| target)
        .collect())
}

fn selected_projections(
    tx: &dyn ReadTx,
    selected: &BTreeSet<String>,
    root_node_id: &str,
) -> Result<Vec<NodeProjection>, PortError> {
    let mut projections = Vec::new();
    for node_id in selected {
        if node_id == root_node_id {
            continue;
        }
        if let Some(projection) = load_node(tx, node_id)? {
            projections.push(projection);
        }
    }
    Ok(projections)
}

fn shortest_outward_path(
    tx: &dyn ReadTx,
    root_node_id: &str,
    target_node_id: &str,
) -> Result<Option<Vec<String>>, PortError> {
    let mut predecessors = BTreeMap::<String, String>::new();
    let mut visited = BTreeSet::from([root_node_id.to_string()]);
    let mut frontier = VecDeque::from([root_node_id.to_string()]);

    while let Some(node_id) = frontier.pop_front() {
        for target in outgoing_targets(tx, &node_id)? {
            if !visited.insert(target.clone()) {
                continue;
            }
            predecessors.insert(target.clone(), node_id.clone());
            if target == target_node_id {
                let mut path = vec![target.clone()];
                let mut current = target.as_str();
                while let Some(previous) = predecessors.get(current) {
                    path.push(previous.clone());
                    current = previous;
                }
                path.reverse();
                return Ok(Some(path));
            }
            frontier.push_back(target);
        }
    }

    Ok(None)
}

impl EmbeddedKernelStore {
    async fn read_catalogue_neighborhood(
        &self,
        request: &NeighborhoodRequest,
        headers: bool,
    ) -> Result<Option<NodeNeighborhood>, PortError> {
        let request = request.clone();
        self.run(move |store| {
            let tx = store.begin_read()?;
            let tx = tx.as_ref();
            let read = |id: &str| {
                if headers {
                    node_admission_header::read(tx, id)
                } else {
                    load_node(tx, id)
                }
            };
            let Some(root) = read(request.root_node_id())? else {
                return Ok(None);
            };
            let (reachable, relations) = OutwardNeighborhoodRead::new(tx, &request).catalogue()?;
            let mut neighbors = Vec::with_capacity(reachable.len());
            for id in reachable {
                if let Some(node) = read(&id)? {
                    neighbors.push(node);
                }
            }
            Ok(Some(NodeNeighborhood {
                root,
                neighbors,
                relations,
            }))
        })
        .await
    }
}

impl GraphNeighborhoodReader for EmbeddedKernelStore {
    async fn graph_read_revision(
        &self,
    ) -> Result<Option<kmp_domain::GraphReadRevision>, PortError> {
        Ok(self.read_revision())
    }

    async fn load_nodes_batch(
        &self,
        node_ids: Vec<String>,
    ) -> Result<Vec<Option<NodeProjection>>, PortError> {
        self.run(move |store| {
            let tx = store.begin_read()?;
            node_ids
                .iter()
                .map(|id| load_node(tx.as_ref(), id))
                .collect()
        })
        .await
    }

    async fn load_bounded_trace(
        &self,
        request: &kmp_domain::TraceSearchRequest,
    ) -> Result<kmp_domain::TraceSearchResult, PortError> {
        let request = request.clone();
        self.run(move |store| {
            let tx = store.begin_read()?;
            kmp_domain::bounded_trace_search(
                &super::trace_snapshot::TraceSnapshot(tx.as_ref()),
                &request,
            )
        })
        .await
    }

    async fn load_memory_nodes(
        &self,
        request: &kmp_domain::MemoryNodesRequest,
    ) -> Result<kmp_domain::MemoryNodesResult, PortError> {
        let request = request.clone();
        self.run(move |store| {
            let pinned = store.pin_snapshot()?;
            let revision = pinned.read_revision().ok_or_else(|| {
                PortError::Conflict(
                    "store changed while pinning the node batch; retry the original read".into(),
                )
            })?;
            if request
                .expect_snapshot
                .as_ref()
                .is_some_and(|expected| expected != &revision)
            {
                return Err(PortError::Conflict(
                    "node batch snapshot changed; discard previous batches and restart the focus"
                        .into(),
                ));
            }
            let tx = pinned.begin_read()?;
            let mut result = kmp_domain::read_memory_nodes(
                &super::trace_snapshot::TraceSnapshot(tx.as_ref()),
                &request,
            )?;
            result.snapshot = Some(revision);
            Ok(result)
        })
        .await
    }

    async fn load_evidence_paths(
        &self,
        request: &kmp_domain::EvidencePathRequest,
    ) -> Result<kmp_domain::EvidencePathResult, PortError> {
        let request = request.clone();
        self.run(move |store| {
            let tx = store.begin_read()?;
            kmp_domain::search_evidence_paths(
                &super::trace_snapshot::TraceSnapshot(tx.as_ref()),
                &request,
            )
        })
        .await
    }

    async fn load_neighborhood(
        &self,
        root_node_id: &str,
        depth: u32,
    ) -> Result<Option<NodeNeighborhood>, PortError> {
        self.load_scoped_neighborhood(&NeighborhoodRequest::new(root_node_id, depth))
            .await
    }

    async fn load_scoped_neighborhood(
        &self,
        request: &NeighborhoodRequest,
    ) -> Result<Option<NodeNeighborhood>, PortError> {
        self.read_catalogue_neighborhood(request, false).await
    }

    async fn load_neighborhood_headers(
        &self,
        request: &NeighborhoodRequest,
    ) -> Result<Option<NodeNeighborhood>, PortError> {
        self.read_catalogue_neighborhood(request, true).await
    }

    async fn load_context_path(
        &self,
        root_node_id: &str,
        target_node_id: &str,
        subtree_depth: u32,
    ) -> Result<Option<ContextPathNeighborhood>, PortError> {
        let root_node_id = root_node_id.to_string();
        let target_node_id = target_node_id.to_string();
        self.run(move |store| {
            let tx = store.begin_read()?;
            let tx = tx.as_ref();

            let Some(root) = load_node(tx, &root_node_id)? else {
                return Ok(None);
            };
            if load_node(tx, &target_node_id)?.is_none() {
                return Ok(None);
            }
            let Some(path_node_ids) = shortest_outward_path(tx, &root_node_id, &target_node_id)?
            else {
                return Ok(None);
            };

            let subtree = NeighborhoodRequest::new(&target_node_id, subtree_depth);
            let (selected, relations) = OutwardNeighborhoodRead::new(tx, &subtree)
                .extending(path_node_ids.iter().cloned().collect())?;

            Ok(Some(ContextPathNeighborhood {
                neighbors: selected_projections(tx, &selected, &root_node_id)?,
                relations,
                path_node_ids,
                root,
            }))
        })
        .await
    }
}

impl NodeRelationshipReader for EmbeddedKernelStore {
    async fn load_node_relationships(
        &self,
        node_id: &str,
    ) -> Result<Option<NodeRelationships>, PortError> {
        let node_id = node_id.to_string();
        self.run(move |store| {
            let tx = store.begin_read()?;
            let tx = tx.as_ref();
            if load_node(tx, &node_id)?.is_none() {
                return Ok(None);
            }

            let mut incoming = Vec::new();
            for ((target, source, relation_type), _) in
                tx.scan_str3_by_first(Table::RelationsByTarget, &node_id)?
            {
                let Some(raw) = tx.get(
                    Table::Relations,
                    Key::Str3(&source, &target, &relation_type),
                )?
                else {
                    return Err(PortError::InvalidState(format!(
                        "embedded store adjacency index points at missing relation \
                         `{source}` -> `{target}` ({relation_type})"
                    )));
                };
                incoming.push(NodeRelationProjection {
                    explanation: decode_explanation(&raw)?,
                    source_node_id: source,
                    target_node_id: target,
                    relation_type,
                });
            }

            Ok(Some(NodeRelationships {
                incoming,
                outgoing: outgoing_rows(tx, &node_id)?,
            }))
        })
        .await
    }
}