query-flow 0.18.0

//! Type-erased cache storage for query results and assets.

use std::any::{Any, TypeId};
use std::hash::Hasher;
use std::sync::Arc;

use whale::ShardedMap;

use crate::asset::{AssetKey, AssetLocator, DurabilityLevel};
use crate::key::{AssetCacheKey, FullCacheKey};
use crate::query::Query;
use crate::runtime::DbDispatch;

/// Cached query result (success or user error).
///
/// This is used internally to store both successful outputs and user errors
/// in the cache, enabling early cutoff optimization for error results.
#[derive(Clone)]
pub enum CachedValue<T> {
    /// Successful query output.
    Ok(T),
    /// User error from `QueryError::UserError`.
    UserError(Arc<anyhow::Error>),
}

/// Type-erased cache entry stored in whale's Node.data.
///
/// This is stored directly in the whale runtime's nodes for atomic
/// update of both cached value and dependency tracking state.
#[derive(Clone)]
pub enum CachedEntry {
    /// Successful query output (type-erased).
    Ok(Arc<dyn Any + Send + Sync>),
    /// User error from query.
    UserError(Arc<anyhow::Error>),
    /// Asset is ready with value (type-erased).
    AssetReady(Arc<dyn Any + Send + Sync>),
    /// Asset error (from locator or resolve_asset_error).
    ///
    /// Unlike query errors which use `UserError`, asset errors are stored
    /// separately to maintain distinct retrieval paths for queries and assets.
    AssetError(Arc<anyhow::Error>),
}

impl CachedEntry {
    /// Convert to typed [`CachedValue`] by downcasting.
    ///
    /// This is intended for query results only. Returns `None` for asset variants
    /// (`AssetReady`, `AssetError`) because queries and assets have distinct
    /// retrieval paths and type semantics:
    /// - Queries use `CachedValue<Arc<T>>` with `Ok`/`UserError` variants
    /// - Assets use dedicated methods like `to_asset_value()` or `to_asset_error()`
    pub fn to_cached_value<T: Send + Sync + 'static>(&self) -> Option<CachedValue<Arc<T>>> {
        match self {
            CachedEntry::Ok(arc) => arc.clone().downcast::<T>().ok().map(CachedValue::Ok),
            CachedEntry::UserError(e) => Some(CachedValue::UserError(e.clone())),
            CachedEntry::AssetReady(_) | CachedEntry::AssetError(_) => None,
        }
    }
}

/// Result of locating an asset (type-erased).
pub(crate) enum ErasedLocateResult {
    /// Asset is immediately available.
    Ready {
        value: Arc<dyn Any + Send + Sync>,
        durability: DurabilityLevel,
    },
    /// Asset needs to be loaded asynchronously.
    Pending,
}

/// Type-erased asset locator.
///
/// This wraps an `AssetLocator<K>` implementation and provides type-erased
/// access via monomorphized function pointers captured at registration time.
///
/// Three function pointers are stored:
/// - `locate_with_ctx_fn`: For calls from within a query (with dependency tracking on query)
/// - `locate_with_runtime_fn`: For direct calls from QueryRuntime (no tracking)
/// - `locate_with_locator_ctx_fn`: For locator-specific dependency tracking (deps on asset)
pub(crate) struct ErasedLocator<T: crate::Tracer> {
    /// The actual locator, type-erased.
    inner: Box<dyn Any + Send + Sync>,
    /// Locate function for LocatorContext calls (deps tracked on asset, not query).
    locate_with_locator_ctx_fn: DynAssetLocatorWithLocatorContext<T>,
}

type DynAssetLocatorWithLocatorContext<T> =
    fn(
        &dyn Any,
        &crate::runtime::LocatorContext<'_, T>,
        &dyn Any,
    ) -> Option<Result<ErasedLocateResult, crate::QueryError>>;

impl<T: crate::Tracer> ErasedLocator<T> {
    /// Create a new erased locator from a concrete implementation.
    pub fn new<K: AssetKey, L: AssetLocator<K>>(locator: L) -> Self {
        Self {
            inner: Box::new(locator),
            locate_with_locator_ctx_fn: erased_locate_with_locator_ctx::<K, L, T>,
        }
    }

    /// Attempt to locate an asset using LocatorContext (deps tracked on asset, not query).
    ///
    /// Returns `None` if the key type doesn't match.
    pub fn locate_with_locator_ctx(
        &self,
        locator_ctx: &crate::runtime::LocatorContext<'_, T>,
        key: &dyn Any,
    ) -> Option<Result<ErasedLocateResult, crate::QueryError>> {
        (self.locate_with_locator_ctx_fn)(&*self.inner, locator_ctx, key)
    }
}

/// Monomorphized locate function for LocatorContext calls.
/// Dependencies are tracked on the asset itself, not the calling query.
fn erased_locate_with_locator_ctx<K: AssetKey, L: AssetLocator<K>, T: crate::Tracer>(
    locator: &dyn Any,
    locator_ctx: &crate::runtime::LocatorContext<'_, T>,
    key: &dyn Any,
) -> Option<Result<ErasedLocateResult, crate::QueryError>> {
    let locator = locator.downcast_ref::<L>()?;
    let key = key.downcast_ref::<K>()?;
    let db = DbDispatch::LocatorContext(locator_ctx);
    Some(locator.locate(&db, key).map(|result| match result {
        crate::asset::LocateResult::Ready { value, durability } => ErasedLocateResult::Ready {
            value: Arc::new(value) as Arc<dyn Any + Send + Sync>,
            durability,
        },
        crate::asset::LocateResult::Pending => ErasedLocateResult::Pending,
    }))
}

/// Thread-safe storage for asset locators.
pub(crate) struct LocatorStorage<T: crate::Tracer> {
    /// Map from AssetKey TypeId to type-erased locator
    locators: ShardedMap<TypeId, Arc<ErasedLocator<T>>>,
}

impl<T: crate::Tracer> Default for LocatorStorage<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: crate::Tracer> LocatorStorage<T> {
    /// Create a new empty locator storage.
    pub fn new() -> Self {
        Self {
            locators: ShardedMap::new(),
        }
    }

    /// Register a locator for a specific asset key type.
    pub fn insert<K: AssetKey, L: AssetLocator<K>>(&self, locator: L) {
        self.locators.insert(
            TypeId::of::<K>(),
            Arc::new(ErasedLocator::new::<K, L>(locator)),
        );
    }

    /// Attempt to locate an asset using the registered locator (with LocatorContext).
    ///
    /// This variant tracks dependencies on the asset itself, not the calling query.
    /// Returns `None` if no locator is registered for the key type or if the key
    /// type doesn't match.
    pub fn locate_with_locator_ctx(
        &self,
        key_type: TypeId,
        locator_ctx: &crate::runtime::LocatorContext<'_, T>,
        key: &dyn Any,
    ) -> Option<Result<ErasedLocateResult, crate::QueryError>> {
        // Clone the locator out so no map lock is held while user code runs.
        self.locators
            .get(&key_type)
            .and_then(|locator| locator.locate_with_locator_ctx(locator_ctx, key))
    }
}

/// Thread-safe storage for pending asset requests.
pub(crate) struct PendingStorage {
    /// Map from AssetCacheKey to type-erased key
    pending: ShardedMap<AssetCacheKey, Arc<dyn Any + Send + Sync>>,
}

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

impl PendingStorage {
    /// Create a new empty pending storage.
    pub fn new() -> Self {
        Self {
            pending: ShardedMap::new(),
        }
    }

    /// Add a pending asset request.
    pub fn insert<K: AssetKey>(&self, asset_key: AssetCacheKey, key: K) {
        self.pending
            .insert(asset_key, Arc::new(key) as Arc<dyn Any + Send + Sync>);
    }

    /// Remove a pending asset request.
    pub fn remove(&self, key: &AssetCacheKey) -> bool {
        self.pending.remove(key).is_some()
    }

    /// Check if there are any pending assets.
    pub fn is_empty(&self) -> bool {
        self.pending.is_empty()
    }

    /// Get all pending assets of a specific type.
    pub fn get_of_type<K: AssetKey>(&self) -> Vec<K> {
        let key_type = TypeId::of::<K>();
        let mut keys = Vec::new();
        self.pending.for_each(|k, v| {
            if k.asset_key_type() == key_type {
                if let Some(key) = v.downcast_ref::<K>() {
                    keys.push(key.clone());
                }
            }
        });
        keys
    }

    /// Get all pending assets as PendingAsset.
    pub fn get_all(&self) -> Vec<crate::asset::PendingAsset> {
        let mut assets = Vec::new();
        self.pending.for_each(|k, v| {
            assets.push(crate::asset::PendingAsset::new_from_parts(
                k.asset_key_type(),
                &k.debug_repr(),
                v.clone(),
            ));
        });
        assets
    }
}

/// Thread-safe registry for tracking query instances by type.
///
/// Used by `list_queries` to enumerate all registered queries of a specific type.
pub(crate) struct QueryRegistry {
    /// Map from Query TypeId to per-type registry
    entries: ShardedMap<TypeId, Arc<QueryTypeRegistry>>,
}

/// Per-type registry for queries.
#[derive(Default)]
struct QueryTypeRegistry {
    /// Map from key_hash to type-erased query instance
    queries: ShardedMap<u64, Arc<dyn Any + Send + Sync>>,
}

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

impl QueryRegistry {
    /// Create a new empty query registry.
    pub fn new() -> Self {
        Self {
            entries: ShardedMap::new(),
        }
    }

    /// Register a query instance. Returns `true` if this was a new entry.
    pub fn register<Q: Query>(&self, query: &Q) -> bool {
        let type_id = TypeId::of::<Q>();
        // Hash the query itself since it's now the cache key
        let mut hasher = ahash::AHasher::default();
        query.dyn_hash(&mut hasher);
        let key_hash = hasher.finish();

        // Get or create the per-type registry
        let (type_registry, _) = self.entries.get_or_insert_with(type_id, Default::default);
        type_registry.queries.compute(key_hash, |slot| {
            if slot.is_some() {
                return false; // Already registered
            }
            *slot = Some(Arc::new(query.clone()) as Arc<dyn Any + Send + Sync>);
            true
        })
    }

    /// Get all query instances of type Q.
    pub fn get_all<Q: Query>(&self) -> Vec<Q> {
        let type_id = TypeId::of::<Q>();
        let Some(type_registry) = self.entries.get(&type_id) else {
            return Vec::new();
        };
        let mut queries = Vec::new();
        type_registry.queries.for_each(|_, arc| {
            if let Some(query) = arc.downcast_ref::<Q>() {
                queries.push(query.clone());
            }
        });
        queries
    }

    /// Remove a query from the registry. Returns `true` if it was present.
    pub fn remove<Q: Query>(&self, query: &Q) -> bool {
        let type_id = TypeId::of::<Q>();
        let mut hasher = ahash::AHasher::default();
        query.dyn_hash(&mut hasher);
        let key_hash = hasher.finish();

        self.entries
            .get(&type_id)
            .is_some_and(|type_registry| type_registry.queries.remove(&key_hash).is_some())
    }
}

/// Thread-safe registry for tracking asset keys by type.
///
/// Used by `list_asset_keys` to enumerate all registered asset keys of a specific type.
pub(crate) struct AssetKeyRegistry {
    /// Map from AssetKey TypeId to per-type registry
    entries: ShardedMap<TypeId, Arc<AssetKeyTypeRegistry>>,
}

/// Per-type registry for asset keys.
#[derive(Default)]
struct AssetKeyTypeRegistry {
    /// Map from key_hash to type-erased asset key instance
    keys: ShardedMap<u64, Arc<dyn Any + Send + Sync>>,
}

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

impl AssetKeyRegistry {
    /// Create a new empty asset key registry.
    pub fn new() -> Self {
        Self {
            entries: ShardedMap::new(),
        }
    }

    /// Register an asset key. Returns `true` if this was a new entry.
    pub fn register<K: AssetKey>(&self, key: &K) -> bool {
        let type_id = TypeId::of::<K>();
        let mut hasher = ahash::AHasher::default();
        key.dyn_hash(&mut hasher);
        let key_hash = hasher.finish();

        let (type_registry, _) = self.entries.get_or_insert_with(type_id, Default::default);
        type_registry.keys.compute(key_hash, |slot| {
            if slot.is_some() {
                return false; // Already registered
            }
            *slot = Some(Arc::new(key.clone()) as Arc<dyn Any + Send + Sync>);
            true
        })
    }

    /// Get all asset keys of type K.
    pub fn get_all<K: AssetKey>(&self) -> Vec<K> {
        let type_id = TypeId::of::<K>();
        let Some(type_registry) = self.entries.get(&type_id) else {
            return Vec::new();
        };
        let mut keys = Vec::new();
        type_registry.keys.for_each(|_, arc| {
            if let Some(key) = arc.downcast_ref::<K>() {
                keys.push(key.clone());
            }
        });
        keys
    }

    /// Remove an asset key. Returns `true` if it was present.
    pub fn remove<K: AssetKey>(&self, key: &K) -> bool {
        let type_id = TypeId::of::<K>();
        let mut hasher = ahash::AHasher::default();
        key.dyn_hash(&mut hasher);
        let key_hash = hasher.finish();

        self.entries
            .get(&type_id)
            .is_some_and(|type_registry| type_registry.keys.remove(&key_hash).is_some())
    }
}

/// Trait for type-erased query verification.
///
/// Allows re-executing a query to verify it hasn't changed.
pub(crate) trait AnyVerifier: Send + Sync + 'static {
    /// Verify the query by re-executing it.
    /// The runtime is passed as `&dyn Any` and will be downcast to the correct type.
    fn verify(&self, runtime: &dyn std::any::Any) -> Result<(), crate::QueryError>;
}

/// Concrete verifier for a specific query type.
pub(crate) struct QueryVerifier<Q: Query, T: crate::Tracer> {
    query: Q,
    _marker: std::marker::PhantomData<T>,
}

impl<Q: Query, T: crate::Tracer> QueryVerifier<Q, T> {
    pub fn new(query: Q) -> Self {
        Self {
            query,
            _marker: std::marker::PhantomData,
        }
    }
}

impl<Q: Query, T: crate::Tracer + 'static> AnyVerifier for QueryVerifier<Q, T> {
    fn verify(&self, runtime: &dyn std::any::Any) -> Result<(), crate::QueryError> {
        // Downcast to the correct runtime type
        if let Some(runtime) = runtime.downcast_ref::<crate::QueryRuntime<T>>() {
            match runtime.query(self.query.clone()) {
                Ok(_) => Ok(()),
                // UserError is a valid cached result, verification succeeded
                Err(crate::QueryError::UserError(_)) => Ok(()),
                // System errors mean verification failed
                Err(e) => Err(e),
            }
        } else {
            // This should never happen if used correctly
            Ok(())
        }
    }
}

/// Concrete verifier for a specific asset type.
///
/// Re-accesses the asset to trigger locator re-execution if the asset's
/// dependencies have changed. Uses early cutoff: if the asset value is
/// unchanged, changed_at stays the same and dependent queries won't re-run.
pub(crate) struct AssetVerifier<K: AssetKey, T: crate::Tracer> {
    key: K,
    _marker: std::marker::PhantomData<T>,
}

impl<K: AssetKey, T: crate::Tracer> AssetVerifier<K, T> {
    pub fn new(key: K) -> Self {
        Self {
            key,
            _marker: std::marker::PhantomData,
        }
    }
}

impl<K: AssetKey, T: crate::Tracer + 'static> AnyVerifier for AssetVerifier<K, T> {
    fn verify(&self, runtime: &dyn std::any::Any) -> Result<(), crate::QueryError> {
        // Downcast to the correct runtime type
        if let Some(runtime) = runtime.downcast_ref::<crate::QueryRuntime<T>>() {
            // Re-access the asset, which triggers locator if asset's deps are invalid
            match runtime.get_asset(self.key.clone()) {
                Ok(_) => Ok(()),
                // UserError is a valid cached result, verification succeeded
                Err(crate::QueryError::UserError(_)) => Ok(()),
                // System errors mean verification failed
                Err(e) => Err(e),
            }
        } else {
            // This should never happen if used correctly
            Ok(())
        }
    }
}

/// Thread-safe storage for query verifiers.
///
/// Verifiers can re-execute queries to verify they haven't changed.
/// Used by the "verify-then-decide" pattern to verify dependencies before
/// deciding whether to recompute a dependent query.
pub(crate) struct VerifierStorage {
    /// Map from FullCacheKey to verifier
    verifiers: ShardedMap<FullCacheKey, Arc<dyn AnyVerifier>>,
}

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

impl VerifierStorage {
    /// Create a new empty verifier storage.
    pub fn new() -> Self {
        Self {
            verifiers: ShardedMap::new(),
        }
    }

    /// Register a verifier for a query.
    pub fn insert<Q: Query, T: crate::Tracer + 'static>(&self, key: FullCacheKey, query: Q) {
        self.verifiers
            .insert(key, Arc::new(QueryVerifier::<Q, T>::new(query)));
    }

    /// Register a verifier for an asset.
    pub fn insert_asset<K: AssetKey, T: crate::Tracer + 'static>(
        &self,
        key: FullCacheKey,
        asset_key: K,
    ) {
        self.verifiers
            .insert(key, Arc::new(AssetVerifier::<K, T>::new(asset_key)));
    }

    /// Get a verifier for a query key.
    pub fn get(&self, key: &FullCacheKey) -> Option<Arc<dyn AnyVerifier>> {
        self.verifiers.get(key)
    }

    /// Remove a verifier.
    pub fn remove(&self, key: &FullCacheKey) -> bool {
        self.verifiers.remove(key).is_some()
    }
}

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

    #[test]
    fn test_cached_entry_to_cached_value() {
        // Test Ok variant
        let entry = CachedEntry::Ok(Arc::new("hello".to_string()) as Arc<dyn Any + Send + Sync>);
        let result: Option<CachedValue<Arc<String>>> = entry.to_cached_value();
        assert!(result.is_some());
        match result.unwrap() {
            CachedValue::Ok(v) => assert_eq!(*v, "hello"),
            CachedValue::UserError(_) => panic!("expected Ok"),
        }

        // Test UserError variant
        let err = Arc::new(anyhow::anyhow!("test error"));
        let entry = CachedEntry::UserError(err.clone());
        let result: Option<CachedValue<Arc<String>>> = entry.to_cached_value();
        assert!(result.is_some());
        match result.unwrap() {
            CachedValue::Ok(_) => panic!("expected UserError"),
            CachedValue::UserError(e) => assert_eq!(e.to_string(), "test error"),
        }

        // Test type mismatch (Ok with wrong type)
        let entry = CachedEntry::Ok(Arc::new(42i32) as Arc<dyn Any + Send + Sync>);
        let result: Option<CachedValue<Arc<String>>> = entry.to_cached_value();
        assert!(result.is_none()); // downcast fails
    }
}