mnemosyne-graph-core 0.1.0

//! R-tree viewport spatial index.
//!
//! Wraps `rstar::RTree` with an incremental-update policy so the
//! caller can stream position changes from the force simulation
//! into the tree without paying full O(n log n) bulk-rebuild cost
//! on every tick. Queries return ids filtered to a viewport
//! axis-aligned bounding box, optionally truncated and ordered by
//! score.
//!
//! ## Update policy
//!
//! `update_positions` appends changes to a dirty buffer and
//! updates a per-id position cache. Once the dirty-buffer size
//! exceeds `max(1000, n / 20)` the tree is rebuilt in bulk.
//! Rebuild is O(n log n); per-update cost is amortised O(1)
//! across a burst of updates. The threshold was picked to keep
//! rebuild-induced pauses rare even on 100k-node graphs
//! (5000 updates before rebuild = ~5% of node count) while still
//! keeping queries close to up-to-date: the default
//! force-simulation tick cadence moves ≤ 100 nodes per frame,
//! so a 5000-update buffer holds ≈ 50 frames of drift.

use glam::Vec2;
use rstar::primitives::GeomWithData;
use rstar::{RTree, AABB};
#[cfg(feature = "serde_support")]
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// A single indexable point. `position` is an `[f32; 2]` so it
/// matches rstar's `Point` impl for arrays directly; we convert
/// to/from `glam::Vec2` at the API boundary only.
#[cfg_attr(feature = "serde_support", derive(Serialize, Deserialize))]
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct IndexPoint {
    pub id: u32,
    pub position: [f32; 2],
    pub score: f32,
}

/// How to order the ids returned by [`ViewportIndex::query`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScoreKey {
    /// Tree-traversal order. Cheapest; caller shouldn't rely on
    /// any particular ordering.
    Natural,
    /// Highest score first.
    Desc,
    /// Lowest score first.
    Asc,
}

// `GeomWithData<[f32; 2], (u32, f32)>`: payload is (id, score).
type Entry = GeomWithData<[f32; 2], (u32, f32)>;

/// R-tree spatial index with a lazy-rebuild incremental-update
/// policy. Construct with [`ViewportIndex::new`] and populate with
/// [`ViewportIndex::rebuild`].
pub struct ViewportIndex {
    tree: RTree<Entry>,
    // Current position and score for every id known to the index.
    // Used to (a) translate incremental updates into the final
    // `IndexPoint` set that a rebuild needs, and (b) let callers
    // update the same id multiple times before a rebuild hits.
    state: HashMap<u32, IndexPoint>,
    // Pending incremental position updates. Position-only — score
    // doesn't drift between ticks so it comes in only via `rebuild`.
    dirty_count: usize,
}

impl ViewportIndex {
    pub fn new() -> Self {
        Self {
            tree: RTree::new(),
            state: HashMap::new(),
            dirty_count: 0,
        }
    }

    /// Number of points currently indexed (or pending in the dirty
    /// buffer — the state map is the source of truth).
    #[inline]
    pub fn len(&self) -> usize {
        self.state.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.state.is_empty()
    }

    /// Bulk-load the index from `points`. O(n log n) via rstar's
    /// STR packing. Clears any pending updates.
    pub fn rebuild(&mut self, points: &[IndexPoint]) {
        self.state.clear();
        self.state.reserve(points.len());
        let entries: Vec<Entry> = points
            .iter()
            .map(|p| {
                self.state.insert(p.id, *p);
                GeomWithData::new(p.position, (p.id, p.score))
            })
            .collect();
        self.tree = RTree::bulk_load(entries);
        self.dirty_count = 0;
    }

    /// Apply a batch of `(id, new_position)` updates. Updates are
    /// stored in an internal position cache immediately; the
    /// R-tree itself is rebuilt only when the dirty count crosses
    /// `max(1000, n / 20)`. Amortised O(1) per update.
    pub fn update_positions(&mut self, changes: &[(u32, Vec2)]) {
        if changes.is_empty() {
            return;
        }
        for &(id, pos) in changes {
            let entry = self.state.entry(id).or_insert(IndexPoint {
                id,
                position: [pos.x, pos.y],
                score: 0.0,
            });
            entry.position = [pos.x, pos.y];
        }
        self.dirty_count += changes.len();

        let threshold = 1000usize.max(self.state.len() / 20);
        if self.dirty_count > threshold {
            // Rebuild from the state map. Materialise once; rstar's
            // bulk_load consumes the vec.
            let entries: Vec<Entry> = self
                .state
                .values()
                .map(|p| GeomWithData::new(p.position, (p.id, p.score)))
                .collect();
            self.tree = RTree::bulk_load(entries);
            self.dirty_count = 0;
        }
    }

    /// Return the ids whose current position falls inside the
    /// axis-aligned box `[min, max]`, optionally truncated and
    /// ordered.
    ///
    /// The tree may be slightly out-of-date relative to the state
    /// map between rebuilds — callers who need point-in-time
    /// accuracy should call `rebuild` before querying. For the
    /// viewport use case a handful of frames of drift is fine
    /// (the user's viewport is typically moving faster than the
    /// simulation is drifting).
    pub fn query(&self, min: Vec2, max: Vec2, limit: usize, order: ScoreKey) -> Vec<u32> {
        if self.tree.size() == 0 {
            return Vec::new();
        }

        let aabb = AABB::from_corners([min.x, min.y], [max.x, max.y]);
        // Collect (id, score) inside the envelope. `score` only
        // matters when `order != Natural`.
        let mut hits: Vec<(u32, f32)> = self
            .tree
            .locate_in_envelope_intersecting(&aabb)
            .map(|e| e.data)
            .collect();

        match order {
            ScoreKey::Natural => {}
            ScoreKey::Desc => {
                hits.sort_by(|a, b| {
                    b.1.partial_cmp(&a.1)
                        .unwrap_or(std::cmp::Ordering::Equal)
                        .then_with(|| a.0.cmp(&b.0))
                });
            }
            ScoreKey::Asc => {
                hits.sort_by(|a, b| {
                    a.1.partial_cmp(&b.1)
                        .unwrap_or(std::cmp::Ordering::Equal)
                        .then_with(|| a.0.cmp(&b.0))
                });
            }
        }

        hits.truncate(limit);
        hits.into_iter().map(|(id, _)| id).collect()
    }

    /// Return the `(id, score)` pairs inside the axis-aligned box,
    /// optionally truncated and ordered.
    ///
    /// Added in Phase 223 Wave 3. Same semantics as [`Self::query`]
    /// but exposes the associated score so the PyO3 surface can
    /// return both without a second lookup. The id-only [`Self::query`]
    /// stays as a convenience for callers (including the WASM crate)
    /// that already ignore the score.
    pub fn query_with_scores(
        &self,
        min: Vec2,
        max: Vec2,
        limit: usize,
        order: ScoreKey,
    ) -> Vec<(u32, f32)> {
        if self.tree.size() == 0 {
            return Vec::new();
        }

        let aabb = AABB::from_corners([min.x, min.y], [max.x, max.y]);
        let mut hits: Vec<(u32, f32)> = self
            .tree
            .locate_in_envelope_intersecting(&aabb)
            .map(|e| e.data)
            .collect();

        match order {
            ScoreKey::Natural => {}
            ScoreKey::Desc => {
                hits.sort_by(|a, b| {
                    b.1.partial_cmp(&a.1)
                        .unwrap_or(std::cmp::Ordering::Equal)
                        .then_with(|| a.0.cmp(&b.0))
                });
            }
            ScoreKey::Asc => {
                hits.sort_by(|a, b| {
                    a.1.partial_cmp(&b.1)
                        .unwrap_or(std::cmp::Ordering::Equal)
                        .then_with(|| a.0.cmp(&b.0))
                });
            }
        }

        hits.truncate(limit);
        hits
    }

    /// Snapshot the indexed points — a clone of the internal state
    /// map's values, in unspecified order.
    ///
    /// Phase 223 Wave 3: used by the PyO3 surface to hand the point
    /// set to bincode for Redis snapshotting. The tree itself isn't
    /// serialised; callers rebuild it in O(n log n) from the returned
    /// slice via [`Self::rebuild`].
    pub fn snapshot(&self) -> Vec<IndexPoint> {
        self.state.values().copied().collect()
    }
}

impl Default for ViewportIndex {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashSet;

    fn grid_points() -> Vec<IndexPoint> {
        // 5x5 grid of points at integer coordinates 0..5, id = x*10 + y,
        // score = x + y.
        let mut out = Vec::new();
        for x in 0..5 {
            for y in 0..5 {
                out.push(IndexPoint {
                    id: (x * 10 + y) as u32,
                    position: [x as f32, y as f32],
                    score: (x + y) as f32,
                });
            }
        }
        out
    }

    #[test]
    fn known_rect_returns_expected_set() {
        let mut idx = ViewportIndex::new();
        idx.rebuild(&grid_points());
        // Top-left 3x3 = x in 0..=2, y in 0..=2.
        let got: HashSet<u32> = idx
            .query(Vec2::new(0.0, 0.0), Vec2::new(2.0, 2.0), 100, ScoreKey::Natural)
            .into_iter()
            .collect();
        let want: HashSet<u32> = (0..3u32)
            .flat_map(|x| (0..3u32).map(move |y| x * 10 + y))
            .collect();
        assert_eq!(got, want);
    }

    #[test]
    fn empty_rect_returns_empty() {
        let mut idx = ViewportIndex::new();
        idx.rebuild(&grid_points());
        let got = idx.query(
            Vec2::new(100.0, 100.0),
            Vec2::new(200.0, 200.0),
            100,
            ScoreKey::Natural,
        );
        assert!(got.is_empty());
    }

    #[test]
    fn limit_and_order_respected() {
        let mut idx = ViewportIndex::new();
        let points: Vec<IndexPoint> = (0..20u32)
            .map(|i| IndexPoint {
                id: i,
                position: [i as f32 % 10.0, (i / 10) as f32],
                score: (i as f32) * 0.5,
            })
            .collect();
        idx.rebuild(&points);

        // Query a rect that covers everything.
        let got = idx.query(
            Vec2::new(-1.0, -1.0),
            Vec2::new(10.0, 10.0),
            5,
            ScoreKey::Desc,
        );
        // Top 5 by score (descending) should be ids 19, 18, 17, 16, 15.
        assert_eq!(got, vec![19, 18, 17, 16, 15]);
    }

    #[test]
    fn incremental_update_matches_rebuild() {
        let mut base_points: Vec<IndexPoint> = (0..100u32)
            .map(|i| IndexPoint {
                id: i,
                position: [(i as f32) * 0.5, (i as f32) * 0.3],
                score: i as f32,
            })
            .collect();

        let mut incremental = ViewportIndex::new();
        incremental.rebuild(&base_points);

        // Apply 30 position changes. Threshold is max(1000, 100/20=5) = 1000,
        // so these stay in the dirty buffer without triggering a rebuild —
        // exactly the "below threshold" branch we want to exercise. To
        // compare query results we then force a rebuild on the incremental
        // index by pushing past the threshold with a no-op re-application
        // of the same updates (position already matches, so no new state).
        let mut updates = Vec::new();
        for i in 0..30u32 {
            let new_pos = Vec2::new(50.0 + i as f32, 50.0 + i as f32);
            updates.push((i, new_pos));
            base_points[i as usize].position = [new_pos.x, new_pos.y];
        }
        incremental.update_positions(&updates);

        let mut rebuilt = ViewportIndex::new();
        rebuilt.rebuild(&base_points);

        // Cross the dirty threshold (1000) with repeated no-op updates so
        // the tree picks up the 30 changes. The state is unchanged; only
        // the tree gets refreshed.
        let mut more_updates = Vec::new();
        for _ in 0..35 {
            more_updates.extend_from_slice(&updates);
        }
        incremental.update_positions(&more_updates);

        // Query a rect covering all final positions and compare id sets.
        let bbox_min = Vec2::new(-10.0, -10.0);
        let bbox_max = Vec2::new(100.0, 100.0);
        let a: HashSet<u32> = incremental
            .query(bbox_min, bbox_max, 1000, ScoreKey::Natural)
            .into_iter()
            .collect();
        let b: HashSet<u32> = rebuilt
            .query(bbox_min, bbox_max, 1000, ScoreKey::Natural)
            .into_iter()
            .collect();
        assert_eq!(a, b);
    }

    #[test]
    fn incremental_update_triggers_rebuild() {
        let points: Vec<IndexPoint> = (0..100u32)
            .map(|i| IndexPoint {
                id: i,
                position: [i as f32, i as f32],
                score: i as f32,
            })
            .collect();
        let mut idx = ViewportIndex::new();
        idx.rebuild(&points);

        // Apply 5001 updates. Threshold is max(1000, 100/20=5) = 1000,
        // so we must cross it. Sweep all 100 ids, updating each
        // multiple times to ≈ position (1000+k, 1000+k).
        let mut updates = Vec::with_capacity(5001);
        for step in 0..51u32 {
            for id in 0..100u32 {
                updates.push((id, Vec2::new(1000.0 + step as f32, 1000.0 + step as f32)));
                if updates.len() == 5001 {
                    break;
                }
            }
            if updates.len() == 5001 {
                break;
            }
        }
        assert_eq!(updates.len(), 5001);
        idx.update_positions(&updates);

        // Final positions all near (1050, 1050). Query a rect around
        // that should return (nearly) all ids. The last update for
        // each id wins.
        let hits = idx.query(
            Vec2::new(900.0, 900.0),
            Vec2::new(1100.0, 1100.0),
            200,
            ScoreKey::Asc,
        );
        assert_eq!(hits.len(), 100, "expected all 100 ids inside viewport");
        // Ascending by score: ids 0..100 have scores 0..100.
        assert_eq!(hits[0], 0);
        assert_eq!(hits[99], 99);

        // Old positions should no longer match.
        let old_hits = idx.query(
            Vec2::new(-1.0, -1.0),
            Vec2::new(100.0, 100.0),
            200,
            ScoreKey::Natural,
        );
        assert!(
            old_hits.is_empty(),
            "stale positions leaked: {old_hits:?}"
        );
    }
}