batpak 0.9.0

pub(crate) mod columnar;
mod entry;
pub(crate) mod idemp;
pub(crate) mod interner;
mod projection_bridge;
mod query;
mod restore;
mod visibility;

use self::columnar::ScanIndex;
pub(crate) use self::entry::{ClockKey, QueryHit};
pub use self::entry::{DiskPos, IndexEntry};
use self::interner::{InternId, StringInterner};
pub(crate) use self::projection_bridge::{
    projection_kind_matches, ProjectionCacheStoreStatus, ProjectionReplayItem, ProjectionReplayPlan,
};
use self::restore::RestoreBase;
pub(crate) use self::restore::{
    recommended_restore_chunk_count, restore_chunk_ranges, RoutingSummary,
};
use self::visibility::SequenceGate;
use crate::store::config::IndexConfig;
use crate::store::hidden_ranges::CancelledVisibilityRanges;
use dashmap::DashMap;
use parking_lot::RwLock;
use std::collections::{BTreeMap, HashMap};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
struct LaneHeadKey {
    entity_id: InternId,
    lane: u32,
}

impl LaneHeadKey {
    fn new(entity_id: InternId, lane: u32) -> Self {
        Self { entity_id, lane }
    }
}

fn lane_visible_from_entries(entries: &[Arc<IndexEntry>]) -> BTreeMap<u32, u64> {
    let mut lanes = BTreeMap::new();
    for entry in entries {
        let visible = entry.global_sequence.saturating_add(1);
        lanes
            .entry(entry.dag_lane)
            .and_modify(|current: &mut u64| *current = (*current).max(visible))
            .or_insert(visible);
    }
    lanes
}

/// StoreIndex: in-memory 2D index + auxiliaries. Not persisted; rebuilt from segments on cold start.
/// [DEP:dashmap::DashMap] — see DEPENDENCY SURFACE for deadlock warnings
///
/// # F6 / FREEZE-4 compact swap-point
///
/// Compaction must never expose a cleared or partially rebuilt index to
/// readers. The protocol enforced here is:
///
/// 1. **Off-side build.** `compact()` allocates a *fresh* [`StoreIndex`],
///    populates it from segments via `rebuild_from_segments`, and
///    only then hands it to [`StoreIndex::replace_contents_from_fresh`].
///    While the fresh index is being built the live index is not touched,
///    so readers keep serving the pre-compact state unchanged.
/// 2. **Single swap point.** The critical section inside
///    `replace_contents_from_fresh` takes the write guard on `swap_gate`
///    and publishes the new contents. That guard is the commit point.
/// 3. **Failure safety.** If the off-side rebuild errors before the swap,
///    the live index is still valid and readable; the caller observes
///    [`crate::store::segment::CompactionOutcome::Failed`] rather than a
///    silently clobbered store.
/// 4. **Reader observation.** Reader-facing methods take the read guard on
///    `swap_gate` at entry. Concurrent readers complete on their snapshot;
///    a concurrent compact-swap waits for them to release before taking
///    the write guard. Either the old index or the new one is fully
///    observable — never a partially rebuilt view.
/// 5. **Segment cleanup AFTER swap.** Callers (see
///    `src/store/lifecycle.rs::compact`) delete the old segment files
///    only after `replace_contents_from_fresh` returns. If the process
///    crashes between swap and cleanup, cold-start reconciliation removes
///    the orphaned files.
pub(crate) struct StoreIndex {
    /// Primary: entity -> ordered events. [DEP:dashmap::DashMap::get_mut] for insert.
    streams: DashMap<Arc<str>, BTreeMap<ClockKey, Arc<IndexEntry>>>,
    /// Base AoS scan maps plus optional overlay views.
    /// Handles by_fact and scope queries while keeping the live topology honest:
    /// the base maps always exist and configured overlays fan out in parallel.
    pub(crate) scan: ScanIndex,
    /// Point lookup: event_id -> entry. O(1) get by ID.
    by_id: DashMap<u128, Arc<IndexEntry>>,
    /// Chain head: (entity, lane) -> latest IndexEntry. For prev_hash in writer step 2.
    latest: DashMap<LaneHeadKey, Arc<IndexEntry>>,
    /// Gated sequence counter: allocator + visibility watermark.
    /// Replaces the former bare `global_sequence: AtomicU64`.
    pub(crate) sequence: SequenceGate,
    /// Total event count.
    len: AtomicUsize,
    /// String interner for compact index keys and checkpoint serialization.
    /// Entity and scope strings are interned on insert; IDs are used by
    /// checkpoint and (future) InternId-based IndexEntry fields.
    pub(crate) interner: Arc<StringInterner>,
    /// F6 / FREEZE-4 compact swap-point lock. Writers (compact) take the
    /// write guard for the single critical section that swaps fresh
    /// contents in; readers (queries) take the read guard at entry so they
    /// see either the old index or the new one, never a partial rebuild.
    /// The lock's inner unit value is intentionally trivial — it exists
    /// purely as a rendezvous barrier around the swap.
    swap_gate: RwLock<()>,
    /// Durable idempotency key → receipt-reconstruction store. A separate
    /// authority from `by_id`: it is NOT cleared by
    /// `replace_contents_from_fresh`, so a keyed retry survives retention
    /// compaction, cold-start, and snapshot independent of event eviction.
    /// justifies: INV-IDEMPOTENCY-DURABLE-WINDOW
    pub(crate) idemp: idemp::IdempotencyStore,
}

impl StoreIndex {
    #[cfg(test)]
    pub(crate) fn new() -> Self {
        Self::with_config(&IndexConfig::default())
    }

    /// Create a StoreIndex with the specified index configuration.
    pub(crate) fn with_config(config: &IndexConfig) -> Self {
        Self {
            streams: DashMap::new(),
            scan: ScanIndex::for_config(config),
            by_id: DashMap::new(),
            latest: DashMap::new(),
            sequence: SequenceGate::new(),
            len: AtomicUsize::new(0),
            interner: Arc::new(StringInterner::new()),
            swap_gate: RwLock::new(()),
            idemp: idemp::IdempotencyStore::new(
                config.idempotency_retention,
                config.idempotency_overflow,
            ),
        }
    }

    /// Reserve N global sequences for batch staging.
    /// Returns the first sequence number; caller allocates `[first, first + n)`.
    /// Used by writer to pre-assign sequences before writing to disk.
    pub(crate) fn reserve_sequences(&self, n: u64) -> u64 {
        self.sequence.reserve(n)
    }

    /// Called by writer step 9. Inserts into ALL indexes atomically.
    /// Caller must be the single writer thread; this is the only writer of
    /// the index, so no per-entity lock is needed.
    /// Advances the allocator by one — used by the live single-event append path.
    ///
    /// F6 / FREEZE-4: the insert path takes the `swap_gate` read guard so
    /// that a concurrent compact-swap (which holds the write guard) cannot
    /// race mid-clear with an in-flight writer insert. The guard is
    /// released before the method returns — writer throughput is unaffected
    /// when no compact is in flight.
    pub(crate) fn insert(&self, entry: IndexEntry) {
        let _read = self.swap_gate.read();
        self.insert_inner(entry);
        // Advance allocator (visibility is advanced separately by publish()).
        self.sequence.advance();
    }

    fn insert_inner(&self, entry: IndexEntry) {
        let lane = entry.dag_lane;

        // Intern entity and scope strings. IDs stored in IndexEntry for
        // compact checkpoint serialization and future InternId-only index.
        debug_assert!(self
            .interner
            .intern(entry.coord.entity())
            .is_ok_and(|id| id == entry.entity_id));
        debug_assert!(self
            .interner
            .intern(entry.coord.scope())
            .is_ok_and(|id| id == entry.scope_id));

        let key = ClockKey {
            wall_ms: entry.wall_ms,
            clock: entry.clock,
            uuid: entry.event_id,
        };

        // Arc: one allocation, shared across all maps.
        let arc_entry = Arc::new(entry);
        let entity = arc_entry.coord.entity_arc();
        let entity_id = arc_entry.entity_id;

        // Primary index: entity -> BTreeMap
        // [DEP:dashmap::DashMap::entry] — holds write lock, release fast
        self.streams
            .entry(Arc::clone(&entity))
            .or_default()
            .insert(key, Arc::clone(&arc_entry));

        // Scan index: by_fact + scope (DashMap or columnar depending on layout)
        self.scan.insert(&arc_entry);

        // Point lookup
        self.by_id
            .insert(arc_entry.event_id, Arc::clone(&arc_entry));

        // Chain head
        self.latest
            .insert(LaneHeadKey::new(entity_id, lane), arc_entry);

        self.len.fetch_add(1, Ordering::Relaxed);
    }

    /// Atomic batch insert: all entries become visible together.
    ///
    /// F6 / FREEZE-4: takes the `swap_gate` read guard for the duration of
    /// the batch so that a concurrent compact-swap cannot race mid-clear
    /// with an in-flight writer batch insert.
    pub(crate) fn insert_batch(&self, entries: Vec<IndexEntry>) {
        let _read = self.swap_gate.read();
        if entries.is_empty() {
            return;
        }

        // Pre-allocate Arcs to minimize work under locks.
        let arc_entries: Vec<Arc<IndexEntry>> = entries.into_iter().map(Arc::new).collect();

        // Insert all entries. Since we have a single writer thread,
        // no other inserts can interleave. Readers will see all or none
        // depending on when they query relative to this loop.
        for arc_entry in &arc_entries {
            let entity = arc_entry.coord.entity_arc();
            let lane = arc_entry.dag_lane;
            let entity_id = arc_entry.entity_id;
            let key = ClockKey {
                wall_ms: arc_entry.wall_ms,
                clock: arc_entry.clock,
                uuid: arc_entry.event_id,
            };

            // Primary index: entity -> BTreeMap
            self.streams
                .entry(Arc::clone(&entity))
                .or_default()
                .insert(key, Arc::clone(arc_entry));

            // Scan index
            self.scan.insert(arc_entry);

            // Point lookup
            self.by_id.insert(arc_entry.event_id, Arc::clone(arc_entry));

            // Chain head
            self.latest
                .insert(LaneHeadKey::new(entity_id, lane), Arc::clone(arc_entry));

            // Global sequence already reserved during batch staging via reserve_sequences()
            self.len.fetch_add(1, Ordering::Relaxed);
        }

        // Global sequence already reserved during batch staging via reserve_sequences()
        // No additional fetch_add needed.
    }

    /// Replace the in-memory index contents from a sorted durable snapshot.
    ///
    /// `entries` must be sorted ascending by `global_sequence`. The allocator is
    /// restored to `max(last_sequence + 1, allocator_hint)` and published only
    /// after every base map and overlay view has been rebuilt.
    fn restore_sorted_entries_impl(
        &self,
        entries: Vec<IndexEntry>,
        allocator_hint: u64,
        chunk_count: usize,
        routing_hint: Option<&RoutingSummary>,
        before_publish: impl FnOnce(&Self),
    ) -> Result<(), crate::store::StoreError> {
        self.streams.clear();
        self.scan.clear();
        self.by_id.clear();
        self.latest.clear();
        self.sequence.clear();

        let restored = RestoreBase::from_sorted_entries(entries, chunk_count, routing_hint);
        let mut by_id =
            HashMap::<u128, Arc<IndexEntry>>::with_capacity(restored.entries_by_sequence.len());
        let mut latest = HashMap::<LaneHeadKey, Arc<IndexEntry>>::new();

        for run in &restored.routing.entity_runs {
            let range = run.usize_range()?;
            let slice = restored.entries_by_entity.get(range).ok_or_else(|| {
                crate::store::StoreError::corrupt_segment_with_detail(
                    0,
                    "routing entity-run range out of bounds for restored entries",
                )
            })?;
            if slice.is_empty() {
                continue;
            }
            let entity = slice[0].coord.entity_arc();
            let stream: BTreeMap<ClockKey, Arc<IndexEntry>> = slice
                .iter()
                .map(|entry| {
                    (
                        ClockKey {
                            wall_ms: entry.wall_ms,
                            clock: entry.clock,
                            uuid: entry.event_id,
                        },
                        Arc::clone(entry),
                    )
                })
                .collect();
            for entry in slice {
                latest.insert(
                    LaneHeadKey::new(entry.entity_id, entry.dag_lane),
                    Arc::clone(entry),
                );
            }
            self.streams.insert(entity, stream);
        }

        self.scan.rebuild_from_restore_base(
            &restored.entries_by_sequence,
            &restored.entries_by_entity,
            &restored.routing,
        )?;
        for entry in &restored.entries_by_sequence {
            by_id.insert(entry.event_id, Arc::clone(entry));
        }
        for (event_id, entry) in by_id {
            self.by_id.insert(event_id, entry);
        }
        for (key, entry) in latest {
            self.latest.insert(key, entry);
        }

        self.len
            .store(restored.entries_by_sequence.len(), Ordering::Relaxed);
        before_publish(self);

        let lane_visible = lane_visible_from_entries(&restored.entries_by_sequence);
        let next_sequence = restored
            .entries_by_sequence
            .last()
            .map(|entry| entry.global_sequence.saturating_add(1))
            .unwrap_or(allocator_hint)
            .max(allocator_hint);
        self.sequence.restore_allocator(next_sequence);
        self.publish(next_sequence, "restore_sorted_entries")?;
        self.sequence.restore_lane_visible(lane_visible);
        Ok(())
    }

    #[cfg(test)]
    pub(crate) fn restore_sorted_entries(
        &self,
        entries: Vec<IndexEntry>,
        allocator_hint: u64,
    ) -> Result<(), crate::store::StoreError> {
        self.restore_sorted_entries_impl(entries, allocator_hint, 1, None, |_| {})
    }

    pub(crate) fn restore_sorted_entries_with_routing(
        &self,
        entries: Vec<IndexEntry>,
        allocator_hint: u64,
        routing: &RoutingSummary,
    ) -> Result<(), crate::store::StoreError> {
        let chunk_count = usize::try_from(routing.chunk_count).unwrap_or(1).max(1);
        self.restore_sorted_entries_impl(
            entries,
            allocator_hint,
            chunk_count,
            Some(routing),
            |_| {},
        )
    }

    #[cfg(test)]
    pub(crate) fn restore_sorted_entries_with_before_publish(
        &self,
        entries: Vec<IndexEntry>,
        allocator_hint: u64,
        before_publish: impl FnOnce(&Self),
    ) -> Result<(), crate::store::StoreError> {
        self.restore_sorted_entries_impl(entries, allocator_hint, 1, None, before_publish)
    }

    /// Return a snapshot of all entries in the index, collected from `by_id`.
    ///
    /// Used by `checkpoint::write_checkpoint` to serialise the full index.
    /// DashMap iteration is not a linearisable snapshot, but that is acceptable
    /// because checkpoints are always written from a quiesced write path.
    pub(crate) fn all_entries(&self) -> Vec<IndexEntry> {
        self.by_id
            .iter()
            .map(|r| r.value().as_ref().clone())
            .collect()
    }

    pub(crate) fn visible_entries(&self) -> Vec<IndexEntry> {
        let _read_guard = self.swap_gate.read();
        let visibility = self.sequence.snapshot();
        self.by_id
            .iter()
            .filter(|r| {
                let entry = r.value();
                visibility.is_visible_on_lane(entry.global_sequence, entry.dag_lane)
            })
            .map(|r| r.value().as_ref().clone())
            .collect()
    }

    pub(crate) fn hlc_for_global_sequence(
        &self,
        global_sequence: u64,
    ) -> Option<crate::store::stats::HlcPoint> {
        self.by_id
            .iter()
            .find(|entry| entry.value().global_sequence == global_sequence)
            .map(|entry| crate::store::stats::HlcPoint {
                wall_ms: entry.value().wall_ms,
                global_sequence,
            })
    }

    /// Current allocator position (next sequence to be assigned).
    /// Used by checkpoint, rebuild, writer, and stats/diagnostics.
    pub(crate) fn global_sequence(&self) -> u64 {
        self.sequence.allocated()
    }

    /// Current visibility watermark (exclusive upper bound).
    /// Entries with `global_sequence < visible_sequence()` are returned by read methods.
    pub(crate) fn visible_sequence(&self) -> u64 {
        self.sequence.visible()
    }

    /// Advance the visibility watermark so readers can observe entries
    /// with `global_sequence < up_to`.
    ///
    /// Called by the writer after `insert()` or `insert_batch()`, and by
    /// checkpoint restore / index rebuild after all entries are loaded.
    pub(crate) fn publish(
        &self,
        up_to: u64,
        operation: &'static str,
    ) -> Result<(), crate::store::StoreError> {
        self.sequence.publish(up_to, operation)
    }

    pub(crate) fn publish_on_lanes(
        &self,
        global_up_to: u64,
        lanes: impl IntoIterator<Item = (u32, u64)>,
        operation: &'static str,
    ) -> Result<(), crate::store::StoreError> {
        self.sequence
            .publish_on_lanes(global_up_to, lanes, operation)
    }

    pub(crate) fn len(&self) -> usize {
        self.len.load(Ordering::Relaxed)
    }

    /// Mark durable idempotency entries whose event frame is no longer present
    /// in the live `by_id` index as evicted. Called at the compaction tail so a
    /// later no-op deduplicated against an evicted event is honestly flagged.
    /// justifies: INV-IDEMPOTENCY-DURABLE-WINDOW
    pub(crate) fn mark_idemp_evicted_against_live(&self) {
        let _read = self.swap_gate.read();
        self.idemp
            .mark_evicted(|event_id| self.by_id.contains_key(&event_id));
    }

    /// F6 / FREEZE-4 compact swap-point.
    ///
    /// Adopt the contents of a freshly-built sibling [`StoreIndex`] as the
    /// live contents of `self`. Called exclusively by
    /// `src/store/lifecycle.rs::compact` after the segment merge has
    /// produced the new on-disk layout and `rebuild_from_segments` has
    /// populated the fresh index off-side.
    ///
    /// Protocol:
    ///
    /// * The **exclusive** write guard on `swap_gate` is held for the
    ///   duration of the transfer. Reader-facing methods on [`StoreIndex`]
    ///   acquire the read guard at entry, so no reader observes the
    ///   intermediate cleared state.
    /// * The interner on `self` is mutated in-place via
    ///   [`StringInterner::replace_from_full_snapshot`] so that external
    ///   references to `index.interner` (cold-start, writer staging) remain
    ///   valid across the swap.
    /// * `fresh` is consumed; its per-field contents are transferred into
    ///   `self` by draining `fresh` and re-inserting. This is an `O(n)`
    ///   operation but it is performed on already-populated in-memory data
    ///   (no disk I/O), which is significantly faster than the previous
    ///   protocol of clearing the live index and replaying segments under
    ///   reader visibility.
    pub(crate) fn replace_contents_from_fresh(
        &self,
        fresh: StoreIndex,
    ) -> Result<(), crate::store::StoreError> {
        let _write = self.swap_gate.write();

        // Reset live fields. `scan.clear()` also drops any overlay slots.
        self.streams.clear();
        self.scan.clear();
        self.by_id.clear();
        self.latest.clear();
        self.sequence.clear();
        self.len.store(0, Ordering::Relaxed);

        // Transfer the fresh interner strings in-place so that external
        // handles to `self.interner` (cold-start checkpoint writer, writer
        // staging) continue to see a populated interner after the swap.
        // `replace_from_full_snapshot` expects `[sentinel, ...strings]`;
        // `to_snapshot` returns only the non-sentinel strings, so we
        // re-prepend the sentinel here (same shape cold-start uses).
        let mut interner_full = vec![String::new()];
        interner_full.extend(fresh.interner.to_snapshot());
        self.interner.replace_from_full_snapshot(&interner_full)?;

        // Streams: drain fresh, insert into self.
        for (entity, stream) in fresh.streams.into_iter() {
            self.streams.insert(entity, stream);
        }
        for (id, entry) in fresh.by_id.into_iter() {
            self.by_id.insert(id, entry);
        }
        for (key, latest_entry) in fresh.latest.into_iter() {
            self.latest.insert(key, latest_entry);
        }

        // Scan overlays: rebuild from the entries we just populated so the
        // overlay topology matches the live configuration. We walk `by_id`
        // once because it is the only map that already contains every
        // entry exactly once.
        for entry in self.by_id.iter() {
            self.scan.insert(entry.value());
        }

        // Sequence gate: restore allocator + visibility + cancelled ranges
        // from the fresh gate. The fresh gate was driven by
        // `rebuild_from_segments` which publishes to the correct watermark.
        let fresh_allocated = fresh.sequence.allocated();
        let fresh_visible = fresh.sequence.visible();
        let fresh_cancelled = fresh.sequence.cancelled_ranges_snapshot();
        let fresh_lane_visible = fresh.sequence.lane_visible_snapshot();
        self.sequence.restore_allocator(fresh_allocated);
        self.sequence
            .publish(fresh_visible, "replace_contents_from_fresh")?;
        self.sequence.restore_lane_visible(fresh_lane_visible);
        self.sequence.restore_cancelled_ranges(fresh_cancelled);

        // Restore len. `fresh.len` is a snapshot of how many entries were
        // populated off-side — reading it after draining `fresh.by_id`
        // would be inconsistent, so we use `self.by_id.len()` which is the
        // live truth at this point.
        self.len.store(self.by_id.len(), Ordering::Relaxed);
        Ok(())
    }

    /// Begin a replay session against this index. Use this for checkpoint
    /// restore and segment rebuild paths to preserve sparse `global_sequence`
    /// values from durable sources (SIDX footers / checkpoint payload) while
    /// synthesizing sequences for entries with no durable source.
    ///
    /// The returned [`ReplayCursor`] holds an exclusive borrow of the index
    /// and **must** be `commit()`-ed to publish entries and restore the
    /// allocator. Forgetting to commit leaves the index unpublished — the
    /// `Drop` impl emits a debug-mode panic to catch this in tests.
    pub(crate) fn topology_name(&self) -> &'static str {
        self.scan.topology_name()
    }

    pub(crate) fn tile_count(&self) -> usize {
        self.scan.tile_count()
    }

    pub(crate) fn begin_visibility_fence(&self) -> Result<u64, crate::store::StoreError> {
        self.sequence.begin_fence()
    }

    pub(crate) fn active_visibility_fence(&self) -> Option<u64> {
        self.sequence.active_fence_token()
    }

    pub(crate) fn finish_visibility_fence_on_lanes(
        &self,
        token: u64,
        publish_to: Option<u64>,
        lanes: impl IntoIterator<Item = (u32, u64)>,
    ) -> Result<(), crate::store::StoreError> {
        self.sequence
            .finish_fence_on_lanes(token, publish_to, lanes)
    }

    pub(crate) fn note_visibility_fence_progress(
        &self,
        token: u64,
        start: u64,
        end: u64,
    ) -> Result<(), crate::store::StoreError> {
        self.sequence.note_fence_progress(token, start, end)
    }

    pub(crate) fn cancel_visibility_fence(
        &self,
        token: u64,
    ) -> Result<(), crate::store::StoreError> {
        let mut lane_ranges = BTreeMap::<u32, Vec<(u64, u64)>>::new();
        if let Some((start, end)) = self.sequence.active_fence_range() {
            for entry in self.by_id.iter() {
                let entry = entry.value();
                if entry.global_sequence >= start && entry.global_sequence < end {
                    lane_ranges.entry(entry.dag_lane).or_default().push((
                        entry.global_sequence,
                        entry.global_sequence.saturating_add(1),
                    ));
                }
            }
        }
        self.sequence.cancel_fence(token, lane_ranges)
    }

    pub(crate) fn cancelled_visibility_ranges(&self) -> CancelledVisibilityRanges {
        self.sequence.cancelled_ranges_snapshot()
    }

    pub(crate) fn restore_cancelled_visibility_ranges(&self, ranges: CancelledVisibilityRanges) {
        self.sequence.restore_cancelled_ranges(ranges);
    }
}

#[cfg(test)]
mod tests;