Optimize test-only sweep for touched nodes (#47)

chutoro-core/src/hnsw/cpu/test_helpers.rs

@@ -17,9 +17,15 @@ impl CpuHnsw {
     /// property-based mutation checks that rely on post-commit healing passes.
     pub fn heal_for_test(&self) {
         let healed = self.write_graph(|graph| {
-            let mut executor = graph.insertion_executor();
-            executor.heal_reachability(self.params.max_connections());
-            executor.enforce_bidirectional_all(self.params.max_connections());
+            {
+                let mut executor = graph.insertion_executor();
+                executor.heal_reachability(self.params.max_connections());
+            }
+            let touched = graph.take_touched_nodes();
+            {
+                let mut executor = graph.insertion_executor();
+                executor.enforce_bidirectional_for_touched(&touched, self.params.max_connections());
+            }
             Ok(())
         });
         if let Err(err) = healed {

chutoro-core/src/hnsw/cpu/unit_tests.rs

@@ -1,7 +1,12 @@
 //! Unit tests for the CPU HNSW index.
 
 use super::*;
-use crate::{MetricDescriptor, datasource::DataSource, error::DataSourceError, hnsw::HnswParams};
+use crate::{
+    MetricDescriptor,
+    datasource::DataSource,
+    error::DataSourceError,
+    hnsw::{HnswParams, graph::NodeContext, insert::test_helpers::add_edge_if_missing},
+};
 use std::{
     sync::{
         Arc,
@@ -48,6 +53,172 @@ fn insert_waits_for_mutex() {
     assert!(finished.load(AtomicOrdering::SeqCst));
 }
 
+#[test]
+fn heal_for_test_repairs_inserted_edges_without_sweeping_unrelated_edges() {
+    let params = HnswParams::new(2, 4).expect("params").with_rng_seed(41);
+    let index = CpuHnsw::with_capacity(params, 4).expect("index");
+    let source = TestSource::new(vec![0.0, 1.0, 2.0, 3.0]);
+
+    index.insert(0, &source).expect("insert entry");
+    index.heal_for_test();
+    index.insert(1, &source).expect("insert neighbour");
+
+    index
+        .write_graph(|graph| {
+            let touched = graph.take_touched_nodes();
+            assert!(
+                touched.contains(&(1, 0)),
+                "insertion must record the new node for localized healing"
+            );
+            graph.record_touched_nodes(touched);
+
+            graph
+                .attach_node(NodeContext {
+                    node: 2,
+                    level: 0,
+                    sequence: 2,
+                })
+                .expect("attach unrelated node");
+            graph
+                .attach_node(NodeContext {
+                    node: 3,
+                    level: 0,
+                    sequence: 3,
+                })
+                .expect("attach unrelated target");
+            graph
+                .node_mut(0)
+                .expect("entry must exist")
+                .neighbours_mut(0)
+                .retain(|neighbour| *neighbour != 1);
+            add_edge_if_missing(graph, 0, 2, 0);
+            add_edge_if_missing(graph, 2, 0, 0);
+            add_edge_if_missing(graph, 0, 3, 0);
+            add_edge_if_missing(graph, 3, 0, 0);
+            add_edge_if_missing(graph, 2, 3, 0);
+            Ok(())
+        })
+        .expect("prepare asymmetric edges");
+
+    index.heal_for_test();
+
+    index.inspect_graph(|graph| {
+        let node0 = graph.node(0).expect("entry must remain");
+        let node1 = graph.node(1).expect("inserted node must remain");
+        let node2 = graph.node(2).expect("unrelated node must remain");
+        let node3 = graph.node(3).expect("unrelated target must remain");
+        assert!(
+            node0.neighbours(0).contains(&1),
+            "healing must restore the reciprocal edge for the touched insertion"
+        );
+        assert!(node1.neighbours(0).contains(&0));
+        assert!(
+            node2.neighbours(0).contains(&3),
+            "the unrelated asymmetric edge is preserved for the locality check"
+        );
+        assert!(
+            !node3.neighbours(0).contains(&2),
+            "healing must not sweep an untouched asymmetric edge"
+        );
+    });
+    let touched = index
+        .write_graph(|graph| Ok(graph.take_touched_nodes()))
+        .expect("read healing queue");
+    assert!(touched.is_empty(), "healing must drain the insertion queue");
+}
+
+#[test]
+fn heal_for_test_drains_tracking_created_by_deletion() {
+    let params = HnswParams::new(2, 4).expect("params").with_rng_seed(43);
+    let mut index = CpuHnsw::with_capacity(params, 3).expect("index");
+    let source = TestSource::new(vec![0.0, 1.0, 2.0]);
+
+    for node in 0..source.len() {
+        index.insert(node, &source).expect("insert node");
+    }
+    index
+        .write_graph(|graph| Ok(graph.take_touched_nodes()))
+        .expect("clear insertion tracking");
+
+    assert!(
+        index.delete_node_for_test(1).expect("delete node"),
+        "existing node should be deleted"
+    );
+    let touched = index
+        .write_graph(|graph| Ok(graph.take_touched_nodes()))
+        .expect("read deletion tracking");
+    assert!(
+        !touched.is_empty(),
+        "deletion must record changed adjacency lists for healing"
+    );
+    index
+        .write_graph(|graph| {
+            graph.record_touched_nodes(touched);
+            Ok(())
+        })
+        .expect("restore deletion tracking");
+
+    index.heal_for_test();
+
+    assert!(index.inspect_graph(|graph| graph.node(1).is_none()));
+    let touched = index
+        .write_graph(|graph| Ok(graph.take_touched_nodes()))
+        .expect("read healing queue");
+    assert!(touched.is_empty(), "healing must drain the deletion queue");
+}
+
+#[test]
+fn heal_for_test_repairs_reachability_and_drains_its_tracking() {
+    let params = HnswParams::new(2, 4).expect("params").with_rng_seed(47);
+    let index = CpuHnsw::with_capacity(params, 3).expect("index");
+
+    index
+        .write_graph(|graph| {
+            graph.insert_first(NodeContext {
+                node: 0,
+                level: 0,
+                sequence: 0,
+            })?;
+            graph.attach_node(NodeContext {
+                node: 1,
+                level: 0,
+                sequence: 1,
+            })?;
+            graph.attach_node(NodeContext {
+                node: 2,
+                level: 0,
+                sequence: 2,
+            })?;
+            add_edge_if_missing(graph, 0, 1, 0);
+            add_edge_if_missing(graph, 1, 0, 0);
+            let _ = graph.take_touched_nodes();
+            Ok(())
+        })
+        .expect("prepare disconnected graph");
+
+    index.heal_for_test();
+
+    index.inspect_graph(|graph| {
+        let entry = graph.node(0).expect("entry must remain");
+        let repaired = graph.node(2).expect("isolated node must remain");
+        assert!(
+            entry.neighbours(0).contains(&2),
+            "reachability repair must link the isolated node from the entry"
+        );
+        assert!(
+            repaired.neighbours(0).contains(&0),
+            "reachability repair must create a reciprocal link"
+        );
+    });
+    let touched = index
+        .write_graph(|graph| Ok(graph.take_touched_nodes()))
+        .expect("read healing queue");
+    assert!(
+        touched.is_empty(),
+        "healing must consume tracking created by reachability repair"
+    );
+}
+
 #[derive(Clone)]
 struct TestSource {
     data: Vec<f32>,

chutoro-core/src/hnsw/graph/core.rs

@@ -9,6 +9,9 @@ use crate::hnsw::{
     types::{EntryPoint, InsertionPlan},
 };
 
+#[cfg(test)]
+use std::collections::BTreeSet;
+
 /// Context for attaching or inserting a node into the HNSW graph.
 ///
 /// The insertion `sequence` is used for deterministic neighbour ordering and
@@ -168,6 +171,8 @@ pub(crate) struct Graph {
     pub(super) params: HnswParams,
     pub(super) nodes: Vec<Option<Node>>,
     pub(super) entry: Option<EntryPoint>,
+    #[cfg(test)]
+    pub(super) touched: BTreeSet<(usize, usize)>,
 }
 
 fn should_promote_entry(current: Option<EntryPoint>, level: usize) -> bool {
@@ -183,6 +188,8 @@ impl Graph {
             params,
             nodes: vec![None; capacity],
             entry: None,
+            #[cfg(test)]
+            touched: BTreeSet::new(),
         }
     }
 

chutoro-core/src/hnsw/graph/test_helpers/mod.rs

@@ -19,6 +19,19 @@ use crate::hnsw::{
 mod tests;
 
 impl Graph {
+    /// Records test-only mutation pairs for the next localized healing pass.
+    pub(crate) fn record_touched_nodes<I>(&mut self, touched: I)
+    where
+        I: IntoIterator<Item = (usize, usize)>,
+    {
+        self.touched.extend(touched);
+    }
+
+    /// Returns and clears the mutation pairs accumulated since the last pass.
+    pub(crate) fn take_touched_nodes(&mut self) -> Vec<(usize, usize)> {
+        std::mem::take(&mut self.touched).into_iter().collect()
+    }
+
     pub(crate) fn set_params(&mut self, params: &HnswParams) {
         self.params = params.clone();
     }
@@ -31,13 +44,15 @@ impl Graph {
 
         let snapshot_nodes = self.nodes.clone();
         let snapshot_entry = self.entry;
+        let snapshot_touched = self.touched.clone();
         let removed_neighbours = collect_neighbour_layers(existing);
 
         let Some(_taken) = self.nodes.get_mut(node).and_then(Option::take) else {
             unreachable!("node presence checked above");
         };
 
-        self.strip_references_to(node);
+        let removed_references = self.strip_references_to(node);
+        self.record_touched_nodes(removed_references);
         self.reconnect_layers(removed_neighbours);
 
         if self.entry.map(|entry| entry.node) == Some(node) {
@@ -47,6 +62,7 @@ impl Graph {
         if let Err(err) = self.ensure_reachability() {
             self.nodes = snapshot_nodes;
             self.entry = snapshot_entry;
+            self.touched = snapshot_touched;
             return Err(err);
         }
 
@@ -127,37 +143,50 @@ impl Graph {
 
     pub(super) fn try_add_edge(&mut self, origin: usize, target: usize, level: usize) -> bool {
         let limit = params::connection_limit_for_level(level, self.params.max_connections());
-        let Some(node) = self.nodes.get_mut(origin).and_then(Option::as_mut) else {
-            return false;
-        };
-        if level >= node.level_count() {
-            return false;
-        }
+        let added = {
+            let Some(node) = self.nodes.get_mut(origin).and_then(Option::as_mut) else {
+                return false;
+            };
+            if level >= node.level_count() {
+                return false;
+            }
 
-        let neighbours = node.neighbours_mut(level);
-        if neighbours.contains(&target) {
-            return true;
-        }
+            let neighbours = node.neighbours_mut(level);
+            if neighbours.contains(&target) {
+                return true;
+            }
+
+            if neighbours.len() >= limit {
+                return false;
+            }
 
-        if neighbours.len() < limit {
             neighbours.push(target);
-            return true;
+            true
+        };
+        if added {
+            self.record_touched_nodes([(origin, level)]);
         }
-
-        false
+        added
     }
 
     pub(super) fn remove_edge(&mut self, origin: usize, target: usize, level: usize) {
-        let Some(node) = self.nodes.get_mut(origin).and_then(Option::as_mut) else {
-            return;
-        };
-        if level >= node.level_count() {
-            return;
-        }
+        let removed = {
+            let Some(node) = self.nodes.get_mut(origin).and_then(Option::as_mut) else {
+                return;
+            };
+            if level >= node.level_count() {
+                return;
+            }
 
-        let neighbours = node.neighbours_mut(level);
-        if let Some(pos) = neighbours.iter().position(|&candidate| candidate == target) {
+            let neighbours = node.neighbours_mut(level);
+            let Some(pos) = neighbours.iter().position(|&candidate| candidate == target) else {
+                return;
+            };
             neighbours.remove(pos);
+            true
+        };
+        if removed {
+            self.record_touched_nodes([(origin, level)]);
         }
     }
 
@@ -170,14 +199,30 @@ impl Graph {
         Ok(())
     }
 
-    fn strip_references_to(&mut self, node: usize) {
-        for maybe_node in self.nodes.iter_mut().flatten() {
-            let levels = maybe_node.level_count();
-            for level in 0..levels {
-                let neighbours = maybe_node.neighbours_mut(level);
-                neighbours.retain(|&target| target != node);
-            }
+    fn strip_references_to(&mut self, node: usize) -> Vec<(usize, usize)> {
+        let mut touched = Vec::new();
+        for (id, maybe_node) in self.nodes.iter_mut().enumerate() {
+            let Some(existing) = maybe_node.as_mut() else {
+                continue;
+            };
+            touched.extend(Self::strip_node_references(id, existing, node));
         }
+        touched
+    }
+
+    fn strip_node_references(
+        id: usize,
+        existing: &mut crate::hnsw::node::Node,
+        target: usize,
+    ) -> Vec<(usize, usize)> {
+        (0..existing.level_count())
+            .filter_map(|level| {
+                let neighbours = existing.neighbours_mut(level);
+                let previous_len = neighbours.len();
+                neighbours.retain(|&candidate| candidate != target);
+                (neighbours.len() != previous_len).then_some((id, level))
+            })
+            .collect()
     }
 
     fn reconnect_layers(&mut self, removed_neighbours: Vec<Vec<usize>>) {