ucp_graph/store/
document.rs

1use std::collections::HashMap;
2
3use chrono::Utc;
4use ucm_core::{BlockId, Document, PortableDocument};
5
6use super::{GraphNodeRecord, GraphStore, GraphStoreObservability, GraphStoreStats};
7use crate::types::GraphEdgeSummary;
8
9#[derive(Debug, Clone)]
10pub struct InMemoryGraphStore {
11    document: Document,
12    parent_by_child: HashMap<BlockId, BlockId>,
13    label_index: HashMap<String, BlockId>,
14    stats: GraphStoreStats,
15}
16
17impl InMemoryGraphStore {
18    pub fn from_document(document: Document) -> Self {
19        let mut parent_by_child = HashMap::new();
20        for (parent, children) in &document.structure {
21            for child in children {
22                parent_by_child.insert(*child, *parent);
23            }
24        }
25
26        let label_index = document
27            .blocks
28            .values()
29            .filter_map(|block| {
30                block
31                    .metadata
32                    .label
33                    .as_ref()
34                    .map(|label| (label.clone(), block.id))
35            })
36            .collect::<HashMap<_, _>>();
37
38        let explicit_edge_count = document
39            .blocks
40            .values()
41            .map(|block| block.edges.len())
42            .sum();
43        let structural_edge_count = parent_by_child.len();
44        let stats = GraphStoreStats {
45            backend: "memory".to_string(),
46            document_id: document.id.0.clone(),
47            root_block_id: document.root,
48            node_count: document.blocks.len(),
49            explicit_edge_count,
50            structural_edge_count,
51            captured_at: Utc::now(),
52            graph_key: None,
53        };
54
55        Self {
56            document,
57            parent_by_child,
58            label_index,
59            stats,
60        }
61    }
62
63    pub fn from_portable(portable: &PortableDocument) -> Result<Self, super::GraphStoreError> {
64        Ok(Self::from_document(portable.to_document()?))
65    }
66
67    pub fn from_json(payload: &str) -> Result<Self, super::GraphStoreError> {
68        let portable: PortableDocument = serde_json::from_str(payload)?;
69        Self::from_portable(&portable)
70    }
71
72    pub fn document(&self) -> &Document {
73        &self.document
74    }
75}
76
77impl GraphStore for InMemoryGraphStore {
78    fn stats(&self) -> GraphStoreStats {
79        self.stats.clone()
80    }
81
82    fn observability(&self) -> GraphStoreObservability {
83        GraphStoreObservability {
84            stats: self.stats(),
85            indexed_fields: vec![
86                "block_id".to_string(),
87                "label".to_string(),
88                "parent".to_string(),
89                "content_type".to_string(),
90            ],
91        }
92    }
93
94    fn root_id(&self) -> BlockId {
95        self.document.root
96    }
97
98    fn node_ids(&self) -> Vec<BlockId> {
99        let mut ids = self.document.blocks.keys().copied().collect::<Vec<_>>();
100        ids.sort_by_key(|id| id.to_string());
101        ids
102    }
103
104    fn node(&self, block_id: BlockId) -> Option<GraphNodeRecord> {
105        let block = self.document.get_block(&block_id)?;
106        Some(GraphNodeRecord {
107            block_id,
108            label: block.metadata.label.clone(),
109            content_type: block.content_type().to_string(),
110            semantic_role: block
111                .metadata
112                .semantic_role
113                .as_ref()
114                .map(ToString::to_string),
115            tags: block.metadata.tags.clone(),
116            parent: self.parent_by_child.get(&block_id).copied(),
117            children: self.document.children(&block_id).len(),
118            outgoing_edges: self.document.edge_index.outgoing_from(&block_id).len(),
119            incoming_edges: self.document.edge_index.incoming_to(&block_id).len(),
120        })
121    }
122
123    fn children(&self, block_id: BlockId) -> Vec<BlockId> {
124        self.document.children(&block_id).to_vec()
125    }
126
127    fn parent(&self, block_id: BlockId) -> Option<BlockId> {
128        self.parent_by_child.get(&block_id).copied()
129    }
130
131    fn outgoing_edges(&self, block_id: BlockId) -> Vec<GraphEdgeSummary> {
132        self.document
133            .edge_index
134            .outgoing_from(&block_id)
135            .iter()
136            .map(|(edge_type, target)| GraphEdgeSummary {
137                source: block_id,
138                target: *target,
139                relation: edge_relation(edge_type),
140                direction: "outgoing".to_string(),
141            })
142            .collect()
143    }
144
145    fn incoming_edges(&self, block_id: BlockId) -> Vec<GraphEdgeSummary> {
146        self.document
147            .edge_index
148            .incoming_to(&block_id)
149            .iter()
150            .map(|(edge_type, source)| GraphEdgeSummary {
151                source: *source,
152                target: block_id,
153                relation: edge_relation(edge_type),
154                direction: "incoming".to_string(),
155            })
156            .collect()
157    }
158
159    fn resolve_selector(&self, selector: &str) -> Option<BlockId> {
160        if selector == "root" {
161            return Some(self.document.root);
162        }
163        selector
164            .parse::<BlockId>()
165            .ok()
166            .filter(|id| self.document.blocks.contains_key(id))
167            .or_else(|| self.label_index.get(selector).copied())
168    }
169
170    fn to_portable_document(&self) -> Result<PortableDocument, super::GraphStoreError> {
171        Ok(self.document.to_portable())
172    }
173}
174
175fn edge_relation(edge_type: &ucm_core::EdgeType) -> String {
176    match edge_type {
177        ucm_core::EdgeType::Custom(value) => value.clone(),
178        _ => serde_json::to_value(edge_type)
179            .ok()
180            .and_then(|value| value.as_str().map(ToOwned::to_owned))
181            .unwrap_or_else(|| format!("{:?}", edge_type).to_lowercase()),
182    }
183}
ucp_graph/store/document.rs

ucp_graph/store/
document.rs
