vest_lib 0.2.2

//! Minimal big-endian base-256 integer conversion and formats.
#[cfg(feature = "alloc")]
use alloc::{vec, vec::Vec};
use vstd::arithmetic::{div_mod::*, power::*, power2::*};
use vstd::bits::*;
#[cfg(verus_only)]
use vstd::std_specs::iter::IteratorSpec;
use vstd::{calc, prelude::*};

verus! {

const USIZE_MODULUS_32: u64 = 0x100000000;

const USIZE_MODULUS_64: u128 = 0x10000000000000000u128;

/// Unsigned big-endian base-256 decoding.
pub open spec fn nat_from_be_bytes(bytes: Seq<u8>) -> nat
    decreases bytes.len(),
{
    if bytes.len() == 0 {
        0
    } else {
        nat_from_be_bytes(bytes.drop_last()) * 256 + bytes.last() as nat
    }
}

/// Unsigned big-endian base-256 encoding.
pub open spec fn nat_to_be_bytes(n: nat) -> Seq<u8>
    decreases n,
{
    if n < 256 {
        seq![n as u8]
    } else {
        nat_to_be_bytes((n / 256) as nat).push((n % 256) as u8)
    }
}

/// Number of bytes in `usize`.
pub open spec fn size_of_usize() -> nat {
    if usize::BITS == 32 {
        4
    } else {
        8
    }
}

proof fn lemma_usize_shr8_is_div256(v: usize)
    ensures
        (v >> 8usize) as nat == v as nat / 256,
{
    lemma_usize_shr_is_div(v, 8usize);
    assert(pow2(8) == 256) by (compute_only);
}

proof fn lemma_usize_low8_is_mod256(v: usize)
    ensures
        (v & 0xffusize) as nat == v as nat % 256,
{
    lemma_usize_low_bits_mask_is_mod(v, 8);
    assert(pow2(8) == 256) by (compute_only);
}

proof fn lemma_nat_from_be_bytes_fits_shr8(bytes: Seq<u8>)
    requires
        bytes.len() <= size_of_usize(),
    ensures
        usize::BITS == 32 ==> nat_from_be_bytes(bytes) < USIZE_MODULUS_32 as nat,
        usize::BITS == 64 ==> nat_from_be_bytes(bytes) < USIZE_MODULUS_64 as nat,
{
    lemma_from_be_bytes_upper_bound(bytes);
    assert(usize::BITS == 32 || usize::BITS == 64);
    if usize::BITS == 32 {
        assert(size_of_usize() == 4);
        reveal_with_fuel(pow, 5);
    } else {
        assert(usize::BITS == 64);
        assert(size_of_usize() == 8);
        reveal_with_fuel(pow, 9);
    }
}

proof fn lemma_usize32_shl8_or_is_base256(v: usize, b: u8)
    by (bit_vector)
    requires
        usize::BITS == 32,
    ensures
        (((v << 8usize) | b as usize) as nat) == (v as nat * 256 + b as nat) % (
        USIZE_MODULUS_32 as nat),
{
}

proof fn lemma_usize64_shl8_or_is_base256(v: usize, b: u8)
    by (bit_vector)
    requires
        usize::BITS == 64,
    ensures
        (((v << 8usize) | b as usize) as nat) == (v as nat * 256 + b as nat) % (
        USIZE_MODULUS_64 as nat),
{
}

pub proof fn lemma_nat_from_be_bytes_fits_usize(bytes: Seq<u8>)
    requires
        bytes.len() <= size_of_usize(),
    ensures
        nat_from_be_bytes(bytes) <= usize::MAX,
{
    // nat_from_be_bytes(bytes) < pow(256, bytes.len()) <= pow(256, size_of_usize())
    // For 32-bit: pow(256, 4) = 2^32 = usize::MAX + 1, so < pow(256,4) means <= usize::MAX.
    // For 64-bit: pow(256, 8) = 2^64 = usize::MAX + 1, same argument.
    lemma_from_be_bytes_upper_bound(bytes);
    if usize::BITS == 32 {
        reveal_with_fuel(pow, 5);  // unfolds pow(256, 0..4)
    } else {
        reveal_with_fuel(pow, 9);  // unfolds pow(256, 0..8)
    }
}

pub proof fn lemma_from_be_bytes_push(bytes: Seq<u8>, b: u8)
    ensures
        nat_from_be_bytes(bytes.push(b)) == nat_from_be_bytes(bytes) * 256 + b as nat,
{
    assert(bytes.push(b).drop_last() == bytes);
}

pub proof fn lemma_from_be_bytes_singleton(b: u8)
    ensures
        nat_from_be_bytes(seq![b]) == b as nat,
{
    lemma_from_be_bytes_push(seq![], b);
}

pub proof fn lemma_pow256_succ(exp: nat)
    ensures
        pow(256, exp + 1) == pow(256, exp) * 256,
{
    lemma_pow_adds(256, exp, 1);
    lemma_pow1(256);
}

pub broadcast proof fn lemma_from_be_bytes_upper_bound(bytes: Seq<u8>)
    ensures
        #[trigger] nat_from_be_bytes(bytes) < pow(256, bytes.len()),
    decreases bytes.len(),
{
    if bytes.len() == 0 {
        lemma_pow0(256);
    } else {
        let prefix = bytes.drop_last();
        lemma_from_be_bytes_upper_bound(prefix);
        lemma_pow256_succ(prefix.len());
    }
}

pub proof fn lemma_from_be_bytes_lower_bound(bytes: Seq<u8>)
    requires
        bytes.len() > 0,
        bytes[0] != 0,
    ensures
        pow(256, (bytes.len() - 1) as nat) <= nat_from_be_bytes(bytes),
    decreases bytes.len(),
{
    if bytes.len() == 1 {
        lemma_pow0(256);
    } else {
        let prefix = bytes.drop_last();
        lemma_from_be_bytes_lower_bound(prefix);
        lemma_pow256_succ((prefix.len() - 1) as nat);
    }
}

pub proof fn lemma_to_be_bytes_props(n: nat)
    ensures
        nat_to_be_bytes(n).len() > 0,
        n > 0 ==> nat_to_be_bytes(n)[0] != 0,
        n > 0 ==> pow(256, (nat_to_be_bytes(n).len() - 1) as nat) <= n,
    decreases n,
{
    if n < 256 {
        lemma_pow0(256);
    } else {
        let q = (n / 256) as nat;
        lemma_to_be_bytes_props(q);
        lemma_pow256_succ((nat_to_be_bytes(q).len() - 1) as nat);
    }
}

pub proof fn lemma_to_be_bytes_len_bound(n: nat, max_len: nat)
    requires
        0 < max_len,
        n < pow(256, max_len),
    ensures
        nat_to_be_bytes(n).len() <= max_len,
{
    if n == 0 {
    } else {
        lemma_to_be_bytes_props(n);
        lemma_pow_strictly_increases_converse(256, (nat_to_be_bytes(n).len() - 1) as nat, max_len);
    }
}

pub proof fn lemma_usize_to_be_bytes_len_bound(n: usize)
    ensures
        usize::BITS == 32 ==> nat_to_be_bytes(n as nat).len() <= 4,
        usize::BITS == 64 ==> nat_to_be_bytes(n as nat).len() <= 8,
{
    if usize::BITS == 32 {
        reveal_with_fuel(pow, 5);
        lemma_to_be_bytes_len_bound(n as nat, 4);
    } else {
        reveal_with_fuel(pow, 9);
        lemma_to_be_bytes_len_bound(n as nat, 8);
    }
}

pub proof fn lemma_to_from_be_bytes_roundtrip(n: nat)
    ensures
        nat_from_be_bytes(nat_to_be_bytes(n)) == n,
    decreases n,
{
    if n < 256 {
        lemma_from_be_bytes_singleton(n as u8);
    } else {
        let q = (n / 256) as nat;
        let r = (n % 256) as nat;
        lemma_to_from_be_bytes_roundtrip(q);
        lemma_from_be_bytes_push(nat_to_be_bytes(q), r as u8);
    }
}

pub proof fn lemma_from_to_be_bytes_roundtrip(bytes: Seq<u8>)
    requires
        bytes.len() > 0,
        bytes.len() > 1 ==> bytes[0] != 0,
    ensures
        nat_to_be_bytes(nat_from_be_bytes(bytes)) == bytes,
    decreases bytes.len(),
{
    if bytes.len() == 1 {
        lemma_from_be_bytes_singleton(bytes[0]);
        assert(bytes == seq![bytes[0]]);
    } else {
        let prefix = bytes.drop_last();
        lemma_from_to_be_bytes_roundtrip(prefix);
    }
}

pub proof fn lemma_from_be_bytes_prepend(bytes: Seq<u8>, b: u8)
    ensures
        nat_from_be_bytes(seq![b] + bytes) == b as nat * pow(256, bytes.len()) + nat_from_be_bytes(
            bytes,
        ),
    decreases bytes.len(),
{
    if bytes.len() == 0 {
        lemma_from_be_bytes_singleton(b);
        lemma_pow0(256);
    } else {
        let prefix = bytes.drop_last();
        let last = bytes.last();
        lemma_from_be_bytes_prepend(prefix, b);
        lemma_from_be_bytes_push(prefix, last);
        lemma_from_be_bytes_push(seq![b] + prefix, last);
        lemma_pow256_succ(prefix.len());
        assert(seq![b] + bytes == (seq![b] + prefix).push(last));
        assert((b as nat * pow(256, prefix.len()) + nat_from_be_bytes(prefix)) * 256 + last as nat
            == b as nat * (pow(256, prefix.len()) * 256) + (nat_from_be_bytes(prefix) * 256
            + last as nat)) by (nonlinear_arith);
    }
}

/// Executable loop-based big-endian base-256 decoding into `usize`.
/// Verified against the `nat_from_be_bytes` specification.
#[verifier::loop_isolation(false)]
pub fn usize_from_be_bytes_exec(bytes: &[u8]) -> (result: usize)
    requires
        bytes.len() <= size_of_usize(),
    ensures
        result as nat == nat_from_be_bytes(bytes.deep_view()),
{
    let mut acc: usize = 0;
    for &b in iter: bytes.iter()
        invariant
            acc == nat_from_be_bytes(bytes@[0..iter.index()]),
    {
        proof {
            let prefix = bytes@[0..iter.index() + 1];
            assert(prefix.drop_last() == bytes@[0..iter.index()]);
            lemma_nat_from_be_bytes_fits_shr8(prefix);
            if usize::BITS == 32 {
                lemma_usize32_shl8_or_is_base256(acc, b);
            } else {
                lemma_usize64_shl8_or_is_base256(acc, b);
            }
        }
        acc = (acc << 8usize) | (b as usize);
    }
    assert(bytes@[0..bytes@.len()] == bytes.deep_view());
    acc
}

pub fn u64_from_be_bytes(bytes: &[u8]) -> (r: u64)
    requires
        usize::BITS == 64,
        bytes.len() <= 8,
    ensures
        r as nat == nat_from_be_bytes(bytes.deep_view()),
{
    usize_from_be_bytes_exec(bytes) as u64
}

/// Executable big-endian base-256 encoding from `usize`.
/// Verified against the `nat_to_be_bytes` specification.
///
/// This allocation-backed compatibility helper builds the result recursively.
/// Serialization hot paths should prefer [`usize_to_be_bytes_in_place`].
#[cfg(feature = "alloc")]
pub fn usize_to_be_bytes_exec(v: usize) -> (buf: Vec<u8>)
    ensures
        buf@ == nat_to_be_bytes(v as nat),
    decreases v,
{
    if v < 256 {
        vec![v as u8]
    } else {
        proof {
            lemma_usize_shr8_is_div256(v);
            lemma_usize_low8_is_mod256(v);
        }
        let mut buf = usize_to_be_bytes_exec(v >> 8);
        buf.push((v & 0xff) as u8);
        buf
    }
}

/// Writes the minimal big-endian base-256 encoding of `v` into an exactly-sized slice.
#[verifier::loop_isolation(false)]
pub fn usize_to_be_bytes_in_place(v: usize, obuf: &mut [u8])
    requires
        old(obuf)@.len() == nat_to_be_bytes(v as nat).len(),
    ensures
        final(obuf)@ == nat_to_be_bytes(v as nat),
{
    let ghost len = obuf@.len();

    let ghost target = nat_to_be_bytes(v as nat);
    proof {
        lemma_to_from_be_bytes_roundtrip(v as nat);  // nat_from_be_bytes(nat_to_be_bytes(v)) == v
        assert(target[0..len] == target);
    }
    let mut current = v;

    // We write the bytes in Big-Endian order, starting from the last byte and moving backwards.
    // The loop invariant maintains that `current` is the remaining value to encode (the higher bytes),
    // and the suffix `obuf[pos..]` is the part of the output buffer that has already been filled with the correct bytes.
    let bytes = obuf.iter_mut();
    let ghost slots = bytes.remaining();
    let reversed = bytes.rev();
    proof {
        vstd::std_specs::iter::rev_postcondition(bytes, reversed);
    }
    for byte in iter: reversed
        invariant
            iter.seq().len() == len,
            current as nat == nat_from_be_bytes(target[0..len - iter.index()]),
            forall|j: int|
                #![auto]
                len as int - iter.index() <= j < len ==> *final(slots[j]) == target[j],
    {
        proof {
            let pos = len as int - iter.index() - 1;
            lemma_usize_shr8_is_div256(current);
            lemma_usize_low8_is_mod256(current);
            assert(target[0..pos + 1].drop_last() == target[0..pos]);
        }
        *byte = (current & 0xff) as u8;
        current >>= 8;
    }
}

#[cfg(feature = "alloc")]
pub fn u64_to_be_bytes(v: u64) -> (buf: Vec<u8>)
    requires
        usize::BITS == 64,
    ensures
        buf@ == nat_to_be_bytes(v as nat),
{
    usize_to_be_bytes_exec(v as usize)
}

/// Executable loop-based byte-length computation.
/// verified against the `nat_to_be_bytes` specification.
pub fn usize_to_be_bytes_len(v: usize) -> (len: usize)
    ensures
        len == nat_to_be_bytes(v as nat).len(),
{
    let mut cur = v;
    let mut len: usize = 1;
    while cur >= 256
        invariant
            len + nat_to_be_bytes(cur as nat).len() == nat_to_be_bytes(v as nat).len() + 1,
        decreases cur,
    {
        proof {
            lemma_usize_shr8_is_div256(cur);
            lemma_usize_to_be_bytes_len_bound(v);
        }
        cur >>= 8;
        len += 1;
    }
    len
}

pub fn u64_to_be_bytes_len(v: u64) -> (len: usize)
    requires
        usize::BITS == 64,
    ensures
        len == nat_to_be_bytes(v as nat).len(),
{
    usize_to_be_bytes_len(v as usize)
}

pub fn u64_to_be_bytes_first(v: u64) -> (first: u8)
    requires
        usize::BITS == 64,
    ensures
        first == nat_to_be_bytes(v as nat)[0],
    decreases v,
{
    if v < 256 {
        v as u8
    } else {
        proof {
            lemma_usize_shr8_is_div256(v as usize);
        }
        u64_to_be_bytes_first(v >> 8)
    }
}

} // verus!
#[cfg(all(test, feature = "alloc"))]
mod tests {
    use super::{usize_to_be_bytes_exec, usize_to_be_bytes_in_place, usize_to_be_bytes_len};

    #[test]
    fn usize_to_be_bytes_in_place_matches_vec_encoding() {
        for value in [0, 1, 0xff, 0x100, 0xffff, 0x10000, usize::MAX] {
            let expected = usize_to_be_bytes_exec(value);
            let len = usize_to_be_bytes_len(value);
            let mut actual = [0u8; size_of::<usize>()];

            usize_to_be_bytes_in_place(value, &mut actual[..len]);

            assert_eq!(&actual[..len], expected.as_slice());
        }
    }
}