// This file is part of ICU4X. For terms of use, please see the file

// called LICENSE at the top level of the ICU4X source tree

// (online at: https://github.com/unicode-org/icu4x/blob/main/LICENSE ).

//! Varint spec for ZeroTrie:

//!

//! - Lead byte: top M (2 or 3) bits are metadata; next is varint extender; rest is value

//! - Trail bytes: top bit is varint extender; rest are low bits of value

//! - Guaranteed uniqueness of varint by adding "latent value" for each extender byte

//! - No maximum, but high bits will be dropped if they don't fit in the platform's `usize`

//!

//! This is best shown by examples.

//!

//! ```txt

//! xxx0'1010 = 10

//! xxx0'1111 = 15 (largest single-byte value with M=3)

//! xxx1'0000 0000'0000 must be 16 (smallest two-byte value with M=3)

//! xxx1'0000 0000'0001 = 17

//! xxx1'1111 0111'1111 = 2063 (largest two-byte value with M=3)

//! xxx1'0000 1000'0000 0000'0000 must be 2064 (smallest three-byte value with M=3)

//! xxx1'0000 1000'0000 0000'0001 = 2065

//! ```

//!

//! The latent values by number of bytes for M=3 are:

//!

//! - 1 byte: 0

//! - 2 bytes: 16 = 0x10 = 0b10000

//! - 3 bytes: 2064 = 0x810 = 0b100000010000

//! - 4 bytes: 264208 = 0x40810 = 0b1000000100000010000

//! - 5 bytes: 33818640 = 0x2040810 = 0b10000001000000100000010000

//! - …

//!

//! For M=2, the latent values are:

//!

//! - 1 byte: 0

//! - 2 bytes: 32 = 0x20 = 0b100000

//! - 3 bytes: 4128 = 0x1020 = 0b1000000100000

//! - 4 bytes: 524320 = 0x81020 = 0b10000001000000100000

//! - 5 bytes: 67637280 = 0x4081020 = 0b100000010000001000000100000

//! - …

use crate::builder::konst::ConstArrayBuilder;

#[cfg(feature = "alloc")]

use crate::builder::nonconst::TrieBuilderStore;

/// Reads a varint with 2 bits of metadata in the lead byte.

///

/// Returns the varint value and a subslice of `remainder` with the varint bytes removed.

///

/// If the varint spills off the end of the slice, a debug assertion will fail,

/// and the function will return the value up to that point.

pub const fn read_varint_meta2(start: u8, remainder: &[u8]) -> (usize, &[u8]) {

    let mut value = (start & 0b00011111) as usize;

    let mut remainder = remainder;

    if (start & 0b00100000) != 0 {

        loop {

            let next;

            (next, remainder) = debug_unwrap!(remainder.split_first(), break, "invalid varint");

            // Note: value << 7 could drop high bits. The first addition can't overflow.

            // The second addition could overflow; in such a case we just inform the

            // developer via the debug assertion.

            value = (value << 7) + ((*next & 0b01111111) as usize) + 32;

            if (*next & 0b10000000) == 0 {

                break;

            }

        }

    }

    (value, remainder)

}

/// Reads a varint with 3 bits of metadata in the lead byte.

///

/// Returns the varint value and a subslice of `remainder` with the varint bytes removed.

///

/// If the varint spills off the end of the slice, a debug assertion will fail,

/// and the function will return the value up to that point.

pub const fn read_varint_meta3(start: u8, remainder: &[u8]) -> (usize, &[u8]) {

    let mut value = (start & 0b00001111) as usize;

    let mut remainder = remainder;

    if (start & 0b00010000) != 0 {

        loop {

            let next;

            (next, remainder) = debug_unwrap!(remainder.split_first(), break, "invalid varint");

            // Note: value << 7 could drop high bits. The first addition can't overflow.

            // The second addition could overflow; in such a case we just inform the

            // developer via the debug assertion.

            value = (value << 7) + ((*next & 0b01111111) as usize) + 16;

            if (*next & 0b10000000) == 0 {

                break;

            }

        }

    }

    (value, remainder)

}

/// Reads and removes a varint with 3 bits of metadata from a [`TrieBuilderStore`].

///

/// Returns the varint value.

#[cfg(feature = "alloc")]

pub(crate) fn try_read_varint_meta3_from_tstore<S: TrieBuilderStore>(

    start: u8,

    remainder: &mut S,

) -> Option<usize> {

    let mut value = (start & 0b00001111) as usize;

    if (start & 0b00010000) != 0 {

        loop {

            let next = remainder.atbs_pop_front()?;

            // Note: value << 7 could drop high bits. The first addition can't overflow.

            // The second addition could overflow; in such a case we just inform the

            // developer via the debug assertion.

            value = (value << 7) + ((next & 0b01111111) as usize) + 16;

            if (next & 0b10000000) == 0 {

                break;

            }

        }

    }

    Some(value)

}

#[cfg(test)]

const MAX_VARINT: usize = usize::MAX;

// *Upper Bound:* Each trail byte stores 7 bits of data, plus the latent value.

// Add an extra 1 since the lead byte holds only 5 bits of data.

const MAX_VARINT_LENGTH: usize = 1 + core::mem::size_of::<usize>() * 8 / 7;

/// Returns a new [`ConstArrayBuilder`] containing a varint with 2 bits of metadata.

pub(crate) const fn write_varint_meta2(value: usize) -> ConstArrayBuilder<MAX_VARINT_LENGTH, u8> {

    let mut result = [0; MAX_VARINT_LENGTH];

    let mut i = MAX_VARINT_LENGTH - 1;

    let mut value = value;

    let mut last = true;

    loop {

        if value < 32 {

            result[i] = value as u8;

            if !last {

                result[i] |= 0b00100000;

            }

            break;

        }

        value -= 32;

        result[i] = (value as u8) & 0b01111111;

        if !last {

            result[i] |= 0b10000000;

        } else {

            last = false;

        }

        value >>= 7;

        i -= 1;

    }

    // The bytes are from i to the end.

    ConstArrayBuilder::from_manual_slice(result, i, MAX_VARINT_LENGTH)

}

/// Returns a new [`ConstArrayBuilder`] containing a varint with 3 bits of metadata.

pub(crate) const fn write_varint_meta3(value: usize) -> ConstArrayBuilder<MAX_VARINT_LENGTH, u8> {

    let mut result = [0; MAX_VARINT_LENGTH];

    let mut i = MAX_VARINT_LENGTH - 1;

    let mut value = value;

    let mut last = true;

    loop {

        if value < 16 {

            result[i] = value as u8;

            if !last {

                result[i] |= 0b00010000;

            }

            break;

        }

        value -= 16;

        result[i] = (value as u8) & 0b01111111;

        if !last {

            result[i] |= 0b10000000;

        } else {

            last = false;

        }

        value >>= 7;

        i -= 1;

    }

    // The bytes are from i to the end.

    ConstArrayBuilder::from_manual_slice(result, i, MAX_VARINT_LENGTH)

}

/// A secondary implementation that separates the latent value while computing the varint.

#[cfg(test)]

pub(crate) const fn write_varint_reference(

    value: usize,

) -> ConstArrayBuilder<MAX_VARINT_LENGTH, u8> {

    let mut result = [0; MAX_VARINT_LENGTH];

    if value < 32 {

        result[0] = value as u8;

        return ConstArrayBuilder::from_manual_slice(result, 0, 1);

    }

    result[0] = 32;

    let mut latent = 32;

    let mut steps = 2;

    loop {

        let next_latent = (latent << 7) + 32;

        if value < next_latent || next_latent == latent {

            break;

        }

        latent = next_latent;

        steps += 1;

    }

    let mut value = value - latent;

    let mut i = steps;

    while i > 0 {

        i -= 1;

        result[i] |= (value as u8) & 0b01111111;

        value >>= 7;

        if i > 0 && i < steps - 1 {

            result[i] |= 0b10000000;

        }

    }

    // The bytes are from 0 to `steps`.

    ConstArrayBuilder::from_manual_slice(result, 0, steps)

}

#[cfg(test)]

mod tests {

    use super::*;

    #[derive(Debug)]

    struct TestCase<'a> {

        bytes: &'a [u8],

        remainder: &'a [u8],

        value: usize,

    }

    static CASES: &[TestCase] = &[

        TestCase {

            bytes: &[0b00000000],

            remainder: &[],

            value: 0,

        },

        TestCase {

            bytes: &[0b00001010],

            remainder: &[],

            value: 10,

        },

        TestCase {

            bytes: &[0b00011111],

            remainder: &[],

            value: 31,

        },

        TestCase {

            bytes: &[0b00011111, 0b10101010],

            remainder: &[0b10101010],

            value: 31,

        },

        TestCase {

            bytes: &[0b00100000, 0b00000000],

            remainder: &[],

            value: 32,

        },

        TestCase {

            bytes: &[0b00100000, 0b00000001],

            remainder: &[],

            value: 33,

        },

        TestCase {

            bytes: &[0b00100000, 0b00100000],

            remainder: &[],

            value: 64,

        },

        TestCase {

            bytes: &[0x20, 0x44],

            remainder: &[],

            value: 100,

        },

        TestCase {

            bytes: &[0b00100000, 0b01111111],

            remainder: &[],

            value: 159,

        },

        TestCase {

            bytes: &[0b00100001, 0b00000000],

            remainder: &[],

            value: 160,

        },

        TestCase {

            bytes: &[0b00100001, 0b00000001],

            remainder: &[],

            value: 161,

        },

        TestCase {

            bytes: &[0x23, 0x54],

            remainder: &[],

            value: 500,

        },

        TestCase {

            bytes: &[0b00111111, 0b01111111],

            remainder: &[],

            value: 4127, // 32 + (1 << 12) - 1

        },

        TestCase {

            bytes: &[0b00100000, 0b10000000, 0b00000000],

            remainder: &[],

            value: 4128, // 32 + (1 << 12)

        },

        TestCase {

            bytes: &[0b00100000, 0b10000000, 0b00000001],

            remainder: &[],

            value: 4129, // 32 + (1 << 12) + 1

        },

        TestCase {

            bytes: &[0b00100000, 0b10000000, 0b01111111],

            remainder: &[],

            value: 4255, // 32 + (1 << 12) + 127

        },

        TestCase {

            bytes: &[0b00100000, 0b10000001, 0b00000000],

            remainder: &[],

            value: 4256, // 32 + (1 << 12) + 128

        },

        TestCase {

            bytes: &[0b00100000, 0b10000001, 0b00000001],

            remainder: &[],

            value: 4257, // 32 + (1 << 12) + 129

        },

        TestCase {

            bytes: &[0x20, 0x86, 0x68],

            remainder: &[],

            value: 5000,

        },

        TestCase {

            bytes: &[0b00100000, 0b11111111, 0b01111111],

            remainder: &[],

            value: 20511, // 32 + (1 << 12) + (1 << 14) - 1

        },

        TestCase {

            bytes: &[0b00100001, 0b10000000, 0b00000000],

            remainder: &[],

            value: 20512, // 32 + (1 << 12) + (1 << 14)

        },

        TestCase {

            bytes: &[0b00111111, 0b11111111, 0b01111111],

            remainder: &[],

            value: 528415, // 32 + (1 << 12) + (1 << 19) - 1

        },

        TestCase {

            bytes: &[0b00100000, 0b10000000, 0b10000000, 0b00000000],

            remainder: &[],

            value: 528416, // 32 + (1 << 12) + (1 << 19)

        },

        TestCase {

            bytes: &[0b00100000, 0b10000000, 0b10000000, 0b00000001],

            remainder: &[],

            value: 528417, // 32 + (1 << 12) + (1 << 19) + 1

        },

        TestCase {

            bytes: &[0b00111111, 0b11111111, 0b11111111, 0b01111111],

            remainder: &[],

            value: 67637279, // 32 + (1 << 12) + (1 << 19) + (1 << 26) - 1

        },

        TestCase {

            bytes: &[0b00100000, 0b10000000, 0b10000000, 0b10000000, 0b00000000],

            remainder: &[],

            value: 67637280, // 32 + (1 << 12) + (1 << 19) + (1 << 26)

        },

    ];

    #[test]

    fn test_read() {

        for cas in CASES {

            let recovered = read_varint_meta2(cas.bytes[0], &cas.bytes[1..]);

            assert_eq!(recovered, (cas.value, cas.remainder), "{:?}", cas);

        }

    }

    #[test]

    fn test_read_write() {

        for cas in CASES {

            let reference_bytes = write_varint_reference(cas.value);

            assert_eq!(

                reference_bytes.len(),

                cas.bytes.len() - cas.remainder.len(),

                "{:?}",

                cas

            );

            assert_eq!(

                reference_bytes.as_slice(),

                &cas.bytes[0..reference_bytes.len()],

                "{:?}",

                cas

            );

            let recovered = read_varint_meta2(cas.bytes[0], &cas.bytes[1..]);

            assert_eq!(recovered, (cas.value, cas.remainder), "{:?}", cas);

            let write_bytes = write_varint_meta2(cas.value);

            assert_eq!(

                reference_bytes.as_slice(),

                write_bytes.as_slice(),

                "{:?}",

                cas

            );

        }

    }

    #[test]

    fn test_roundtrip() {

        let mut i = 0usize;

        while i < MAX_VARINT {

            let bytes = write_varint_meta2(i);

            let recovered = read_varint_meta2(bytes.as_slice()[0], &bytes.as_slice()[1..]);

            assert_eq!(i, recovered.0, "{:?}", bytes.as_slice());

            i <<= 1;

            i += 1;

        }

    }

    #[test]

    fn test_extended_roundtrip() {

        let mut i = 0usize;

        while i < MAX_VARINT {

            let bytes = write_varint_meta3(i);

            let recovered = read_varint_meta3(bytes.as_slice()[0], &bytes.as_slice()[1..]);

            assert_eq!(i, recovered.0, "{:?}", bytes.as_slice());

            i <<= 1;

            i += 1;

        }

    }

    #[test]

    fn test_max() {

        let reference_bytes = write_varint_reference(MAX_VARINT);

        let write_bytes = write_varint_meta2(MAX_VARINT);

        assert_eq!(reference_bytes.len(), MAX_VARINT_LENGTH);

        assert_eq!(reference_bytes.as_slice(), write_bytes.as_slice());

        let subarray = write_bytes

            .as_const_slice()

            .get_subslice_or_panic(1, write_bytes.len());

        let (recovered_value, remainder) = read_varint_meta2(

            *write_bytes.as_const_slice().first().unwrap(),

            subarray.as_slice(),

        );

        assert!(remainder.is_empty());

        assert_eq!(recovered_value, MAX_VARINT);

        assert_eq!(

            write_bytes.as_slice(),

            &[

                0b00100001, //

                0b11011111, //

                0b11011111, //

                0b11011111, //

                0b11011111, //

                0b11011111, //

                0b11011111, //

                0b11011111, //

                0b11011111, //

                0b01011111, //

            ]

        );

    }

    #[test]

    fn text_extended_max() {

        let write_bytes = write_varint_meta3(MAX_VARINT);

        assert_eq!(write_bytes.len(), MAX_VARINT_LENGTH);

        let (lead, trailing) = write_bytes.as_slice().split_first().unwrap();

        let (recovered_value, remainder) = read_varint_meta3(*lead, trailing);

        assert!(remainder.is_empty());

        assert_eq!(recovered_value, MAX_VARINT);

        assert_eq!(

            write_bytes.as_slice(),

            &[

                0b00010001, //

                0b11101111, //

                0b11101111, //

                0b11101111, //

                0b11101111, //

                0b11101111, //

                0b11101111, //

                0b11101111, //

                0b11101111, //

                0b01101111, //

            ]

        );

    }

    #[test]

    fn test_latent_values() {

        // Same values documented in the module docs: M=2

        let m2 = read_varint_meta2;

        assert_eq!(m2(0, &[]).0, 0);

        assert_eq!(m2(0x20, &[0x00]).0, 32);

        assert_eq!(m2(0x20, &[0x80, 0x00]).0, 4128);

        assert_eq!(m2(0x20, &[0x80, 0x80, 0x00]).0, 528416);

        assert_eq!(m2(0x20, &[0x80, 0x80, 0x80, 0x00]).0, 67637280);

        // Same values documented in the module docs: M=3

        let m3 = read_varint_meta3;

        assert_eq!(m3(0, &[]).0, 0);

        assert_eq!(m3(0x10, &[0x00]).0, 16);

        assert_eq!(m3(0x10, &[0x80, 0x00]).0, 2064);

        assert_eq!(m3(0x10, &[0x80, 0x80, 0x00]).0, 264208);

        assert_eq!(m3(0x10, &[0x80, 0x80, 0x80, 0x00]).0, 33818640);

    }

}