Skip to main content

mentedb_graph/
csr.rs

1//! Compressed Sparse Row/Column graph storage with delta log for incremental updates.
2
3use ahash::HashMap;
4use mentedb_core::edge::{EdgeType, MemoryEdge};
5use mentedb_core::error::{MenteError, MenteResult};
6use mentedb_core::types::{MemoryId, Timestamp};
7use serde::{Deserialize, Serialize};
8
9/// Compact edge data stored in CSR/CSC arrays.
10#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub struct StoredEdge {
12    /// The relationship type.
13    pub edge_type: EdgeType,
14    /// Edge weight (0.0 to 1.0).
15    pub weight: f32,
16    /// When this edge was created.
17    pub created_at: Timestamp,
18    /// When this relationship became valid. None = since creation.
19    #[serde(default, skip_serializing_if = "Option::is_none")]
20    pub valid_from: Option<Timestamp>,
21    /// When this relationship stopped being valid. None = still valid.
22    #[serde(default, skip_serializing_if = "Option::is_none")]
23    pub valid_until: Option<Timestamp>,
24    /// Semantic label for the relationship (e.g. "owns", "attends").
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub label: Option<String>,
27}
28
29impl StoredEdge {
30    /// Converts a [`MemoryEdge`] into a compact stored representation.
31    pub fn from_memory_edge(edge: &MemoryEdge) -> Self {
32        Self {
33            edge_type: edge.edge_type,
34            weight: edge.weight,
35            created_at: edge.created_at,
36            valid_from: edge.valid_from,
37            valid_until: edge.valid_until,
38            label: edge.label.clone(),
39        }
40    }
41
42    /// Returns true if this edge is temporally valid at the given timestamp.
43    pub fn is_valid_at(&self, at: Timestamp) -> bool {
44        let from = self.valid_from.unwrap_or(0);
45        match self.valid_until {
46            Some(until) => at >= from && at < until,
47            None => at >= from,
48        }
49    }
50
51    /// Returns true if this edge has been invalidated.
52    pub fn is_invalidated(&self) -> bool {
53        self.valid_until.is_some()
54    }
55}
56
57/// A pending edge in the delta log before compaction into CSR.
58#[derive(Debug, Clone, Serialize, Deserialize)]
59struct DeltaEdge {
60    source_idx: u32,
61    target_idx: u32,
62    data: StoredEdge,
63}
64
65/// Compressed Sparse Row storage for one direction (outgoing or incoming).
66#[derive(Debug, Clone, Default, Serialize, Deserialize)]
67struct CompressedStorage {
68    /// Length = num_nodes + 1. row_offsets[i]..row_offsets[i+1] gives the range in col_indices/edge_data.
69    row_offsets: Vec<u32>,
70    /// Column indices (target node for CSR, source node for CSC).
71    col_indices: Vec<u32>,
72    /// Edge metadata parallel to col_indices.
73    edge_data: Vec<StoredEdge>,
74}
75
76impl CompressedStorage {
77    #[allow(dead_code)]
78    fn new(num_nodes: usize) -> Self {
79        Self {
80            row_offsets: vec![0; num_nodes + 1],
81            col_indices: Vec::new(),
82            edge_data: Vec::new(),
83        }
84    }
85
86    /// Get neighbors and edge data for a given row index.
87    fn neighbors(&self, row: u32) -> &[u32] {
88        let row = row as usize;
89        if row + 1 >= self.row_offsets.len() {
90            return &[];
91        }
92        let start = self.row_offsets[row] as usize;
93        let end = self.row_offsets[row + 1] as usize;
94        &self.col_indices[start..end]
95    }
96
97    fn edge_data_for(&self, row: u32) -> &[StoredEdge] {
98        let row = row as usize;
99        if row + 1 >= self.row_offsets.len() {
100            return &[];
101        }
102        let start = self.row_offsets[row] as usize;
103        let end = self.row_offsets[row + 1] as usize;
104        &self.edge_data[start..end]
105    }
106}
107
108/// Bidirectional graph with CSR (outgoing) and CSC (incoming) plus a delta log.
109#[derive(Debug, Clone, Serialize, Deserialize)]
110pub struct CsrGraph {
111    /// Maps MemoryId -> internal u32 index.
112    id_to_idx: HashMap<MemoryId, u32>,
113    /// Maps internal u32 index -> MemoryId.
114    idx_to_id: Vec<MemoryId>,
115
116    /// CSR for outgoing edges.
117    csr: CompressedStorage,
118    /// CSC for incoming edges.
119    csc: CompressedStorage,
120
121    /// Recent edges not yet merged into the compressed storage.
122    delta_edges: Vec<DeltaEdge>,
123    /// Edges marked for removal (source_idx, target_idx).
124    removed_edges: Vec<(u32, u32)>,
125
126    /// Live edge count per type, maintained incrementally so counts (e.g. the
127    /// dashboard conflict badge) are O(1) instead of a walk over every node ever
128    /// created. The detectors store one directed edge per conflict pair, so a
129    /// directed count equals the pair count. Grounded authoritatively on
130    /// `compact` and after load; adjusted on add and per-node removal.
131    #[serde(default)]
132    edge_type_counts: HashMap<EdgeType, usize>,
133
134    /// Live conflict-class edges (contradicts, supersedes) with their weight,
135    /// maintained incrementally so the dashboard conflict view enumerates the
136    /// actual pairs in O(conflicts) instead of walking every node's out-edges
137    /// (which rescans the delta log per node, O(nodes x delta)). Keyed by
138    /// `(source, target, type)` so a pair carrying two different edge types does
139    /// not collide. Skipped on serialize (a tuple key is not a JSON map key) and
140    /// reground from the live edges by `recompute_edge_counts` on load, so it
141    /// costs nothing on disk and is always correct after a restart.
142    #[serde(skip)]
143    conflict_edges: HashMap<(MemoryId, MemoryId, EdgeType), f32>,
144}
145
146/// The edge types the dashboard treats as conflicts, maintained as a live pair
147/// list for O(conflicts) enumeration.
148fn is_conflict_type(t: EdgeType) -> bool {
149    matches!(t, EdgeType::Contradicts | EdgeType::Supersedes)
150}
151
152impl CsrGraph {
153    /// Creates a new empty CSR graph.
154    pub fn new() -> Self {
155        Self {
156            id_to_idx: HashMap::default(),
157            idx_to_id: Vec::new(),
158            csr: CompressedStorage::default(),
159            csc: CompressedStorage::default(),
160            delta_edges: Vec::new(),
161            removed_edges: Vec::new(),
162            edge_type_counts: HashMap::default(),
163            conflict_edges: HashMap::default(),
164        }
165    }
166
167    /// Register a node. Returns its internal index.
168    pub fn add_node(&mut self, id: MemoryId) -> u32 {
169        if let Some(&idx) = self.id_to_idx.get(&id) {
170            return idx;
171        }
172        let idx = self.idx_to_id.len() as u32;
173        self.id_to_idx.insert(id, idx);
174        self.idx_to_id.push(id);
175        idx
176    }
177
178    /// Remove a node and all its edges.
179    pub fn remove_node(&mut self, id: MemoryId) {
180        let Some(&idx) = self.id_to_idx.get(&id) else {
181            return;
182        };
183
184        // Collect the types of this node's still-live edges so the running
185        // per-type counts can be decremented. `is_removed` guards against
186        // decrementing an edge twice when both of its endpoints are removed.
187        let mut decrements: Vec<EdgeType> = Vec::new();
188        {
189            let out_neighbors = self.csr.neighbors(idx).to_vec();
190            let out_types: Vec<EdgeType> = self
191                .csr
192                .edge_data_for(idx)
193                .iter()
194                .map(|e| e.edge_type)
195                .collect();
196            for (i, &neighbor) in out_neighbors.iter().enumerate() {
197                if !self.is_removed(idx, neighbor) {
198                    decrements.push(out_types[i]);
199                }
200            }
201            let in_neighbors = self.csc.neighbors(idx).to_vec();
202            let in_types: Vec<EdgeType> = self
203                .csc
204                .edge_data_for(idx)
205                .iter()
206                .map(|e| e.edge_type)
207                .collect();
208            for (i, &neighbor) in in_neighbors.iter().enumerate() {
209                if !self.is_removed(neighbor, idx) {
210                    decrements.push(in_types[i]);
211                }
212            }
213            for e in &self.delta_edges {
214                if e.source_idx == idx || e.target_idx == idx {
215                    decrements.push(e.data.edge_type);
216                }
217            }
218        }
219        for et in decrements {
220            if let Some(c) = self.edge_type_counts.get_mut(&et) {
221                *c = c.saturating_sub(1);
222            }
223        }
224
225        // Mark all outgoing and incoming edges for removal
226        for &neighbor in self.csr.neighbors(idx) {
227            self.removed_edges.push((idx, neighbor));
228        }
229        for &neighbor in self.csc.neighbors(idx) {
230            self.removed_edges.push((neighbor, idx));
231        }
232        // Also remove from delta
233        self.delta_edges
234            .retain(|e| e.source_idx != idx && e.target_idx != idx);
235        // Drop any conflict edges touching this node.
236        self.conflict_edges
237            .retain(|(s, t, _), _| *s != id && *t != id);
238        self.id_to_idx.remove(&id);
239    }
240
241    /// Add an edge to the delta log.
242    ///
243    /// Deduplicates: if a currently-valid edge with the same source, target,
244    /// and type already exists (in CSR or the delta log), the graph is left
245    /// unchanged. Write inference and pipeline passes may infer the same
246    /// relationship more than once; parallel duplicates carry no information.
247    pub fn add_edge(&mut self, edge: &MemoryEdge) {
248        let source_idx = self.add_node(edge.source);
249        let target_idx = self.add_node(edge.target);
250
251        let duplicate = self.outgoing_by_idx(source_idx).into_iter().any(|(t, e)| {
252            t == edge.target && e.edge_type == edge.edge_type && e.valid_until.is_none()
253        });
254        if duplicate {
255            return;
256        }
257
258        self.delta_edges.push(DeltaEdge {
259            source_idx,
260            target_idx,
261            data: StoredEdge::from_memory_edge(edge),
262        });
263        *self.edge_type_counts.entry(edge.edge_type).or_default() += 1;
264        if is_conflict_type(edge.edge_type) {
265            self.conflict_edges
266                .insert((edge.source, edge.target, edge.edge_type), edge.weight);
267        }
268    }
269
270    /// Strengthen an edge by incrementing its weight (Hebbian learning).
271    ///
272    /// Updates the existing edge rather than appending a parallel duplicate:
273    /// a delta edge is bumped in place; a compressed (CSR) edge is marked
274    /// removed and replaced with a delta override, so compaction keeps
275    /// exactly one copy.
276    pub fn strengthen_edge(&mut self, source: MemoryId, target: MemoryId, delta: f32) {
277        let (Some(&source_idx), Some(&target_idx)) =
278            (self.id_to_idx.get(&source), self.id_to_idx.get(&target))
279        else {
280            return;
281        };
282
283        // Keep the conflict-edge weights in step with the live edge (a pair may
284        // carry either conflict type).
285        for t in [EdgeType::Contradicts, EdgeType::Supersedes] {
286            if let Some(w) = self.conflict_edges.get_mut(&(source, target, t)) {
287                *w = (*w + delta).min(1.0);
288            }
289        }
290
291        // Bump an existing delta edge in place.
292        if let Some(existing) = self
293            .delta_edges
294            .iter_mut()
295            .find(|e| e.source_idx == source_idx && e.target_idx == target_idx)
296        {
297            existing.data.weight = (existing.data.weight + delta).min(1.0);
298            return;
299        }
300
301        // Edge lives in compressed storage: suppress the CSR copy and push a
302        // delta override. Removed-edge filtering only applies to compressed
303        // edges, so the override remains visible.
304        if let Some((_, stored)) = self
305            .outgoing_by_idx(source_idx)
306            .into_iter()
307            .find(|(id, _)| *id == target)
308        {
309            let new_weight = (stored.weight + delta).min(1.0);
310            self.removed_edges.push((source_idx, target_idx));
311            self.delta_edges.push(DeltaEdge {
312                source_idx,
313                target_idx,
314                data: StoredEdge {
315                    weight: new_weight,
316                    ..stored
317                },
318            });
319        }
320    }
321
322    /// Mark an edge for removal.
323    pub fn remove_edge(&mut self, source: MemoryId, target: MemoryId) {
324        let (Some(&src_idx), Some(&tgt_idx)) =
325            (self.id_to_idx.get(&source), self.id_to_idx.get(&target))
326        else {
327            return;
328        };
329        self.removed_edges.push((src_idx, tgt_idx));
330        self.delta_edges
331            .retain(|e| !(e.source_idx == src_idx && e.target_idx == tgt_idx));
332        // Pair-granular removal drops every edge between the pair.
333        self.conflict_edges
334            .retain(|(s, t, _), _| !(*s == source && *t == target));
335    }
336
337    /// Remove every edge whose type is in `types`, preserving all other edges,
338    /// including edges of a different type between the same pair. Returns the
339    /// number of edges removed.
340    ///
341    /// The removal model suppresses a compressed edge at pair granularity, so
342    /// when a suppressed pair also carries a non-matching edge, that sibling is
343    /// re-added to the delta log to keep it. Callers should `compact()` after to
344    /// materialize the result into both CSR and CSC. This is a one-time cleanup
345    /// primitive (used to purge conflict edges the old write-time heuristic
346    /// created at ~0% precision), not a hot path.
347    pub fn remove_edges_of_types(&mut self, types: &[EdgeType]) -> usize {
348        let matches = |et: EdgeType| types.contains(&et);
349        let mut removed = 0usize;
350
351        // Delta edges: drop matching ones directly.
352        let before = self.delta_edges.len();
353        self.delta_edges.retain(|e| !matches(e.data.edge_type));
354        removed += before - self.delta_edges.len();
355
356        // Compressed edges: suppress matching ones. Suppression is pair level,
357        // so restore any non-matching sibling of a suppressed pair via delta.
358        let mut restore: Vec<DeltaEdge> = Vec::new();
359        let n = self.idx_to_id.len() as u32;
360        for src_idx in 0..n {
361            let row: Vec<(u32, StoredEdge)> = {
362                let neighbors = self.csr.neighbors(src_idx);
363                let edges = self.csr.edge_data_for(src_idx);
364                let mut v = Vec::new();
365                for (i, &t) in neighbors.iter().enumerate() {
366                    if !self.is_removed(src_idx, t) {
367                        v.push((t, edges[i].clone()));
368                    }
369                }
370                v
371            };
372            let matched_pairs: std::collections::HashSet<u32> = row
373                .iter()
374                .filter(|(_, e)| matches(e.edge_type))
375                .map(|(t, _)| *t)
376                .collect();
377            if matched_pairs.is_empty() {
378                continue;
379            }
380            for &tgt_idx in &matched_pairs {
381                self.removed_edges.push((src_idx, tgt_idx));
382            }
383            for (tgt_idx, e) in row {
384                if !matched_pairs.contains(&tgt_idx) {
385                    continue;
386                }
387                if matches(e.edge_type) {
388                    removed += 1;
389                } else {
390                    restore.push(DeltaEdge {
391                        source_idx: src_idx,
392                        target_idx: tgt_idx,
393                        data: e,
394                    });
395                }
396            }
397        }
398        self.delta_edges.extend(restore);
399        // Bulk type removal touches many edges via suppression and restore; just
400        // recount from the resulting live set rather than track each delta.
401        self.recompute_edge_counts();
402        removed
403    }
404
405    /// Get all outgoing edges from a node (CSR + delta, minus removed).
406    pub fn outgoing(&self, id: MemoryId) -> Vec<(MemoryId, StoredEdge)> {
407        let Some(&idx) = self.id_to_idx.get(&id) else {
408            return Vec::new();
409        };
410        self.outgoing_by_idx(idx)
411    }
412
413    /// Get outgoing edges that are temporally valid at the given timestamp.
414    pub fn outgoing_valid_at(&self, id: MemoryId, at: Timestamp) -> Vec<(MemoryId, StoredEdge)> {
415        self.outgoing(id)
416            .into_iter()
417            .filter(|(_, e)| e.is_valid_at(at))
418            .collect()
419    }
420
421    pub(crate) fn outgoing_by_idx(&self, idx: u32) -> Vec<(MemoryId, StoredEdge)> {
422        let mut results = Vec::new();
423
424        // From compressed storage
425        let neighbors = self.csr.neighbors(idx);
426        let edges = self.csr.edge_data_for(idx);
427        for (i, &neighbor) in neighbors.iter().enumerate() {
428            if !self.is_removed(idx, neighbor)
429                && let Some(&id) = self.idx_to_id.get(neighbor as usize)
430            {
431                results.push((id, edges[i].clone()));
432            }
433        }
434
435        // From delta
436        for delta in &self.delta_edges {
437            if delta.source_idx == idx
438                && let Some(&id) = self.idx_to_id.get(delta.target_idx as usize)
439            {
440                results.push((id, delta.data.clone()));
441            }
442        }
443
444        results
445    }
446
447    /// `(contradicts, supersedes)` live conflict-edge counts, read O(1) from the
448    /// running per-type counts. The detectors store one directed edge per
449    /// conflict pair, so a directed count is the pair count.
450    pub fn conflict_edge_counts(&self) -> (usize, usize) {
451        (
452            self.edge_type_counts
453                .get(&EdgeType::Contradicts)
454                .copied()
455                .unwrap_or(0),
456            self.edge_type_counts
457                .get(&EdgeType::Supersedes)
458                .copied()
459                .unwrap_or(0),
460        )
461    }
462
463    /// Every live conflict-class edge as `(source, target, type, weight)`, read
464    /// in O(conflicts) from the maintained pair list, so the dashboard conflict
465    /// view never walks the whole graph. Grounded on compact/load, so it stays
466    /// correct across a restart.
467    pub fn conflict_edges(&self) -> Vec<(MemoryId, MemoryId, EdgeType, f32)> {
468        self.conflict_edges
469            .iter()
470            .map(|(&(s, t, et), &w)| (s, t, et, w))
471            .collect()
472    }
473
474    /// Recount live edges by type from scratch. Authoritative but O(nodes +
475    /// edges); called only when the graph is already being walked wholesale
476    /// (`compact`, bulk type removal) or grounded after load, never on a read.
477    pub fn recompute_edge_counts(&mut self) {
478        let mut counts: HashMap<EdgeType, usize> = HashMap::default();
479        let mut conflicts: HashMap<(MemoryId, MemoryId, EdgeType), f32> = HashMap::default();
480        let node_count = self.idx_to_id.len() as u32;
481        for idx in 0..node_count {
482            let neighbors = self.csr.neighbors(idx);
483            let data = self.csr.edge_data_for(idx);
484            for (i, &nbr) in neighbors.iter().enumerate() {
485                if !self.is_removed(idx, nbr) {
486                    let et = data[i].edge_type;
487                    *counts.entry(et).or_default() += 1;
488                    if is_conflict_type(et)
489                        && let (Some(&s), Some(&t)) = (
490                            self.idx_to_id.get(idx as usize),
491                            self.idx_to_id.get(nbr as usize),
492                        )
493                    {
494                        conflicts.insert((s, t, et), data[i].weight);
495                    }
496                }
497            }
498        }
499        for d in &self.delta_edges {
500            if !self.is_removed(d.source_idx, d.target_idx) {
501                let et = d.data.edge_type;
502                *counts.entry(et).or_default() += 1;
503                if is_conflict_type(et)
504                    && let (Some(&s), Some(&t)) = (
505                        self.idx_to_id.get(d.source_idx as usize),
506                        self.idx_to_id.get(d.target_idx as usize),
507                    )
508                {
509                    conflicts.insert((s, t, et), d.data.weight);
510                }
511            }
512        }
513        self.edge_type_counts = counts;
514        self.conflict_edges = conflicts;
515    }
516
517    /// Get all incoming edges to a node (CSC + delta, minus removed).
518    pub fn incoming(&self, id: MemoryId) -> Vec<(MemoryId, StoredEdge)> {
519        let Some(&idx) = self.id_to_idx.get(&id) else {
520            return Vec::new();
521        };
522        self.incoming_by_idx(idx)
523    }
524
525    /// Get incoming edges that are temporally valid at the given timestamp.
526    pub fn incoming_valid_at(&self, id: MemoryId, at: Timestamp) -> Vec<(MemoryId, StoredEdge)> {
527        self.incoming(id)
528            .into_iter()
529            .filter(|(_, e)| e.is_valid_at(at))
530            .collect()
531    }
532
533    pub(crate) fn incoming_by_idx(&self, idx: u32) -> Vec<(MemoryId, StoredEdge)> {
534        let mut results = Vec::new();
535
536        // From compressed storage (CSC)
537        let neighbors = self.csc.neighbors(idx);
538        let edges = self.csc.edge_data_for(idx);
539        for (i, &neighbor) in neighbors.iter().enumerate() {
540            if !self.is_removed(neighbor, idx)
541                && let Some(&id) = self.idx_to_id.get(neighbor as usize)
542            {
543                results.push((id, edges[i].clone()));
544            }
545        }
546
547        // From delta
548        for delta in &self.delta_edges {
549            if delta.target_idx == idx
550                && let Some(&id) = self.idx_to_id.get(delta.source_idx as usize)
551            {
552                results.push((id, delta.data.clone()));
553            }
554        }
555
556        results
557    }
558
559    /// Check if a node exists in the graph.
560    pub fn contains_node(&self, id: MemoryId) -> bool {
561        self.id_to_idx.contains_key(&id)
562    }
563
564    /// Number of registered nodes.
565    pub fn node_count(&self) -> usize {
566        self.idx_to_id.len()
567    }
568
569    /// Resolve a MemoryId to its internal index.
570    pub(crate) fn get_idx(&self, id: MemoryId) -> Option<u32> {
571        self.id_to_idx.get(&id).copied()
572    }
573
574    /// Resolve an internal index to its MemoryId.
575    #[allow(dead_code)]
576    pub(crate) fn get_id(&self, idx: u32) -> Option<MemoryId> {
577        self.idx_to_id.get(idx as usize).copied()
578    }
579
580    /// All registered node IDs.
581    pub fn node_ids(&self) -> &[MemoryId] {
582        &self.idx_to_id
583    }
584
585    fn is_removed(&self, source: u32, target: u32) -> bool {
586        self.removed_edges
587            .iter()
588            .any(|&(s, t)| s == source && t == target)
589    }
590
591    /// Merge all delta edges and removals into the compressed CSR/CSC storage.
592    pub fn compact(&mut self) {
593        let num_nodes = self.idx_to_id.len();
594
595        // Collect all edges: existing (minus removed) + delta
596        let mut all_edges: Vec<(u32, u32, StoredEdge)> = Vec::new();
597
598        // Existing CSR edges
599        for row in 0..num_nodes {
600            let row = row as u32;
601            let neighbors = self.csr.neighbors(row);
602            let edges = self.csr.edge_data_for(row);
603            for (i, &col) in neighbors.iter().enumerate() {
604                if !self.is_removed(row, col) {
605                    all_edges.push((row, col, edges[i].clone()));
606                }
607            }
608        }
609
610        // Delta edges
611        for delta in &self.delta_edges {
612            all_edges.push((delta.source_idx, delta.target_idx, delta.data.clone()));
613        }
614
615        // Build CSR (sorted by source)
616        self.csr = Self::build_compressed(&all_edges, num_nodes, false);
617
618        // Build CSC (sorted by target)
619        self.csc = Self::build_compressed(&all_edges, num_nodes, true);
620
621        self.delta_edges.clear();
622        self.removed_edges.clear();
623
624        // Ground the running per-type counts against the freshly compacted
625        // edge set, so any incremental drift since the last compaction is reset.
626        self.recompute_edge_counts();
627    }
628
629    fn build_compressed(
630        edges: &[(u32, u32, StoredEdge)],
631        num_nodes: usize,
632        transpose: bool,
633    ) -> CompressedStorage {
634        // Count edges per row
635        let mut counts = vec![0u32; num_nodes];
636        for &(src, tgt, ref _data) in edges {
637            let row = if transpose { tgt } else { src };
638            if (row as usize) < num_nodes {
639                counts[row as usize] += 1;
640            }
641        }
642
643        // Build offsets via prefix sum
644        let mut row_offsets = vec![0u32; num_nodes + 1];
645        for i in 0..num_nodes {
646            row_offsets[i + 1] = row_offsets[i] + counts[i];
647        }
648
649        let total = row_offsets[num_nodes] as usize;
650        let mut col_indices = vec![0u32; total];
651        let mut edge_data = vec![
652            StoredEdge {
653                edge_type: EdgeType::Related,
654                weight: 0.0,
655                created_at: 0,
656                valid_from: None,
657                valid_until: None,
658                label: None,
659            };
660            total
661        ];
662
663        // Fill using write cursors
664        let mut cursors = row_offsets[..num_nodes].to_vec();
665        for &(src, tgt, ref data) in edges {
666            let (row, col) = if transpose { (tgt, src) } else { (src, tgt) };
667            if (row as usize) < num_nodes {
668                let pos = cursors[row as usize] as usize;
669                col_indices[pos] = col;
670                edge_data[pos] = data.clone();
671                cursors[row as usize] += 1;
672            }
673        }
674
675        CompressedStorage {
676            row_offsets,
677            col_indices,
678            edge_data,
679        }
680    }
681    /// Save the graph snapshot to a file.
682    pub fn save(&self, path: &std::path::Path) -> MenteResult<()> {
683        let data =
684            serde_json::to_vec(self).map_err(|e| MenteError::Serialization(e.to_string()))?;
685        // Atomic snapshot: write to a temp file, fsync, then rename over the
686        // old snapshot so a crash mid-save never leaves a truncated graph.
687        let tmp_path = path.with_extension("json.tmp");
688        {
689            use std::io::Write;
690            let mut file = std::fs::File::create(&tmp_path)?;
691            file.write_all(&data)?;
692            file.sync_data()?;
693        }
694        std::fs::rename(&tmp_path, path)?;
695        Ok(())
696    }
697
698    /// Load the graph from a file.
699    pub fn load(path: &std::path::Path) -> MenteResult<Self> {
700        let data = std::fs::read(path)?;
701        let mut graph: Self =
702            serde_json::from_slice(&data).map_err(|e| MenteError::Serialization(e.to_string()))?;
703        // Ground the per-type counts from the loaded edges. Snapshots written
704        // before this field existed carry an empty map; recomputing makes the
705        // counts correct regardless, and the edge-log replay in
706        // `GraphManager::open` then adjusts for changes since the snapshot.
707        graph.recompute_edge_counts();
708        Ok(graph)
709    }
710}
711
712impl Default for CsrGraph {
713    fn default() -> Self {
714        Self::new()
715    }
716}
717
718#[cfg(test)]
719mod tests {
720    use super::*;
721
722    fn make_edge(src: MemoryId, tgt: MemoryId, etype: EdgeType) -> MemoryEdge {
723        MemoryEdge {
724            source: src,
725            target: tgt,
726            edge_type: etype,
727            weight: 0.8,
728            created_at: 1000,
729            valid_from: None,
730            valid_until: None,
731            label: None,
732        }
733    }
734
735    #[test]
736    fn test_add_node_idempotent() {
737        let mut g = CsrGraph::new();
738        let id = MemoryId::new();
739        let idx1 = g.add_node(id);
740        let idx2 = g.add_node(id);
741        assert_eq!(idx1, idx2);
742        assert_eq!(g.node_count(), 1);
743    }
744
745    #[test]
746    fn edge_type_counts_track_add_compact_and_remove() {
747        let mut g = CsrGraph::new();
748        let (a, b, c, d) = (
749            MemoryId::new(),
750            MemoryId::new(),
751            MemoryId::new(),
752            MemoryId::new(),
753        );
754
755        // Adds land in the delta log and bump the running counts. Related edges
756        // are not conflicts and must not be counted.
757        g.add_edge(&make_edge(a, b, EdgeType::Contradicts));
758        g.add_edge(&make_edge(c, d, EdgeType::Supersedes));
759        g.add_edge(&make_edge(a, c, EdgeType::Related));
760        assert_eq!(g.conflict_edge_counts(), (1, 1));
761        // An exact-duplicate add is a no-op and must not double count.
762        g.add_edge(&make_edge(a, b, EdgeType::Contradicts));
763        assert_eq!(g.conflict_edge_counts(), (1, 1));
764
765        // Counts survive compaction (the CSR is rebuilt and the counts are
766        // reground from the compacted edge set).
767        g.compact();
768        assert_eq!(g.conflict_edge_counts(), (1, 1));
769
770        // A new edge after compaction lands in the delta log again.
771        g.add_edge(&make_edge(a, d, EdgeType::Supersedes));
772        assert_eq!(g.conflict_edge_counts(), (1, 2));
773
774        // Removing a node drops its edges from the counts: a->b (contradicts,
775        // now compacted) and a->d (supersedes, in delta). c->d supersedes stays.
776        g.remove_node(a);
777        assert_eq!(g.conflict_edge_counts(), (0, 1));
778    }
779
780    #[test]
781    fn conflict_edges_enumerated_maintained_and_persist() {
782        let mut g = CsrGraph::new();
783        let (a, b, c, d) = (
784            MemoryId::new(),
785            MemoryId::new(),
786            MemoryId::new(),
787            MemoryId::new(),
788        );
789        g.add_edge(&make_edge(a, b, EdgeType::Contradicts));
790        g.add_edge(&make_edge(c, d, EdgeType::Supersedes));
791        g.add_edge(&make_edge(a, c, EdgeType::Related)); // not a conflict
792
793        // Enumeration returns only the two conflict pairs, not the related edge.
794        let pairs = g.conflict_edges();
795        assert_eq!(pairs.len(), 2);
796        assert!(
797            pairs
798                .iter()
799                .any(|(s, t, k, _)| *s == a && *t == b && *k == EdgeType::Contradicts)
800        );
801        assert!(
802            pairs
803                .iter()
804                .any(|(s, t, k, _)| *s == c && *t == d && *k == EdgeType::Supersedes)
805        );
806
807        // Survives compaction (reground from the compacted set).
808        g.compact();
809        assert_eq!(g.conflict_edges().len(), 2);
810
811        // Persisted and reground on load, so a restart keeps the view O(conflicts).
812        let dir = tempfile::tempdir().unwrap();
813        let path = dir.path().join("graph.json");
814        g.save(&path).unwrap();
815        let reloaded = CsrGraph::load(&path).unwrap();
816        assert_eq!(reloaded.conflict_edges().len(), 2);
817
818        // Removing a node drops the conflict edge that touched it.
819        g.remove_node(a);
820        let pairs = g.conflict_edges();
821        assert_eq!(pairs.len(), 1);
822        assert_eq!(pairs[0].0, c);
823        assert_eq!(pairs[0].2, EdgeType::Supersedes);
824    }
825
826    #[test]
827    fn test_add_and_query_edges() {
828        let mut g = CsrGraph::new();
829        let a = MemoryId::new();
830        let b = MemoryId::new();
831        let c = MemoryId::new();
832
833        g.add_edge(&make_edge(a, b, EdgeType::Caused));
834        g.add_edge(&make_edge(a, c, EdgeType::Related));
835
836        let out = g.outgoing(a);
837        assert_eq!(out.len(), 2);
838
839        let inc_b = g.incoming(b);
840        assert_eq!(inc_b.len(), 1);
841        assert_eq!(inc_b[0].0, a);
842    }
843
844    #[test]
845    fn test_remove_edge() {
846        let mut g = CsrGraph::new();
847        let a = MemoryId::new();
848        let b = MemoryId::new();
849
850        g.add_edge(&make_edge(a, b, EdgeType::Caused));
851        assert_eq!(g.outgoing(a).len(), 1);
852
853        g.remove_edge(a, b);
854        assert_eq!(g.outgoing(a).len(), 0);
855    }
856
857    #[test]
858    fn test_remove_edges_of_types_preserves_siblings() {
859        let mut g = CsrGraph::new();
860        let a = MemoryId::new();
861        let b = MemoryId::new();
862        let c = MemoryId::new();
863
864        // a->b carries BOTH a Supersedes and a Related edge (the collateral
865        // case: pair-level suppression must not drop the Related sibling).
866        g.add_edge(&make_edge(a, b, EdgeType::Supersedes));
867        g.add_edge(&make_edge(a, b, EdgeType::Related));
868        // a->c is a lone Contradicts edge.
869        g.add_edge(&make_edge(a, c, EdgeType::Contradicts));
870        // b->c is unrelated and must survive.
871        g.add_edge(&make_edge(b, c, EdgeType::Caused));
872
873        // Compact first so the targets live in compressed storage, exercising
874        // the suppress-pair + restore-sibling path rather than delta retain.
875        g.compact();
876
877        let removed = g.remove_edges_of_types(&[EdgeType::Contradicts, EdgeType::Supersedes]);
878        g.compact();
879        assert_eq!(removed, 2, "one Supersedes + one Contradicts");
880
881        // a->b: Supersedes gone, Related preserved.
882        let ab: Vec<_> = g.outgoing(a).into_iter().filter(|(t, _)| *t == b).collect();
883        assert_eq!(ab.len(), 1);
884        assert_eq!(ab[0].1.edge_type, EdgeType::Related);
885
886        // a->c contradiction gone entirely.
887        assert!(g.outgoing(a).into_iter().all(|(t, _)| t != c));
888
889        // Unrelated edge intact; incoming reflects the cleanup (only b->c left).
890        assert_eq!(g.outgoing(b).len(), 1);
891        let inc_c = g.incoming(c);
892        assert_eq!(inc_c.len(), 1);
893        assert_eq!(inc_c[0].0, b);
894    }
895
896    #[test]
897    fn test_compact() {
898        let mut g = CsrGraph::new();
899        let a = MemoryId::new();
900        let b = MemoryId::new();
901        let c = MemoryId::new();
902
903        g.add_edge(&make_edge(a, b, EdgeType::Caused));
904        g.add_edge(&make_edge(b, c, EdgeType::Before));
905        g.compact();
906
907        let out_a = g.outgoing(a);
908        assert_eq!(out_a.len(), 1);
909        assert_eq!(out_a[0].0, b);
910
911        let inc_c = g.incoming(c);
912        assert_eq!(inc_c.len(), 1);
913        assert_eq!(inc_c[0].0, b);
914    }
915
916    #[test]
917    fn test_compact_with_removals() {
918        let mut g = CsrGraph::new();
919        let a = MemoryId::new();
920        let b = MemoryId::new();
921        let c = MemoryId::new();
922
923        g.add_edge(&make_edge(a, b, EdgeType::Caused));
924        g.add_edge(&make_edge(a, c, EdgeType::Related));
925        g.compact();
926
927        g.remove_edge(a, b);
928        g.compact();
929
930        let out = g.outgoing(a);
931        assert_eq!(out.len(), 1);
932        assert_eq!(out[0].0, c);
933    }
934
935    #[test]
936    fn test_remove_node_cleans_id_to_idx() {
937        let mut g = CsrGraph::new();
938        let a = MemoryId::new();
939        let b = MemoryId::new();
940
941        g.add_edge(&make_edge(a, b, EdgeType::Caused));
942        assert!(g.contains_node(a));
943        assert!(g.contains_node(b));
944
945        g.remove_node(a);
946        assert!(
947            !g.contains_node(a),
948            "removed node should not be in id_to_idx"
949        );
950        assert!(g.contains_node(b), "unrelated node should still exist");
951
952        // Edges involving the removed node should be gone
953        assert!(g.outgoing(a).is_empty());
954        assert!(g.incoming(b).is_empty());
955    }
956
957    #[test]
958    fn test_remove_node_then_readd() {
959        let mut g = CsrGraph::new();
960        let a = MemoryId::new();
961        let b = MemoryId::new();
962        let c = MemoryId::new();
963
964        g.add_edge(&make_edge(a, b, EdgeType::Caused));
965        g.remove_node(a);
966
967        // Re-adding the same ID should get a fresh index
968        g.add_edge(&make_edge(a, c, EdgeType::Related));
969        assert!(g.contains_node(a));
970        let out = g.outgoing(a);
971        assert_eq!(out.len(), 1);
972        assert_eq!(out[0].0, c);
973    }
974}