armdb 0.14.2

//! Field-boundary key math for typed scopes.
use std::mem::size_of;
use std::ops::Bound;

use armour_core::{KeyScheme, KeyType};

use crate::key::Key;

pub(crate) const fn typed_fields(scheme: KeyScheme) -> &'static [KeyType] {
    match scheme {
        KeyScheme::Typed(fields) => fields,
        KeyScheme::Bytes => panic!("scope::<N> requires KeyScheme::Typed"),
    }
}

pub(crate) const fn field_offset(fields: &[KeyType], n: usize) -> usize {
    assert!(
        n <= fields.len(),
        "scope/range field count exceeds KEY_SCHEME fields"
    );
    let mut off = 0usize;
    let mut i = 0usize;
    while i < n {
        off += fields[i].byte_len();
        i += 1;
    }
    off
}

pub(crate) const fn assert_key_width<K: Key>() {
    let scheme = K::KEY_SCHEME;
    match scheme.static_size() {
        Some(w) => assert!(
            w == size_of::<K>(),
            "KEY_SCHEME.static_size() != size_of::<K>()"
        ),
        None => panic!("scope::<N> requires KeyScheme::Typed"),
    }
}

pub(crate) const fn assert_scope_n<K: Key>(n: usize) {
    assert_key_width::<K>();
    let fields = typed_fields(K::KEY_SCHEME);
    assert!(
        n <= fields.len(),
        "scope::<N>: N exceeds KEY_SCHEME field count"
    );
}

/// Full `range::<M>` contract: key width, `M <= fields`, and `M > N`.
/// Callable from a `const { … }` block once `N` is a const-generic on [`Scope`](crate::Scope).
pub(crate) const fn assert_range_m<K: Key>(n: usize, m: usize) {
    assert_key_width::<K>();
    let fields = typed_fields(K::KEY_SCHEME);
    assert!(
        m <= fields.len(),
        "range::<M>: M exceeds KEY_SCHEME field count"
    );
    assert!(m > n, "range::<M>: M must be > scope N");
}

/// Exclusive end bound for a byte prefix — the single copy used by every family.
pub(crate) 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)
    }
}

/// Byte bounds for a field prefix. Caller must already have checked `n` via
/// `assert_scope_n` in a `const { … }` block on the const-generic `N`.
pub(crate) fn scope_bounds<K: Key>(prefix_key: &K, n: usize) -> (Bound<K>, Bound<K>) {
    debug_assert!(n <= typed_fields(K::KEY_SCHEME).len());
    let fields = typed_fields(K::KEY_SCHEME);
    let len = field_offset(fields, n);
    let mut start = K::zeroed();
    start.as_bytes_mut()[..len].copy_from_slice(&prefix_key.as_bytes()[..len]);
    let end = prefix_to_end_bound::<K>(&prefix_key.as_bytes()[..len]);
    (Bound::Included(start), end)
}

/// Const-generic entry used by `scope::<N>()`.
pub(crate) fn scope_bounds_n<K: Key, const N: usize>(prefix_key: &K) -> (Bound<K>, Bound<K>) {
    const { assert_scope_n::<K>(N) };
    scope_bounds(prefix_key, N)
}

/// Copy of `key` with only the first `n` typed fields retained; tail zeroed.
pub(crate) fn prefix_key<K: Key>(key: &K, n: usize) -> K {
    let fields = typed_fields(K::KEY_SCHEME);
    let len = field_offset(fields, n);
    let mut out = K::zeroed();
    out.as_bytes_mut()[..len].copy_from_slice(&key.as_bytes()[..len]);
    out
}

pub(crate) fn saturate_range_bounds<K: Key>(
    scope_key: &K,
    n: usize,
    m: usize,
    start: Bound<&K>,
    end: Bound<&K>,
) -> (Bound<K>, Bound<K>) {
    let fields = typed_fields(K::KEY_SCHEME);
    let n_off = field_offset(fields, n);
    let m_off = field_offset(fields, m);
    // Spec: Unbounded means the corresponding *scope* edge, not the whole collection.
    let (scope_start, scope_end) = scope_bounds(scope_key, n);
    let sat = |b: Bound<&K>, is_start: bool| -> Bound<K> {
        match b {
            Bound::Unbounded => {
                if is_start {
                    scope_start
                } else {
                    scope_end
                }
            }
            Bound::Included(k) | Bound::Excluded(k) => {
                let fill = match (&b, is_start) {
                    (Bound::Included(_), true) => 0x00u8,
                    (Bound::Excluded(_), true) => 0xFF,
                    (Bound::Included(_), false) => 0xFF,
                    (Bound::Excluded(_), false) => 0x00,
                    _ => unreachable!(),
                };
                let mut out = K::zeroed();
                out.as_bytes_mut()[..n_off].copy_from_slice(&scope_key.as_bytes()[..n_off]);
                out.as_bytes_mut()[n_off..m_off].copy_from_slice(&k.as_bytes()[n_off..m_off]);
                out.as_bytes_mut()[m_off..].fill(fill);
                match b {
                    Bound::Included(_) => Bound::Included(out),
                    Bound::Excluded(_) => Bound::Excluded(out),
                    Bound::Unbounded => unreachable!(),
                }
            }
        }
    };
    (sat(start, true), sat(end, false))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::impl_key_zerocopy;
    use armour_core::{KeyScheme, KeyType};
    use zerocopy::{FromBytes, Immutable, IntoBytes};

    #[derive(Clone, Copy, Debug, PartialEq, Eq, FromBytes, IntoBytes, Immutable)]
    #[repr(C)]
    struct K3 {
        a: u64,
        b: u64,
        c: u64,
    }
    impl_key_zerocopy!(
        K3,
        KeyScheme::Typed(&[KeyType::U64, KeyType::U64, KeyType::U64])
    );

    #[test]
    fn included_start_fills_tail_zero() {
        let scope = K3 { a: 1, b: 0, c: 0 };
        let bound = K3 { a: 9, b: 5, c: 9 }; // a ignored (N=1), c ignored (tail)
        let (start, _) =
            saturate_range_bounds(&scope, 1, 2, Bound::Included(&bound), Bound::Unbounded);
        match start {
            Bound::Included(k) => {
                assert_eq!(k.a, 1);
                assert_eq!(k.b, 5);
                assert_eq!(k.c, 0);
            }
            _ => panic!("expected Included"),
        }
    }

    #[test]
    fn excluded_start_fills_tail_ff() {
        let scope = K3 { a: 1, b: 0, c: 0 };
        let bound = K3 { a: 0, b: 5, c: 0 };
        let (start, _) =
            saturate_range_bounds(&scope, 1, 2, Bound::Excluded(&bound), Bound::Unbounded);
        match start {
            Bound::Excluded(k) => {
                assert_eq!(k.a, 1);
                assert_eq!(k.b, 5);
                assert_eq!(k.c, u64::MAX);
            }
            _ => panic!("expected Excluded"),
        }
    }

    #[test]
    fn included_end_fills_tail_ff() {
        let scope = K3 { a: 1, b: 0, c: 0 };
        let bound = K3 { a: 0, b: 5, c: 0 };
        let (_, end) =
            saturate_range_bounds(&scope, 1, 2, Bound::Unbounded, Bound::Included(&bound));
        match end {
            Bound::Included(k) => {
                assert_eq!(k.a, 1);
                assert_eq!(k.b, 5);
                assert_eq!(k.c, u64::MAX);
            }
            _ => panic!("expected Included"),
        }
    }

    #[test]
    fn excluded_end_fills_tail_zero() {
        let scope = K3 { a: 1, b: 0, c: 0 };
        let bound = K3 { a: 0, b: 5, c: 0 };
        let (_, end) =
            saturate_range_bounds(&scope, 1, 2, Bound::Unbounded, Bound::Excluded(&bound));
        match end {
            Bound::Excluded(k) => {
                assert_eq!(k.a, 1);
                assert_eq!(k.b, 5);
                assert_eq!(k.c, 0);
            }
            _ => panic!("expected Excluded"),
        }
    }

    #[test]
    fn all_ff_prefix_end_is_unbounded() {
        let key = K3 {
            a: u64::MAX,
            b: 0,
            c: 0,
        };
        let (_start, end) = scope_bounds(&key, 1);
        assert!(matches!(end, Bound::Unbounded));
    }

    #[test]
    fn unbounded_range_bound_is_scope_edge() {
        let scope = K3 { a: 1, b: 0, c: 0 };
        let (scope_start, scope_end) = scope_bounds(&scope, 1);
        let (start, end) = saturate_range_bounds(&scope, 1, 2, Bound::Unbounded, Bound::Unbounded);
        assert_eq!(bound_bytes(&start), bound_bytes(&scope_start));
        assert_eq!(bound_bytes(&end), bound_bytes(&scope_end));
        // Bytes alone would not catch Included <-> Excluded: an `Excluded(scope_start)`
        // silently drops the first key of every scope.
        assert_eq!(
            std::mem::discriminant(&start),
            std::mem::discriminant(&scope_start)
        );
        assert!(matches!(start, Bound::Included(_)));
        assert_eq!(
            std::mem::discriminant(&end),
            std::mem::discriminant(&scope_end)
        );
    }

    fn bound_bytes(b: &Bound<K3>) -> Option<Vec<u8>> {
        match b {
            Bound::Included(k) | Bound::Excluded(k) => Some(Key::as_bytes(k).to_vec()),
            Bound::Unbounded => None,
        }
    }
}