stemedb/crates/stemedb-lens/src/epoch_aware/mod.rs

//! Epoch-Aware Lens: Filters superseded epochs before resolution.
//!
//! This lens implements the "Paradigm" concept from the architecture: epochs
//! can supersede each other, and assertions from superseded epochs should be
//! excluded from query results.
//!
//! # Design Philosophy
//!
//! Follows the "Deep Module" principle:
//! - Simple interface: `resolve_async(&[Assertion])` returns winner
//! - Complex implementation: Walks epoch supersession chains, filters candidates
//! - Composable: Wraps any inner lens for final resolution
//!
//! # Resolution Strategy
//!
//! 1. Collect all unique epoch IDs from candidate assertions
//! 2. For each epoch, read `E:{epoch_id}` from store
//! 3. Walk the `supersedes` chain to build a set of superseded epoch IDs
//! 4. Filter candidates: exclude any assertion whose epoch is superseded
//! 5. Delegate filtered candidates to inner lens for final winner selection
//!
//! # Fail-Open Semantics
//!
//! Missing or corrupted epoch records do NOT cause assertion exclusion:
//! - Assertion references epoch X, but `E:X` doesn't exist → include assertion
//! - Epoch record fails to deserialize → include assertion
//! - Cycle detected in supersession chain → stop walking, include assertions
//!
//! This ensures data availability is not compromised by metadata issues.
//!
//! # Example
//!
//! ```ignore
//! use stemedb_lens::{EpochAwareLens, RecencyLens};
//! use stemedb_storage::HybridStore;
//! use std::sync::Arc;
//!
//! let store = Arc::new(HybridStore::open("./data").expect("store"));
//! let lens = EpochAwareLens::with_recency(store);
//!
//! let resolution = lens.resolve_async(&candidates).await;
//! ```

use crate::recency::RecencyLens;
use crate::traits::{Lens, Resolution};
use crate::vote_aware_consensus::AsyncLens;
use async_trait::async_trait;
use std::collections::HashSet;
use std::sync::Arc;
use stemedb_core::serde::deserialize;
use stemedb_core::types::{Assertion, Epoch, EpochId};
use stemedb_storage::{key_codec, KVStore};
use tracing::{debug, instrument, warn};

/// Wrapper to use a sync Lens with EpochAwareLens.
///
/// This allows using synchronous lenses (like `RecencyLens`, `ConsensusLens`)
/// with the async `EpochAwareLens`.
pub struct SyncLensWrapper<L>(pub L);

#[async_trait]
impl<L: Lens + 'static> AsyncLens for SyncLensWrapper<L> {
    async fn resolve_async(&self, candidates: &[Assertion]) -> Resolution {
        self.0.resolve(candidates)
    }

    fn name(&self) -> &'static str {
        self.0.name()
    }
}

/// Maximum depth for walking supersession chains.
///
/// Prevents infinite loops and excessive I/O for pathological chains.
/// In practice, epoch chains should be much shorter (3-7 levels).
const MAX_SUPERSESSION_DEPTH: usize = 100;

/// Epoch-Aware Lens: Filters assertions from superseded epochs before delegation.
///
/// This is a decorator/wrapper lens that:
/// 1. Reads epoch records to determine which epochs are superseded
/// 2. Filters out assertions from superseded epochs
/// 3. Delegates remaining candidates to an inner lens
///
/// # Type Parameters
///
/// - `S`: The KVStore implementation for reading epoch records
/// - `L`: The inner lens for final resolution after filtering
pub struct EpochAwareLens<S, L> {
    store: Arc<S>,
    inner: L,
}

impl<S: KVStore> EpochAwareLens<S, SyncLensWrapper<RecencyLens>> {
    /// Create an EpochAwareLens with RecencyLens as the inner resolution strategy.
    ///
    /// This is the common case: filter superseded epochs, then pick the most recent.
    pub fn with_recency(store: Arc<S>) -> Self {
        Self { store, inner: SyncLensWrapper(RecencyLens) }
    }
}

impl<S: KVStore, L> EpochAwareLens<S, L> {
    /// Create an EpochAwareLens with a custom inner lens.
    ///
    /// # Arguments
    ///
    /// * `store` - Arc to KVStore for reading `E:{epoch_id}` keys
    /// * `inner` - Inner lens to delegate to after filtering superseded epochs
    pub fn new(store: Arc<S>, inner: L) -> Self {
        Self { store, inner }
    }

    /// Read an epoch record from the store.
    ///
    /// Returns `None` if the epoch doesn't exist or fails to deserialize.
    async fn read_epoch(&self, epoch_id: &EpochId) -> Option<Epoch> {
        let key = key_codec::epoch_key(&hex::encode(epoch_id));

        match self.store.get(&key).await {
            Ok(Some(bytes)) => match deserialize::<Epoch>(&bytes) {
                Ok(epoch) => Some(epoch),
                Err(e) => {
                    warn!(
                        epoch_id = %hex::encode(epoch_id),
                        error = %e,
                        "Failed to deserialize epoch record, treating as missing"
                    );
                    None
                }
            },
            Ok(None) => {
                debug!(
                    epoch_id = %hex::encode(epoch_id),
                    "Epoch record not found in store"
                );
                None
            }
            Err(e) => {
                warn!(
                    epoch_id = %hex::encode(epoch_id),
                    error = %e,
                    "Failed to read epoch from store, treating as missing"
                );
                None
            }
        }
    }

    /// Check if an epoch is superseded using O(1) marker lookup.
    ///
    /// The IngestWorker writes `SUPERSEDED:{epoch_id}` markers at epoch ingestion
    /// time for the full transitive closure of superseded epochs. This enables
    /// constant-time "is superseded?" checks instead of O(chain_length) walks.
    ///
    /// # Fail-Open Semantics
    ///
    /// - Marker exists → epoch is superseded (return true)
    /// - Marker doesn't exist → epoch is NOT superseded (return false)
    /// - Storage error → treat as NOT superseded (fail-open)
    async fn is_epoch_superseded(&self, epoch_id: &EpochId) -> bool {
        let key = key_codec::superseded_key(&hex::encode(epoch_id));
        match self.store.get(&key).await {
            Ok(Some(_)) => {
                debug!(epoch_id = %hex::encode(epoch_id), "Epoch is superseded (marker found)");
                true
            }
            Ok(None) => false,
            Err(e) => {
                warn!(
                    epoch_id = %hex::encode(epoch_id),
                    error = %e,
                    "Failed to check superseded marker, treating as not superseded (fail-open)"
                );
                false
            }
        }
    }

    /// Compute the set of superseded epoch IDs using O(1) marker lookups.
    ///
    /// For each unique epoch in the candidate set, checks for a `SUPERSEDED:`
    /// marker key. If present, the epoch is superseded and should be filtered.
    ///
    /// # Performance
    ///
    /// This is O(unique_epochs) with O(1) per epoch, compared to the previous
    /// O(unique_epochs * chain_length) approach that walked supersession chains.
    ///
    /// # Fail-Open Semantics
    ///
    /// Missing markers (e.g., for epochs ingested before cascade logic was added)
    /// are treated as "not superseded" - assertions pass through. This ensures
    /// backward compatibility with existing data.
    async fn compute_superseded_epochs(&self, epochs: &HashSet<EpochId>) -> HashSet<EpochId> {
        let mut superseded = HashSet::new();

        for epoch_id in epochs {
            if self.is_epoch_superseded(epoch_id).await {
                superseded.insert(*epoch_id);
            }
        }

        superseded
    }

    /// Walk the supersession chain starting from `epoch_id` (legacy fallback).
    ///
    /// This method is kept for backward compatibility with existing tests and
    /// for scenarios where cascade markers haven't been written yet.
    /// The primary `compute_superseded_epochs` now uses O(1) marker lookups.
    #[allow(dead_code)]
    async fn walk_supersession_chain_legacy(
        &self,
        start_epoch_id: &EpochId,
        superseded: &mut HashSet<EpochId>,
        visited: &mut HashSet<EpochId>,
    ) {
        let mut current_id = *start_epoch_id;
        let mut depth = 0;

        loop {
            // Cycle detection
            if !visited.insert(current_id) {
                debug!(
                    epoch_id = %hex::encode(current_id),
                    "Cycle detected in supersession chain, stopping walk"
                );
                break;
            }

            // Max depth guard
            if depth >= MAX_SUPERSESSION_DEPTH {
                warn!(
                    start_epoch = %hex::encode(start_epoch_id),
                    depth,
                    "Supersession chain exceeded max depth, stopping walk"
                );
                break;
            }

            // Read epoch record
            let epoch = match self.read_epoch(&current_id).await {
                Some(e) => e,
                None => break, // Missing epoch - stop walking
            };

            // Check for supersession
            match epoch.supersedes {
                Some(superseded_id) => {
                    // This epoch supersedes another - add the superseded one to the set
                    superseded.insert(superseded_id);
                    current_id = superseded_id;
                    depth += 1;
                }
                None => break, // End of chain
            }
        }
    }
}

#[async_trait]
impl<S: KVStore + 'static, L: AsyncLens + 'static> AsyncLens for EpochAwareLens<S, L> {
    #[instrument(skip(self, candidates), fields(
        candidates_count = candidates.len(),
        lens = "EpochAware"
    ))]
    async fn resolve_async(&self, candidates: &[Assertion]) -> Resolution {
        if candidates.is_empty() {
            return Resolution::empty();
        }

        // Fast path: single candidate passes through
        if candidates.len() == 1 {
            return self.inner.resolve_async(candidates).await;
        }

        // Collect unique epochs from candidates
        let epochs: HashSet<EpochId> = candidates.iter().filter_map(|a| a.epoch).collect();

        // If no assertions have epochs, delegate directly to inner lens
        if epochs.is_empty() {
            debug!("No epochs in candidates, delegating directly to inner lens");
            return self.inner.resolve_async(candidates).await;
        }

        debug!(unique_epochs = epochs.len(), "Computing superseded epochs");

        // Build the superseded set
        let superseded = self.compute_superseded_epochs(&epochs).await;

        debug!(superseded_count = superseded.len(), "Found superseded epochs");

        // Filter candidates: exclude assertions from superseded epochs
        let filtered: Vec<&Assertion> = candidates
            .iter()
            .filter(|a| {
                match a.epoch {
                    // No epoch: always include
                    None => true,
                    // Has epoch: include if NOT superseded
                    Some(epoch_id) => !superseded.contains(&epoch_id),
                }
            })
            .collect();

        let filtered_count = filtered.len();
        let excluded_count = candidates.len() - filtered_count;

        debug!(filtered_count, excluded_count, "Filtered candidates by epoch supersession");

        // Handle case where all candidates were filtered
        if filtered.is_empty() {
            debug!("All candidates were from superseded epochs, returning empty resolution");
            return Resolution::empty();
        }

        // Clone filtered assertions for delegation to inner lens
        let filtered_owned: Vec<Assertion> = filtered.into_iter().cloned().collect();

        // Delegate to inner lens
        self.inner.resolve_async(&filtered_owned).await
    }

    fn name(&self) -> &'static str {
        "EpochAware"
    }
}

impl<S: KVStore, L: Lens> EpochAwareLens<S, SyncLensWrapper<L>> {
    /// Create an EpochAwareLens wrapping a synchronous lens.
    ///
    /// This allows using any `Lens` implementation (like `ConsensusLens`)
    /// with epoch-aware filtering.
    pub fn with_sync_lens(store: Arc<S>, inner: L) -> Self {
        Self { store, inner: SyncLensWrapper(inner) }
    }
}

#[cfg(test)]
mod tests;