armdb 0.5.3

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

use crate::key::Key;

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

/// Iteration direction for map ordered scans (internal; not re-exported).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Dir {
    Asc,
    Desc,
}

pub(crate) fn dir_from(reversed: bool) -> Dir {
    if reversed { Dir::Desc } else { 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()
    }
}

/// 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 {
            src,
            dir,
            batch: DEFAULT_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),
        )
    }

    /// 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 = 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 = 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);
        }
    }
}

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.
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,
    }
}

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))),
    }
}

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))),
    }
}

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)
            }
        }
    }
}

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)
            }
        }
    }
}

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,
    }
}

fn prefix_to_end_bound<K: Key>(prefix: &[u8]) -> Bound<K> {
    let mut incremented = prefix.to_vec();
    let mut carry = true;
    for byte in incremented.iter_mut().rev() {
        if carry {
            if *byte == 0xFF {
                *byte = 0x00;
            } else {
                *byte += 1;
                carry = false;
                break;
            }
        }
    }
    if carry {
        Bound::Unbounded
    } else {
        let mut end = K::zeroed();
        end.as_bytes_mut()[..incremented.len()].copy_from_slice(&incremented);
        Bound::Excluded(end)
    }
}

/// 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::*;

    #[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]]);
    }
}