armdb 0.14.2

use std::marker::PhantomData;
use std::mem::size_of;

use crate::Key;
use crate::codec::Codec;
use crate::error::{DbError, DbResult};
use crate::lsm::map_fjall;

/// Key-only iterator over an `LsmTree`. Honors an explicit `reversed` flag
/// (not the tree's storage order) and uses [`fjall::Guard::key`] so value bytes
/// are not loaded on KV-separated storage.
pub(crate) struct LsmKeyIter<K> {
    inner: fjall::Iter,
    reversed: bool,
    _marker: PhantomData<K>,
}

/// Strict key-only iterator over an `LsmTree`. Unlike [`LsmKeyIter`], this
/// surfaces guard and key-layout errors and stops permanently after the first
/// failure, so callers cannot mistake a partial scan for a complete one.
pub struct TryLsmKeyIter<K> {
    inner: fjall::Iter,
    reversed: bool,
    failed: bool,
    #[cfg(test)]
    force_guard_error: bool,
    _marker: PhantomData<K>,
}

impl<K> TryLsmKeyIter<K> {
    pub(crate) fn new(inner: fjall::Iter, reversed: bool) -> Self {
        Self {
            inner,
            reversed,
            failed: false,
            #[cfg(test)]
            force_guard_error: false,
            _marker: PhantomData,
        }
    }

    #[cfg(test)]
    pub(crate) fn new_with_forced_guard_error(inner: fjall::Iter, reversed: bool) -> Self {
        Self {
            force_guard_error: true,
            ..Self::new(inner, reversed)
        }
    }
}

impl<K: Key> Iterator for TryLsmKeyIter<K> {
    type Item = DbResult<K>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.failed {
            return None;
        }
        #[cfg(test)]
        if self.force_guard_error {
            self.force_guard_error = false;
            self.failed = true;
            return Some(Err(DbError::Fjall(
                "test-injected lsm key guard error".to_owned(),
            )));
        }
        let guard = if self.reversed {
            self.inner.next_back()?
        } else {
            self.inner.next()?
        };
        let result = map_fjall(guard.key()).and_then(|raw| {
            (raw.len() == size_of::<K>())
                .then(|| K::from_bytes(raw.as_ref()))
                .ok_or(DbError::CorruptedEntry { offset: 0 })
        });
        self.failed = result.is_err();
        Some(result)
    }
}

impl<K> LsmKeyIter<K> {
    pub(crate) fn new(inner: fjall::Iter, reversed: bool) -> Self {
        Self {
            inner,
            reversed,
            _marker: PhantomData,
        }
    }
}

impl<K: Key> Iterator for LsmKeyIter<K> {
    type Item = K;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let g = if self.reversed {
                self.inner.next_back()
            } else {
                self.inner.next()
            }?;
            match g.key() {
                Ok(raw) => return Some(K::from_bytes(raw.as_ref())),
                Err(e) => tracing::warn!(%e, "lsm_key_iter: guard error, skipping"),
            }
        }
    }
}

/// Decoding iterator over an `LsmTree`. Honors the tree's `reversed` flag by
/// pulling from the back of fjall's `DoubleEndedIterator`. Rows whose value
/// fails to decode are skipped (logged at debug), matching `VarTypedIter`.
pub struct LsmIter<'a, K, T, C> {
    inner: fjall::Iter,
    reversed: bool,
    codec: &'a C,
    _marker: PhantomData<fn() -> (K, T)>,
}

impl<'a, K, T, C> LsmIter<'a, K, T, C> {
    pub(crate) fn new(inner: fjall::Iter, reversed: bool, codec: &'a C) -> Self {
        Self {
            inner,
            reversed,
            codec,
            _marker: PhantomData,
        }
    }
}

impl<K: Key, T, C: Codec<T>> LsmIter<'_, K, T, C> {
    /// Decode a raw fjall guard into a `(K, T)` pair. Returns `None` on a guard
    /// error or value decode error (both are logged).
    fn decode_guard(&self, g: fjall::Guard) -> Option<(K, T)> {
        match g.into_inner() {
            Ok((uk, uv)) => match self.codec.decode_from(uv.as_ref()) {
                Ok(v) => Some((K::from_bytes(uk.as_ref()), v)),
                Err(_) => {
                    tracing::debug!(value_len = uv.len(), "lsm_iter: decode error, skipping");
                    None
                }
            },
            Err(e) => {
                tracing::warn!(%e, "lsm_iter: guard error, skipping");
                None
            }
        }
    }
}

impl<K: Key, T, C: Codec<T>> Iterator for LsmIter<'_, K, T, C> {
    type Item = (K, T);

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let g = if self.reversed {
                self.inner.next_back()
            } else {
                self.inner.next()
            }?;
            if let Some(kv) = self.decode_guard(g) {
                return Some(kv);
            }
        }
    }
}

impl<K: Key, T, C: Codec<T>> DoubleEndedIterator for LsmIter<'_, K, T, C> {
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            let g = if self.reversed {
                self.inner.next()
            } else {
                self.inner.next_back()
            }?;
            if let Some(kv) = self.decode_guard(g) {
                return Some(kv);
            }
        }
    }
}

/// Strict decoding iterator over an `LsmTree`. Honors the tree's `reversed` flag
/// and surfaces storage/decode errors instead of skipping corrupt rows.
pub struct TryLsmIter<'a, K, T, C> {
    inner: fjall::Iter,
    codec: &'a C,
    reversed: bool,
    failed: bool,
    _marker: PhantomData<fn() -> (K, T)>,
}

impl<'a, K, T, C> TryLsmIter<'a, K, T, C> {
    pub(crate) fn new(inner: fjall::Iter, reversed: bool, codec: &'a C) -> Self {
        Self {
            inner,
            codec,
            reversed,
            failed: false,
            _marker: PhantomData,
        }
    }
}

impl<K: Key, T, C: Codec<T>> TryLsmIter<'_, K, T, C> {
    fn decode_guard(&self, g: fjall::Guard) -> DbResult<(K, T)> {
        let (uk, uv) = map_fjall(g.into_inner())?;
        if uk.len() != size_of::<K>() {
            return Err(DbError::CorruptedEntry { offset: 0 });
        }
        let key = K::from_bytes(uk.as_ref());
        let value = self.codec.decode_from(uv.as_ref())?;
        Ok((key, value))
    }
}

impl<K: Key, T, C: Codec<T>> Iterator for TryLsmIter<'_, K, T, C> {
    type Item = DbResult<(K, T)>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.failed {
            return None;
        }
        let raw = if self.reversed {
            self.inner.next_back()?
        } else {
            self.inner.next()?
        };
        let result = self.decode_guard(raw);
        if result.is_err() {
            self.failed = true;
        }
        Some(result)
    }
}