tidaldb/tidal/src/db/open.rs

//! Schema-aware database open logic.

use std::sync::{Arc, atomic::AtomicU64};

use roaring::RoaringBitmap;

use super::{
    config::StorageMode,
    state_rebuild::{rebuild_entity_state, rebuild_item_indexes},
    storage_box::StorageBox,
    wal_bridge::WalHandleWriter,
};
use crate::{
    entities::{
        CreatorItemsBitmap, HardNegIndex, InteractionLedger, PreferenceVectors, UserStateIndex,
    },
    ranking::{builtins::register_builtins, profile::RankingProfile, registry::ProfileRegistry},
    schema::{DurabilityError, EntityId, Schema, TidalError, Timestamp},
    signals::{NoopWalWriter, SignalLedger, SignalTypeId},
    storage::{
        InMemoryBackend,
        indexes::{bitmap::BitmapIndex, range::RangeIndex},
        vector::registry::EmbeddingSlotRegistry,
    },
    wal::{WalConfig, WalHandle},
};

/// Default preference-vector dimensionality used when no schema embedding slot
/// declares one.
///
/// The cold-start path ([`TidalDb::new`], no schema) and the
/// no-embedding-slot schema path both fall back to this single value so a future
/// change to the default cannot drift between the two open branches and the
/// schema-less constructor. Keep it in sync with the typical embedding width the
/// engine ships against.
pub const DEFAULT_PREFERENCE_DIM: usize = 128;

/// Resolve the preference-vector dimensionality from a schema.
///
/// Uses the first declared embedding slot's dimensions, falling back to
/// [`DEFAULT_PREFERENCE_DIM`] when the schema declares no embedding slot. This is
/// the single source of truth for the fallback used by every schema-aware open
/// branch.
pub fn preference_dim_for(schema: &Schema) -> usize {
    schema
        .embedding_slots()
        .first()
        .map_or(DEFAULT_PREFERENCE_DIM, |s| s.dimensions)
}

/// Bundle returned by [`TidalDb::open_with_schema`] to avoid a fragile tuple.
///
/// Each field is consumed by [`TidalDb::from_parts`] during construction.
pub struct OpenResult {
    pub storage: StorageBox,
    pub ledger: Arc<SignalLedger>,
    pub wal: Option<WalHandle>,
    pub last_seq: Arc<AtomicU64>,
    pub profile_registry: ProfileRegistry,
    pub category_index: BitmapIndex,
    pub format_index: BitmapIndex,
    pub creator_index: BitmapIndex,
    pub tag_index: BitmapIndex,
    pub duration_index: RangeIndex<u32>,
    pub created_at_index: RangeIndex<u64>,
    pub universe: RoaringBitmap,
    pub embedding_registry: EmbeddingSlotRegistry,
    pub schema_def: Schema,
    // M3 entity state rebuilt from durable storage.
    pub creator_items: CreatorItemsBitmap,
    pub user_state: UserStateIndex,
    pub hard_negatives: HardNegIndex,
    pub interaction_ledger: InteractionLedger,
    pub preference_vectors: PreferenceVectors,
    /// Session journal events recovered on startup (for session crash recovery).
    pub session_events: Vec<crate::wal::format::SessionWalEvent>,
}

impl super::TidalDb {
    /// Open storage and ledger components for the given schema.
    ///
    /// Called from `TidalDbBuilder::open()` when `with_schema()` was called.
    /// Returns an `OpenResult` containing all components needed by `from_parts`.
    #[allow(clippy::too_many_lines)] // Construction logic is inherently verbose; splitting would scatter initialization flow.
    pub(crate) fn open_with_schema(
        config: &super::Config,
        schema: Schema,
        schema_profiles: Vec<RankingProfile>,
    ) -> crate::Result<OpenResult> {
        let last_seq = Arc::new(AtomicU64::new(0));

        // Initialize profile registry with builtins.
        let mut profile_registry = ProfileRegistry::new();
        register_builtins(&mut profile_registry).map_err(|e| {
            TidalError::internal("open", format!("failed to register builtin profiles: {e}"))
        })?;

        // Bind the schema so schema-defined profiles are validated against the
        // actual signal vocabulary (builtins are exempt -- already registered,
        // and they intentionally reference a superset of any one schema).
        profile_registry.with_schema(schema.clone());

        // Register schema-defined profiles, overriding matching builtins.
        for profile in schema_profiles {
            let name = profile.name.clone();
            profile_registry.override_register(profile).map_err(|e| {
                TidalError::internal(
                    "open",
                    format!("failed to register schema profile '{name}': {e}"),
                )
            })?;
        }

        // Initialize M2 indexes (empty -- populated as items are written).
        let category_index = BitmapIndex::new("category");
        let format_index = BitmapIndex::new("format");
        let creator_index = BitmapIndex::new("creator");
        let tag_index = BitmapIndex::new("tags");
        let duration_index = RangeIndex::new("duration");
        let created_at_index = RangeIndex::new("created_at");
        let universe = RoaringBitmap::new();
        let embedding_registry = EmbeddingSlotRegistry::new();

        let schema_def = schema.clone();

        match config.mode {
            StorageMode::Ephemeral => {
                let storage = StorageBox::Memory {
                    items: InMemoryBackend::default(),
                    users: InMemoryBackend::default(),
                    creators: InMemoryBackend::default(),
                };
                let ledger = Arc::new(SignalLedger::new(schema, Box::new(NoopWalWriter)));

                // Read preference vector dimensionality from the schema.
                let pref_dim = preference_dim_for(&schema_def);

                Ok(OpenResult {
                    storage,
                    ledger,
                    wal: None,
                    last_seq,
                    profile_registry,
                    category_index,
                    format_index,
                    creator_index,
                    tag_index,
                    duration_index,
                    created_at_index,
                    universe,
                    embedding_registry,
                    schema_def,
                    creator_items: CreatorItemsBitmap::new(),
                    user_state: UserStateIndex::new(),
                    hard_negatives: HardNegIndex::new(),
                    interaction_ledger: InteractionLedger::new(),
                    preference_vectors: PreferenceVectors::new(pref_dim),
                    session_events: Vec::new(),
                })
            }
            StorageMode::Persistent => {
                let data_dir = config.data_dir.as_ref().ok_or_else(|| {
                    TidalError::internal("open", "persistent mode requires data_dir")
                })?;

                // Open fjall storage.
                let fjall_storage =
                    crate::storage::FjallStorage::open(data_dir).map_err(TidalError::from)?;
                let storage = StorageBox::Fjall(fjall_storage);

                // Build WAL config. `WalConfig.dir` is the *parent* of the
                // "wal/" subdirectory; the WAL therefore defaults to
                // {data_dir}/wal. When the builder supplied an explicit `wal_dir`
                // override, honor it via `wal_dir_override` so the WAL is created
                // at exactly that path — and so the readers below
                // (`config.resolved_wal_dir()`) point at the same place. The
                // builder resolves `config.wal_dir` to the override-or-default at
                // open time, so this threads one source of truth to the writer.
                let wal_config = WalConfig {
                    dir: data_dir.clone(),
                    wal_dir_override: config.wal_dir.clone(),
                    ..WalConfig::default()
                };

                let (wal, replayed_events, session_events) =
                    WalHandle::open(wal_config).map_err(|e| {
                        TidalError::Durability(DurabilityError {
                            message: format!("WAL open failed: {e}"),
                        })
                    })?;

                // Build the WAL bridge for the ledger.
                let wal_writer =
                    Box::new(WalHandleWriter::new(wal.sender(), Arc::clone(&last_seq)));

                // Construct the ledger.
                let ledger = Arc::new(SignalLedger::new(schema, wal_writer));

                // Restore signal state from the last checkpoint.
                if let Err(e) = ledger.restore(storage.items_engine()) {
                    tracing::warn!(
                        error = %e,
                        "signal ledger restore failed; starting from empty state"
                    );
                }

                // Replay WAL events that post-date the checkpoint.
                for event in replayed_events {
                    let type_id = SignalTypeId::new(u16::from(event.signal_type));
                    let entity_id = EntityId::new(event.entity_id);
                    let weight = f64::from(event.weight);
                    let timestamp = Timestamp::from_nanos(event.timestamp_nanos);
                    ledger.apply_wal_event(type_id, entity_id, weight, timestamp);
                }

                // Rebuild in-memory entity state from durable storage.
                // This reconstructs blocked/hidden/follows user state and
                // interaction weights from persisted relationship edges. The
                // creator-items bitmap is populated by `rebuild_item_indexes`
                // below (shared with the write path) so it cannot drift.
                let user_state = UserStateIndex::new();
                let creator_items = CreatorItemsBitmap::new();
                let interaction_ledger = InteractionLedger::new();
                rebuild_entity_state(&storage, &user_state, &interaction_ledger)?;

                // Restore the signal-driven UserStateIndex maps (completion /
                // saved-timestamp / liked) from their durable Tag::UserState rows
                // so the InProgress filter and DateSaved sort survive a restart.
                if let Err(e) = user_state.restore(storage.items_engine()) {
                    tracing::warn!(
                        error = %e,
                        "user-state restore failed; InProgress/DateSaved start empty"
                    );
                }

                // Restore hard negatives from durable Tag::HardNeg rows so a
                // skip/dislike-rejected item stays excluded after a crash.
                let hard_negatives = HardNegIndex::new();
                if let Err(e) = hard_negatives.restore(storage.items_engine()) {
                    tracing::warn!(
                        error = %e,
                        "hard-negative restore failed; rejected items may resurface"
                    );
                }

                // Rebuild item indexes (universe bitmap + bitmap/range indexes +
                // creator-items) from persisted metadata so queries and /health
                // work immediately. The shared `ItemIndexes::index` is the SAME
                // path the live write uses, so a restart cannot diverge.
                let mut universe = universe;
                rebuild_item_indexes(
                    &storage,
                    &mut universe,
                    &super::state_rebuild::ItemIndexes {
                        category: &category_index,
                        format: &format_index,
                        creator: &creator_index,
                        tag: &tag_index,
                        duration: &duration_index,
                        created_at: &created_at_index,
                        creator_items: &creator_items,
                    },
                )?;

                // Read preference vector dimensionality from the schema.
                let pref_dim = preference_dim_for(&schema_def);

                // Restore preference vectors from their Tag::Preference checkpoint
                // so personalized ranking does not snap to cold-start after a crash.
                let preference_vectors = PreferenceVectors::new(pref_dim);
                if let Err(e) = preference_vectors.restore(storage.items_engine()) {
                    tracing::warn!(
                        error = %e,
                        "preference vector restore failed; starting from cold-start"
                    );
                }

                Ok(OpenResult {
                    storage,
                    ledger,
                    wal: Some(wal),
                    last_seq,
                    profile_registry,
                    category_index,
                    format_index,
                    creator_index,
                    tag_index,
                    duration_index,
                    created_at_index,
                    universe,
                    embedding_registry,
                    schema_def,
                    creator_items,
                    user_state,
                    hard_negatives,
                    interaction_ledger,
                    preference_vectors,
                    session_events,
                })
            }
        }
    }
}