armdb 0.11.2

use std::collections::{BTreeSet, HashMap, VecDeque};
use std::hash::{Hash, Hasher};
use std::mem::size_of;
use std::num::NonZeroU32;
use std::ops::{Bound, Deref};

use crate::error::{DbError, DbResult};
use crate::key::Key;

/// Internal batch size for unbounded `MapIter` streams.
const DEFAULT_BATCH: usize = 64;

/// Iteration direction for ordered scans.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Dir {
    Asc,
    Desc,
}

pub(crate) fn dir_from(reversed: bool) -> Dir {
    if reversed { Dir::Desc } else { Dir::Asc }
}

impl Dir {
    pub(crate) fn flip(self) -> Dir {
        match self {
            Dir::Asc => Dir::Desc,
            Dir::Desc => Dir::Asc,
        }
    }
}

/// Per-map hook the k-way merge drives. Each Map implements it over its shards.
/// Not re-exported from the crate root — use `iter_view()` on a map collection.
pub(crate) trait ShardSource {
    type Key: Key;
    type Val;
    fn shard_count(&self) -> usize;
    /// Up to `n` keys in shard `shard`, strictly after `after` (or from the
    /// edge if `None`) in direction `dir`, paired with their values.
    /// Keys are constrained to `[range_start, range_end)` per `Bound` semantics.
    fn next_batch(
        &self,
        shard: usize,
        after: Option<&Self::Key>,
        dir: Dir,
        n: usize,
        range_start: Bound<OrdKey<Self::Key>>,
        range_end: Bound<OrdKey<Self::Key>>,
    ) -> Vec<(Self::Key, Self::Val)>;

    /// Keys-only batch: up to `n` keys in shard `shard` strictly after `after`,
    /// constrained to `[range_start, range_end)`, **without** materializing any
    /// value. The default derives keys from [`next_batch`], but every Map should
    /// override this to read straight from the companion `BTreeSet` — crucial for
    /// `VarMap`/`VarTypedMap`, whose `next_batch` performs the off-disk value read.
    fn next_keys_batch(
        &self,
        shard: usize,
        after: Option<&Self::Key>,
        dir: Dir,
        n: usize,
        range_start: Bound<OrdKey<Self::Key>>,
        range_end: Bound<OrdKey<Self::Key>>,
    ) -> Vec<Self::Key> {
        self.next_batch(shard, after, dir, n, range_start, range_end)
            .into_iter()
            .map(|(k, _)| k)
            .collect()
    }

    /// Fallible batch read for strict inspection iterators. Inline map sources
    /// delegate to [`next_batch`]; `VarMap` overrides with strict read semantics.
    fn try_next_batch(
        &self,
        shard: usize,
        after: Option<&Self::Key>,
        dir: Dir,
        n: usize,
        range_start: Bound<OrdKey<Self::Key>>,
        range_end: Bound<OrdKey<Self::Key>>,
    ) -> DbResult<Vec<(Self::Key, Self::Val)>> {
        Ok(self.next_batch(shard, after, dir, n, range_start, range_end))
    }
}

/// Per-shard batch size for read-only inspection ordered scans.
///
/// Derived from the requested page limit and shard count so a single page does
/// not over-fetch `DEFAULT_BATCH` rows from every shard.
pub(crate) fn inspection_batch(limit: NonZeroU32, shard_count: usize) -> usize {
    debug_assert!(shard_count > 0, "every opened map has at least one shard");
    let limit = u64::from(limit.get());
    let batch = (limit + 1).div_ceil(shard_count as u64) as usize;
    batch.max(1)
}

/// One page of key/value entries returned by `paginate`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Page<K, V> {
    /// The entries on this page, in the collection's iteration order.
    pub items: Vec<(K, V)>,
    /// Cursor for the next page: pass it as `after` to continue. `Some` only when
    /// the page was filled to `limit` (a further page may exist); `None` at the end.
    pub next: Option<K>,
}

/// Keys-only pagination result, returned by `keys_paginate`. Carries no values,
/// so it is cheap to produce even for `VarMap`/`VarTypedMap` (no off-disk read).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KeyPage<K> {
    pub keys: Vec<K>,
    pub next: Option<K>,
}

/// View over an iterable map collection; exposes ordered iteration.
pub(crate) struct MapIterView<'a, S: ShardSource> {
    src: &'a S,
    dir: Dir,
    batch: usize,
}

impl<'a, S: ShardSource> MapIterView<'a, S> {
    pub(crate) fn new(src: &'a S, dir: Dir) -> Self {
        Self::with_batch(src, dir, DEFAULT_BATCH)
    }

    pub(crate) fn with_batch(src: &'a S, dir: Dir, batch: usize) -> Self {
        debug_assert!(batch > 0);
        Self { src, dir, batch }
    }

    /// Full ordered scan in the collection's configured direction.
    pub fn iter(&self) -> MapIter<'a, S> {
        MapIter::new(self.src, self.dir, self.batch)
    }

    /// Iterate entries in `[start, end)` — start inclusive, end exclusive.
    pub fn range(&self, start: &S::Key, end: &S::Key) -> MapIter<'a, S> {
        self.range_bounds(Bound::Included(start), Bound::Excluded(end))
    }

    /// Iterate entries in range defined by `start` and `end` bounds.
    pub fn range_bounds(&self, start: Bound<&S::Key>, end: Bound<&S::Key>) -> MapIter<'a, S> {
        MapIter::bounded(
            self.src,
            self.dir,
            self.batch,
            bound_ordkey(start),
            bound_ordkey(end),
        )
    }

    /// Strict bounded range iterator for inspection — surfaces storage errors.
    pub(crate) fn try_range_bounds(
        &self,
        start: Bound<&S::Key>,
        end: Bound<&S::Key>,
    ) -> TryMapIter<'a, S> {
        TryMapIter::bounded(
            self.src,
            self.dir,
            self.batch,
            bound_ordkey(start),
            bound_ordkey(end),
        )
    }

    /// Iterate entries whose keys start with `prefix`.
    pub fn prefix_iter(&self, prefix: &[u8]) -> MapIter<'a, S> {
        if prefix.len() > size_of::<S::Key>() {
            return MapIter::empty(self.src, self.dir, self.batch);
        }
        let mut search = S::Key::zeroed();
        search.as_bytes_mut()[..prefix.len()].copy_from_slice(prefix);
        let end = crate::scope::key_math::prefix_to_end_bound::<S::Key>(prefix);
        let end_ref = match &end {
            Bound::Excluded(k) => Bound::Excluded(k),
            Bound::Unbounded => Bound::Unbounded,
            Bound::Included(k) => Bound::Included(k),
        };
        self.range_bounds(Bound::Included(&search), end_ref)
    }

    /// Keyset pagination in the collection's configured direction.
    ///
    /// `after` is exclusive — the next page starts strictly after this key.
    /// `next` is set to the last key in the page when `items.len() == limit`.
    pub fn paginate(&self, after: Option<&S::Key>, limit: usize) -> Page<S::Key, S::Val> {
        let iter = match (after, self.dir) {
            (None, _) => self.range_bounds(Bound::Unbounded, Bound::Unbounded),
            (Some(k), Dir::Asc) => self.range_bounds(Bound::Excluded(k), Bound::Unbounded),
            (Some(k), Dir::Desc) => self.range_bounds(Bound::Unbounded, Bound::Excluded(k)),
        };
        let items: Vec<_> = iter.take(limit).collect();
        let next = if items.len() == limit {
            items.last().map(|(k, _)| *k)
        } else {
            None
        };
        Page { items, next }
    }

    pub(crate) fn src(&self) -> &'a S {
        self.src
    }

    // -- Keys-only iteration -------------------------------------------------
    //
    // These mirror the value methods above but drive `next_keys_batch`, so no
    // value is materialized. For `VarMap`/`VarTypedMap` this avoids the off-disk
    // read entirely; for inline maps it skips the per-key `HashMap` lookup.

    /// Full ordered key scan in the collection's configured direction.
    pub fn keys(&self) -> MapKeyIter<'a, S> {
        MapKeyIter::new(self.src, self.dir, self.batch)
    }

    /// Iterate keys in `[start, end)` — start inclusive, end exclusive.
    pub fn keys_range(&self, start: &S::Key, end: &S::Key) -> MapKeyIter<'a, S> {
        self.keys_range_bounds(Bound::Included(start), Bound::Excluded(end))
    }

    /// Iterate keys in range defined by `start` and `end` bounds.
    pub fn keys_range_bounds(
        &self,
        start: Bound<&S::Key>,
        end: Bound<&S::Key>,
    ) -> MapKeyIter<'a, S> {
        MapKeyIter::bounded(
            self.src,
            self.dir,
            self.batch,
            bound_ordkey(start),
            bound_ordkey(end),
        )
    }

    /// Iterate keys that start with `prefix`.
    pub fn keys_prefix(&self, prefix: &[u8]) -> MapKeyIter<'a, S> {
        if prefix.len() > size_of::<S::Key>() {
            return MapKeyIter::empty(self.src, self.dir, self.batch);
        }
        let mut search = S::Key::zeroed();
        search.as_bytes_mut()[..prefix.len()].copy_from_slice(prefix);
        let end = crate::scope::key_math::prefix_to_end_bound::<S::Key>(prefix);
        let end_ref = match &end {
            Bound::Excluded(k) => Bound::Excluded(k),
            Bound::Unbounded => Bound::Unbounded,
            Bound::Included(k) => Bound::Included(k),
        };
        self.keys_range_bounds(Bound::Included(&search), end_ref)
    }

    /// Keyset pagination yielding keys only.
    ///
    /// `after` is exclusive — the next page starts strictly after this key.
    /// `next` is set to the last key in the page when `keys.len() == limit`.
    pub fn keys_paginate(&self, after: Option<&S::Key>, limit: usize) -> KeyPage<S::Key> {
        let iter = match (after, self.dir) {
            (None, _) => self.keys_range_bounds(Bound::Unbounded, Bound::Unbounded),
            (Some(k), Dir::Asc) => self.keys_range_bounds(Bound::Excluded(k), Bound::Unbounded),
            (Some(k), Dir::Desc) => self.keys_range_bounds(Bound::Unbounded, Bound::Excluded(k)),
        };
        let keys: Vec<_> = iter.take(limit).collect();
        let next = if keys.len() == limit {
            keys.last().copied()
        } else {
            None
        };
        KeyPage { keys, next }
    }
}

/// Batched k-way merge iterator across map shards.
pub(crate) struct MapIter<'a, S: ShardSource> {
    src: &'a S,
    dir: Dir,
    buf: Vec<VecDeque<(S::Key, S::Val)>>,
    last: Vec<Option<S::Key>>,
    exhausted: Vec<bool>,
    batch: usize,
    range_start: Bound<OrdKey<S::Key>>,
    range_end: Bound<OrdKey<S::Key>>,
    done: bool,
}

impl<'a, S: ShardSource> MapIter<'a, S> {
    fn new(src: &'a S, dir: Dir, batch: usize) -> Self {
        Self::bounded(src, dir, batch, Bound::Unbounded, Bound::Unbounded)
    }

    fn bounded(
        src: &'a S,
        dir: Dir,
        batch: usize,
        range_start: Bound<OrdKey<S::Key>>,
        range_end: Bound<OrdKey<S::Key>>,
    ) -> Self {
        let n = src.shard_count();
        Self {
            src,
            dir,
            buf: (0..n).map(|_| VecDeque::new()).collect(),
            last: vec![None; n],
            exhausted: vec![false; n],
            batch,
            range_start,
            range_end,
            done: false,
        }
    }

    fn empty(src: &'a S, dir: Dir, batch: usize) -> Self {
        let n = src.shard_count();
        Self {
            src,
            dir,
            buf: (0..n).map(|_| VecDeque::new()).collect(),
            last: vec![None; n],
            exhausted: vec![true; n],
            batch,
            range_start: Bound::Unbounded,
            range_end: Bound::Unbounded,
            done: true,
        }
    }

    fn refill(&mut self, shard: usize) {
        if self.exhausted[shard] {
            return;
        }
        let batch = self.src.next_batch(
            shard,
            self.last[shard].as_ref(),
            self.dir,
            self.batch,
            self.range_start,
            self.range_end,
        );
        if batch.is_empty() {
            self.exhausted[shard] = true;
        } else {
            self.buf[shard].extend(batch);
        }
    }
}

/// Batched k-way merge iterator that surfaces the first storage error from
/// [`ShardSource::try_next_batch`] instead of treating it as shard exhaustion.
pub(crate) struct TryMapIter<'a, S: ShardSource> {
    src: &'a S,
    dir: Dir,
    buf: Vec<VecDeque<(S::Key, S::Val)>>,
    last: Vec<Option<S::Key>>,
    exhausted: Vec<bool>,
    batch: usize,
    range_start: Bound<OrdKey<S::Key>>,
    range_end: Bound<OrdKey<S::Key>>,
    done: bool,
    failed: bool,
}

impl<'a, S: ShardSource> TryMapIter<'a, S> {
    fn bounded(
        src: &'a S,
        dir: Dir,
        batch: usize,
        range_start: Bound<OrdKey<S::Key>>,
        range_end: Bound<OrdKey<S::Key>>,
    ) -> Self {
        let n = src.shard_count();
        Self {
            src,
            dir,
            buf: (0..n).map(|_| VecDeque::new()).collect(),
            last: vec![None; n],
            exhausted: vec![false; n],
            batch,
            range_start,
            range_end,
            done: false,
            failed: false,
        }
    }

    fn refill(&mut self, shard: usize) -> Result<(), DbError> {
        if self.exhausted[shard] {
            return Ok(());
        }
        match self.src.try_next_batch(
            shard,
            self.last[shard].as_ref(),
            self.dir,
            self.batch,
            self.range_start,
            self.range_end,
        ) {
            Ok(batch) => {
                if batch.is_empty() {
                    self.exhausted[shard] = true;
                } else {
                    self.buf[shard].extend(batch);
                }
                Ok(())
            }
            Err(error) => {
                for buf in &mut self.buf {
                    buf.clear();
                }
                self.done = true;
                for ex in &mut self.exhausted {
                    *ex = true;
                }
                Err(error)
            }
        }
    }
}

impl<S: ShardSource> Iterator for TryMapIter<'_, S> {
    type Item = DbResult<(S::Key, S::Val)>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.failed {
            return None;
        }
        if self.done {
            return None;
        }

        for i in 0..self.buf.len() {
            if self.buf[i].is_empty() && !self.exhausted[i] {
                match self.refill(i) {
                    Ok(()) => {}
                    Err(error) => {
                        self.failed = true;
                        return Some(Err(error));
                    }
                }
            }
        }

        let mut best: Option<(usize, OrdKey<S::Key>)> = None;
        for (i, buf) in self.buf.iter().enumerate() {
            if let Some((k, _)) = buf.front() {
                let ok = OrdKey(*k);
                let better = match best {
                    None => true,
                    Some((_, ref b)) => match self.dir {
                        Dir::Asc => ok < *b,
                        Dir::Desc => ok > *b,
                    },
                };
                if better {
                    best = Some((i, ok));
                }
            }
        }

        let (shard, _) = best?;
        let (k, v) = self.buf[shard]
            .pop_front()
            .expect("shard buffer non-empty when selected by merge");
        self.last[shard] = Some(k);
        Some(Ok((k, v)))
    }
}

impl<S: ShardSource> Iterator for MapIter<'_, S> {
    type Item = (S::Key, S::Val);

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }
        for i in 0..self.buf.len() {
            if self.buf[i].is_empty() && !self.exhausted[i] {
                self.refill(i);
            }
        }

        let mut best: Option<(usize, OrdKey<S::Key>)> = None;
        for (i, buf) in self.buf.iter().enumerate() {
            if let Some((k, _)) = buf.front() {
                let ok = OrdKey(*k);
                let better = match best {
                    None => true,
                    Some((_, ref b)) => match self.dir {
                        Dir::Asc => ok < *b,
                        Dir::Desc => ok > *b,
                    },
                };
                if better {
                    best = Some((i, ok));
                }
            }
        }

        let (shard, _) = best?;
        let (k, v) = self.buf[shard]
            .pop_front()
            .expect("shard buffer non-empty when selected by merge");
        self.last[shard] = Some(k);
        Some((k, v))
    }
}

/// Batched k-way merge iterator yielding **keys only** across map shards.
///
/// Structurally identical to [`MapIter`] but drives [`ShardSource::next_keys_batch`],
/// so it never materializes a value (no off-disk read for Var collections, no
/// per-key `HashMap` lookup for inline maps).
pub(crate) struct MapKeyIter<'a, S: ShardSource> {
    src: &'a S,
    dir: Dir,
    buf: Vec<VecDeque<S::Key>>,
    last: Vec<Option<S::Key>>,
    exhausted: Vec<bool>,
    batch: usize,
    range_start: Bound<OrdKey<S::Key>>,
    range_end: Bound<OrdKey<S::Key>>,
    done: bool,
}

impl<'a, S: ShardSource> MapKeyIter<'a, S> {
    fn new(src: &'a S, dir: Dir, batch: usize) -> Self {
        Self::bounded(src, dir, batch, Bound::Unbounded, Bound::Unbounded)
    }

    fn bounded(
        src: &'a S,
        dir: Dir,
        batch: usize,
        range_start: Bound<OrdKey<S::Key>>,
        range_end: Bound<OrdKey<S::Key>>,
    ) -> Self {
        let n = src.shard_count();
        Self {
            src,
            dir,
            buf: (0..n).map(|_| VecDeque::new()).collect(),
            last: vec![None; n],
            exhausted: vec![false; n],
            batch,
            range_start,
            range_end,
            done: false,
        }
    }

    fn empty(src: &'a S, dir: Dir, batch: usize) -> Self {
        let n = src.shard_count();
        Self {
            src,
            dir,
            buf: (0..n).map(|_| VecDeque::new()).collect(),
            last: vec![None; n],
            exhausted: vec![true; n],
            batch,
            range_start: Bound::Unbounded,
            range_end: Bound::Unbounded,
            done: true,
        }
    }

    fn refill(&mut self, shard: usize) {
        if self.exhausted[shard] {
            return;
        }
        let batch = self.src.next_keys_batch(
            shard,
            self.last[shard].as_ref(),
            self.dir,
            self.batch,
            self.range_start,
            self.range_end,
        );
        if batch.is_empty() {
            self.exhausted[shard] = true;
        } else {
            self.buf[shard].extend(batch);
        }
    }
}

impl<S: ShardSource> Iterator for MapKeyIter<'_, S> {
    type Item = S::Key;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }
        for i in 0..self.buf.len() {
            if self.buf[i].is_empty() && !self.exhausted[i] {
                self.refill(i);
            }
        }

        let mut best: Option<(usize, OrdKey<S::Key>)> = None;
        for (i, buf) in self.buf.iter().enumerate() {
            if let Some(k) = buf.front() {
                let ok = OrdKey(*k);
                let better = match best {
                    None => true,
                    Some((_, ref b)) => match self.dir {
                        Dir::Asc => ok < *b,
                        Dir::Desc => ok > *b,
                    },
                };
                if better {
                    best = Some((i, ok));
                }
            }
        }

        let (shard, _) = best?;
        let k = self.buf[shard]
            .pop_front()
            .expect("shard buffer non-empty when selected by merge");
        self.last[shard] = Some(k);
        Some(k)
    }
}

/// Collect up to `n` keys from `ki` strictly after `after` in direction `dir`,
/// constrained to `[range_start, range_end)`.
pub(crate) fn key_index_batch<K: Key + Hash + Eq>(
    ki: &BTreeSet<OrdKey<K>>,
    after: Option<&K>,
    dir: Dir,
    n: usize,
    range_start: Bound<OrdKey<K>>,
    range_end: Bound<OrdKey<K>>,
) -> Vec<K> {
    let mut out = Vec::with_capacity(n.min(ki.len()));
    match dir {
        Dir::Asc => {
            let lower = merge_lower(range_start, after);
            // `BTreeSet::range` panics on a degenerate range (start > end, or
            // start == end with both bounds excluded); such a range is empty.
            if range_is_degenerate(&lower, &range_end) {
                return out;
            }
            let iter = ki.range((lower, range_end));
            for ok in iter.take(n) {
                out.push(ok.0);
            }
        }
        Dir::Desc => {
            let upper = merge_upper(range_end, after);
            if range_is_degenerate(&range_start, &upper) {
                return out;
            }
            let iter = ki.range((range_start, upper)).rev();
            for ok in iter.take(n) {
                out.push(ok.0);
            }
        }
    }
    out
}

/// Whether `(start, end)` is a range `BTreeSet::range` would reject by panicking:
/// `start > end`, or `start == end` with both bounds excluded. Both denote an
/// empty range, so callers return no keys instead of panicking.
pub(crate) fn range_is_degenerate<K: Key>(
    start: &Bound<OrdKey<K>>,
    end: &Bound<OrdKey<K>>,
) -> bool {
    let s = match start {
        Bound::Included(k) | Bound::Excluded(k) => k,
        Bound::Unbounded => return false,
    };
    let e = match end {
        Bound::Included(k) | Bound::Excluded(k) => k,
        Bound::Unbounded => return false,
    };
    match s.cmp(e) {
        std::cmp::Ordering::Greater => true,
        std::cmp::Ordering::Equal => {
            matches!(start, Bound::Excluded(_)) && matches!(end, Bound::Excluded(_))
        }
        std::cmp::Ordering::Less => false,
    }
}

pub(crate) fn merge_lower<K: Key>(
    range_start: Bound<OrdKey<K>>,
    after: Option<&K>,
) -> Bound<OrdKey<K>> {
    match after {
        None => range_start,
        Some(a) => max_bound(range_start, Bound::Excluded(OrdKey(*a))),
    }
}

pub(crate) fn merge_upper<K: Key>(
    range_end: Bound<OrdKey<K>>,
    after: Option<&K>,
) -> Bound<OrdKey<K>> {
    match after {
        None => range_end,
        Some(a) => min_bound(range_end, Bound::Excluded(OrdKey(*a))),
    }
}

pub(crate) fn max_bound<K: Key>(a: Bound<OrdKey<K>>, b: Bound<OrdKey<K>>) -> Bound<OrdKey<K>> {
    match (a, b) {
        (Bound::Unbounded, x) | (x, Bound::Unbounded) => x,
        (Bound::Included(x), Bound::Included(y)) => Bound::Included(if x >= y { x } else { y }),
        (Bound::Excluded(x), Bound::Excluded(y)) => Bound::Excluded(if x >= y { x } else { y }),
        (Bound::Included(x), Bound::Excluded(y)) => {
            if x > y {
                Bound::Included(x)
            } else {
                Bound::Excluded(y)
            }
        }
        (Bound::Excluded(x), Bound::Included(y)) => {
            if x >= y {
                Bound::Excluded(x)
            } else {
                Bound::Included(y)
            }
        }
    }
}

pub(crate) fn min_bound<K: Key>(a: Bound<OrdKey<K>>, b: Bound<OrdKey<K>>) -> Bound<OrdKey<K>> {
    match (a, b) {
        (Bound::Unbounded, x) | (x, Bound::Unbounded) => x,
        (Bound::Included(x), Bound::Included(y)) => Bound::Included(if x <= y { x } else { y }),
        (Bound::Excluded(x), Bound::Excluded(y)) => Bound::Excluded(if x <= y { x } else { y }),
        (Bound::Included(x), Bound::Excluded(y)) => {
            if x < y {
                Bound::Included(x)
            } else {
                Bound::Excluded(y)
            }
        }
        (Bound::Excluded(x), Bound::Included(y)) => {
            if x <= y {
                Bound::Excluded(x)
            } else {
                Bound::Included(y)
            }
        }
    }
}

pub(crate) fn bound_ordkey<K: Key>(b: Bound<&K>) -> Bound<OrdKey<K>> {
    match b {
        Bound::Included(k) => Bound::Included(OrdKey(*k)),
        Bound::Excluded(k) => Bound::Excluded(OrdKey(*k)),
        Bound::Unbounded => Bound::Unbounded,
    }
}

/// A key wrapper whose ordering is byte-lexicographic over `K::as_bytes()` —
/// the same order Tree/SkipList collections use, so "sorted by key" means the
/// same thing across Map and Tree.
#[derive(Clone, Copy, Debug)]
pub(crate) struct OrdKey<K>(pub K);

impl<K: Key> PartialEq for OrdKey<K> {
    fn eq(&self, other: &Self) -> bool {
        self.0.as_bytes() == other.0.as_bytes()
    }
}
impl<K: Key> Eq for OrdKey<K> {}
impl<K: Key> PartialOrd for OrdKey<K> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}
impl<K: Key> Ord for OrdKey<K> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.0.as_bytes().cmp(other.0.as_bytes())
    }
}
impl<K: Key> Hash for OrdKey<K> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.as_bytes().hash(state);
    }
}

/// Per-shard index: the authoritative `HashMap` plus an optional key-ordered
/// companion (`Some` only when the collection was opened `iterable`). Both live
/// under one shard `Mutex`, so they never diverge.
///
/// Reads go through `Deref<Target = HashMap>`. There is intentionally **no**
/// `DerefMut`: every keyset mutation must use [`upsert`]/[`remove`], which keep
/// the companion in sync (R2). Value-only mutation uses [`get_mut`].
pub(crate) struct ShardIndex<K: Key + Hash + Eq, E> {
    map: HashMap<K, E>,
    key_index: Option<BTreeSet<OrdKey<K>>>,
}

impl<K: Key + Hash + Eq, E> ShardIndex<K, E> {
    pub(crate) fn new(iterable: bool) -> Self {
        Self {
            map: HashMap::new(),
            key_index: iterable.then(BTreeSet::new),
        }
    }

    #[allow(dead_code)] // used by map collections in follow-up tasks
    pub(crate) fn is_iterable(&self) -> bool {
        self.key_index.is_some()
    }

    /// Insert or overwrite. Companion is updated only when the key is newly
    /// added (an overwrite leaves the keyset unchanged).
    pub(crate) fn upsert(&mut self, key: K, entry: E) -> Option<E> {
        let prev = self.map.insert(key, entry);
        if prev.is_none()
            && let Some(ki) = self.key_index.as_mut()
        {
            ki.insert(OrdKey(key));
        }
        prev
    }

    pub(crate) fn remove(&mut self, key: &K) -> Option<E> {
        let prev = self.map.remove(key);
        if prev.is_some()
            && let Some(ki) = self.key_index.as_mut()
        {
            ki.remove(&OrdKey(*key));
        }
        prev
    }

    /// Value-only mutation — keyset unchanged, companion untouched.
    pub(crate) fn get_mut(&mut self, key: &K) -> Option<&mut E> {
        self.map.get_mut(key)
    }

    /// Test/recovery helper: turn the companion on and fill it from the map.
    pub(crate) fn set_iterable_and_rebuild(&mut self) {
        let ki = self.key_index.get_or_insert_with(BTreeSet::new);
        ki.clear();
        ki.extend(self.map.keys().map(|k| OrdKey(*k)));
    }

    /// Snapshot of the companion in `[start..)` direction-agnostic order.
    /// Used by `MapIter`. Returns up to `limit` keys ≥ `start` (ascending).
    pub(crate) fn key_index(&self) -> Option<&BTreeSet<OrdKey<K>>> {
        self.key_index.as_ref()
    }
}

impl<K: Key + Hash + Eq, E> Deref for ShardIndex<K, E> {
    type Target = HashMap<K, E>;
    fn deref(&self) -> &HashMap<K, E> {
        &self.map
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::DbResult;
    use std::cell::RefCell;
    use std::io;
    use std::num::NonZeroU32;
    use std::sync::atomic::{AtomicUsize, Ordering};

    #[test]
    fn ordkey_is_byte_lexicographic() {
        let mut v = [OrdKey([2u8, 0]), OrdKey([1u8, 9]), OrdKey([1u8, 0])];
        v.sort();
        assert_eq!(v, [OrdKey([1u8, 0]), OrdKey([1u8, 9]), OrdKey([2u8, 0])]);
    }

    #[test]
    fn upsert_remove_maintain_companion_when_iterable() {
        let mut si: ShardIndex<[u8; 2], u32> = ShardIndex::new(true);
        assert_eq!(si.upsert([1, 0], 10), None);
        assert_eq!(si.upsert([1, 0], 11), Some(10)); // overwrite: keyset unchanged
        si.upsert([2, 0], 20);
        let keys: Vec<_> = si.key_index().unwrap().iter().map(|k| k.0).collect();
        assert_eq!(keys, [[1, 0], [2, 0]]);
        si.remove(&[1, 0]);
        let keys: Vec<_> = si.key_index().unwrap().iter().map(|k| k.0).collect();
        assert_eq!(keys, [[2, 0]]);
    }

    #[test]
    fn companion_absent_when_not_iterable() {
        let mut si: ShardIndex<[u8; 2], u32> = ShardIndex::new(false);
        si.upsert([1, 0], 10);
        assert!(si.key_index().is_none());
        assert_eq!(*si.get(&[1, 0]).unwrap(), 10); // Deref read still works
    }

    #[test]
    fn rebuild_fills_from_map() {
        let mut si: ShardIndex<[u8; 2], u32> = ShardIndex::new(false);
        si.upsert([2, 0], 20);
        si.upsert([1, 0], 10);
        si.set_iterable_and_rebuild();
        let keys: Vec<_> = si.key_index().unwrap().iter().map(|k| k.0).collect();
        assert_eq!(keys, [[1, 0], [2, 0]]);
    }

    #[test]
    fn dir_flip_is_involutive() {
        assert_eq!(Dir::Asc.flip(), Dir::Desc);
        assert_eq!(Dir::Desc.flip(), Dir::Asc);
        for dir in [Dir::Asc, Dir::Desc] {
            assert_eq!(dir.flip().flip(), dir);
        }
    }

    /// Test double with two shards; `try_next_batch` fails on the second refill of shard 0.
    struct FaultSource {
        shards: Vec<Vec<[u8; 2]>>,
        refill_count: RefCell<Vec<usize>>,
        fail_on: Option<(usize, usize)>,
        recorded_batches: RefCell<Vec<usize>>,
        total_rows_returned: AtomicUsize,
    }

    impl FaultSource {
        fn key(shard: usize, idx: u16) -> [u8; 2] {
            [shard as u8, idx as u8]
        }

        fn with_error_after_first_batch() -> Self {
            Self {
                shards: vec![
                    (0..5).map(|i| Self::key(0, i)).collect(),
                    (0..5).map(|i| Self::key(1, i)).collect(),
                ],
                refill_count: RefCell::new(vec![0, 0]),
                fail_on: Some((0, 2)),
                recorded_batches: RefCell::new(Vec::new()),
                total_rows_returned: AtomicUsize::new(0),
            }
        }

        fn with_many_shards(shard_count: usize, keys_per_shard: usize) -> Self {
            let shards = (0..shard_count)
                .map(|s| {
                    (0..keys_per_shard)
                        .map(|i| Self::key(s, i as u16))
                        .collect()
                })
                .collect();
            Self {
                shards,
                refill_count: RefCell::new(vec![0; shard_count]),
                fail_on: None,
                recorded_batches: RefCell::new(Vec::new()),
                total_rows_returned: AtomicUsize::new(0),
            }
        }

        fn keys_in_range(
            &self,
            shard: usize,
            after: Option<&[u8; 2]>,
            dir: Dir,
            n: usize,
        ) -> Vec<([u8; 2], [u8; 2])> {
            let keys = &self.shards[shard];
            let start = match (after, dir) {
                (None, _) => 0,
                (Some(a), Dir::Asc) => keys.partition_point(|k| k <= a),
                (Some(a), Dir::Desc) => keys.partition_point(|k| k >= a),
            };
            let slice: Vec<[u8; 2]> = match dir {
                Dir::Asc => keys.iter().copied().skip(start).take(n).collect(),
                Dir::Desc => keys.iter().copied().rev().skip(start).take(n).collect(),
            };
            slice.into_iter().map(|k| (k, k)).collect()
        }
    }

    impl ShardSource for FaultSource {
        type Key = [u8; 2];
        type Val = [u8; 2];

        fn shard_count(&self) -> usize {
            self.shards.len()
        }

        fn next_batch(
            &self,
            shard: usize,
            after: Option<&Self::Key>,
            dir: Dir,
            n: usize,
            _range_start: Bound<OrdKey<Self::Key>>,
            _range_end: Bound<OrdKey<Self::Key>>,
        ) -> Vec<(Self::Key, Self::Val)> {
            self.keys_in_range(shard, after, dir, n)
        }

        fn try_next_batch(
            &self,
            shard: usize,
            after: Option<&Self::Key>,
            dir: Dir,
            n: usize,
            range_start: Bound<OrdKey<Self::Key>>,
            range_end: Bound<OrdKey<Self::Key>>,
        ) -> DbResult<Vec<(Self::Key, Self::Val)>> {
            self.recorded_batches.borrow_mut().push(n);
            let mut counts = self.refill_count.borrow_mut();
            counts[shard] += 1;
            let refill = counts[shard];
            drop(counts);
            if self.fail_on == Some((shard, refill)) {
                return Err(DbError::Io(io::Error::new(
                    io::ErrorKind::Other,
                    "fault injection",
                )));
            }
            let batch = self.next_batch(shard, after, dir, n, range_start, range_end);
            self.total_rows_returned
                .fetch_add(batch.len(), Ordering::Relaxed);
            Ok(batch)
        }
    }

    #[test]
    fn try_map_iter_returns_error_instead_of_exhausting_faulted_shard() {
        let source = FaultSource::with_error_after_first_batch();
        let rows = MapIterView::with_batch(&source, Dir::Asc, 1)
            .try_range_bounds(Bound::Unbounded, Bound::Unbounded)
            .collect::<DbResult<Vec<_>>>();
        assert!(matches!(rows, Err(DbError::Io(_))));
    }

    #[test]
    fn map_iter_uses_infallible_next_batch_despite_fault_source() {
        let source = FaultSource::with_error_after_first_batch();
        let rows: Vec<_> = MapIterView::with_batch(&source, Dir::Asc, 1)
            .range_bounds(Bound::Unbounded, Bound::Unbounded)
            .collect();
        assert_eq!(rows.len(), 10);
    }

    #[test]
    fn inspection_batch_is_limit_and_shard_derived() {
        assert_eq!(inspection_batch(NonZeroU32::new(1).unwrap(), 16), 1);
        assert_eq!(inspection_batch(NonZeroU32::new(10).unwrap(), 16), 1);
        assert_eq!(inspection_batch(NonZeroU32::new(64).unwrap(), 4), 17);
        assert!(inspection_batch(NonZeroU32::new(u32::MAX).unwrap(), 16) > 0);
    }

    #[test]
    fn inspection_batch_limits_per_shard_fetch_size() {
        let shard_count = 16;
        let limit = NonZeroU32::new(10).unwrap();
        let batch = inspection_batch(limit, shard_count);
        assert_eq!(batch, 1);

        let source = FaultSource::with_many_shards(shard_count, 3);
        let collected: DbResult<Vec<_>> = MapIterView::with_batch(&source, Dir::Asc, batch)
            .try_range_bounds(Bound::Unbounded, Bound::Unbounded)
            .take(limit.get() as usize)
            .collect();
        assert!(collected.is_ok());
        assert!(
            source.recorded_batches.borrow().iter().all(|n| *n == batch),
            "every refill must use the derived batch size"
        );
        let fetched = source.total_rows_returned.load(Ordering::Relaxed);
        let max_expected = limit.get() as usize + shard_count + shard_count;
        assert!(
            fetched <= max_expected,
            "fetched {fetched} rows, expected O(limit + shard_count)"
        );
    }
}