entropic 0.1.1

use core::cmp;
use core::ops::RangeInclusive;

use crate::prelude::*;
use crate::Int;

/// The default means by which base types are converted to/from entropy bytes.
#[derive(Clone, Default)]
pub struct DefaultEntropyScheme;

impl EntropyScheme for DefaultEntropyScheme {
    #[inline]
    fn get_uniform_range<'a, T: Int, I: Iterator<Item = &'a u8>>(
        &mut self,
        entropy: &mut I,
        range: RangeInclusive<T>,
    ) -> Result<T, EntropicError> {
        let (start, end) = (*range.start(), *range.end());
        if start > end {
            return Err(EntropicError::InvalidRange);
        }

        let delta = end.to_unsigned().wrapping_sub(start.to_unsigned());

        // Short-circuit to prevent any entropy being consumed.
        if delta == T::Unsigned::ZERO {
            return Ok(start);
        }

        // Get random number in range of 0..=delta
        let mut rem_delta = (delta >> 4) >> 4; // `>> 8` panics for i8
        let mut offset = T::Unsigned::from_u8(self.get_byte(entropy)?);

        while rem_delta > T::Unsigned::ZERO {
            rem_delta = rem_delta >> 8;
            offset = offset << 8;
            offset = offset
                .checked_add(T::Unsigned::from_u8(self.get_byte(entropy)?))
                .ok_or(EntropicError::Internal)?;
        }

        match delta.checked_add(T::Unsigned::ONE) {
            // range was all-inclusive
            None => Ok(T::from_unsigned(offset)),
            // Range is constrained
            Some(modulus) => {
                offset = offset % modulus; // Constrain offset to range
                Ok(start.wrapping_add(T::from_unsigned(offset)))
            }
        }
    }

    #[inline]
    fn put_uniform_range<'a, T: Int, I: Iterator<Item = &'a mut u8>>(
        &mut self,
        entropy: &mut I,
        range: RangeInclusive<T>,
        value: T,
    ) -> Result<usize, EntropicError> {
        let (start, end) = (*range.start(), *range.end());
        if start > end {
            return Err(EntropicError::InvalidRange);
        }

        if value < start || value > end {
            return Err(EntropicError::ValueOutOfRange);
        }

        let delta = end.to_unsigned().wrapping_sub(start.to_unsigned());
        let offset = value.to_unsigned().wrapping_sub(start.to_unsigned());

        // Short-circuit to prevent any entropy being consumed.
        if delta == T::Unsigned::ZERO {
            return Ok(0);
        }

        // Determine how many bytes `delta` fills
        // ilog2 is index of leftmost set bit (0-indexed from the right)
        // Add 1 and divide by 8 to get the number of bytes that have set bits
        let delta_bytes = (delta
            .checked_ilog2()
            .unwrap_or(0) // Shouldn't ever happen
            .checked_add(8) // Need to add 1 to account for ilog2, then 7 more to have `/ 8` give us the ceiling rather than floor
            .ok_or(EntropicError::Internal)?
            / 8) as usize;

        // Insert offset bytes in reverse order from which they were taken
        for i in (0..delta_bytes).rev() {
            self.put_byte(entropy, (offset >> (i * 8)).to_u8())?;
        }

        Ok(delta_bytes)
    }

    #[inline]
    fn get_uniform_ranges<'a, T: Int, I: Iterator<Item = &'a u8>, const L: usize>(
        &mut self,
        entropy: &mut I,
        ranges: &[RangeInclusive<T>; L],
    ) -> Result<T, EntropicError> {
        let Some((first_range, remaining_ranges)) = ranges.split_first() else {
            return Err(EntropicError::InvalidRange); // array of size 0 not allowed
        };

        if first_range.start() > first_range.end() {
            return Err(EntropicError::InvalidRange);
        }

        let mut delta = first_range
            .end()
            .to_unsigned()
            .wrapping_sub(first_range.start().to_unsigned());
        let mut prev_end = *first_range.end();

        for range in remaining_ranges {
            if prev_end >= *range.start() || range.start() > range.end() {
                return Err(EntropicError::InvalidRange);
            }

            let rdelta = range
                .end()
                .to_unsigned()
                .wrapping_sub(range.start().to_unsigned());

            // Invariant: ranges that cannot overlap and are strictly increasing
            // will not yield a total delta greater than T::Unsigned::MAX
            delta = delta
                .checked_add(rdelta)
                .and_then(|d| d.checked_add(T::Unsigned::ONE))
                .ok_or(EntropicError::Internal)?;
            prev_end = *range.end();
        }

        let mut offset = self.get_uniform_range(entropy, T::Unsigned::ZERO..=delta)?;
        for range in ranges {
            let rdelta = range
                .end()
                .to_unsigned()
                .wrapping_sub(range.start().to_unsigned());
            match offset.checked_sub(rdelta) {
                Some(s) => match s.checked_sub(T::Unsigned::ONE) {
                    None => return Ok(*range.end()),
                    Some(new_offset) => offset = new_offset,
                },
                None => {
                    return (*range.start())
                        .checked_add(T::from_unsigned(offset))
                        .ok_or(EntropicError::Internal)
                }
            }
        }

        // Invariant: method will reach value at some point while iterating ranges
        Err(EntropicError::Internal) // Unreachable
    }

    #[inline]
    fn put_uniform_ranges<'a, T: Int, I: Iterator<Item = &'a mut u8>, const L: usize>(
        &mut self,
        entropy: &mut I,
        ranges: &[RangeInclusive<T>; L],
        value: T,
    ) -> Result<usize, EntropicError> {
        let Some((first_range, remaining_ranges)) = ranges.split_first() else {
            return Err(EntropicError::InvalidRange); // array of size 0 not allowed
        };

        if first_range.start() > first_range.end() {
            return Err(EntropicError::InvalidRange);
        }

        let mut delta = first_range
            .end()
            .to_unsigned()
            .wrapping_sub(first_range.start().to_unsigned());
        let mut prev_end = *first_range.end();

        for range in remaining_ranges {
            if prev_end >= *range.start() || range.start() > range.end() {
                return Err(EntropicError::InvalidRange);
            }
            if range.start() > range.end() {
                return Err(EntropicError::InvalidRange);
            }

            let rdelta = range
                .end()
                .to_unsigned()
                .wrapping_sub(range.start().to_unsigned());

            // Invariant: ranges that cannot overlap and are strictly increasing
            // will not yield a total delta greater than T::Unsigned::MAX
            delta = delta
                .checked_add(rdelta)
                .and_then(|d| d.checked_add(T::Unsigned::ONE))
                .ok_or(EntropicError::Internal)?;
            prev_end = *range.end();
        }

        let mut offset = T::Unsigned::ZERO;
        for range in ranges {
            if value < *range.start() {
                return Err(EntropicError::ValueOutOfRange);
            }

            if value <= *range.end() {
                // Invariant: `offset` is strictly less than or equal to `delta`
                offset = offset
                    .checked_add(
                        value
                            .to_unsigned()
                            .wrapping_sub(range.start().to_unsigned()),
                    )
                    .ok_or(EntropicError::Internal)?;
                return self.put_uniform_range(entropy, T::Unsigned::ZERO..=delta, offset);
            } else {
                let rdelta = range
                    .end()
                    .to_unsigned()
                    .wrapping_sub(range.start().to_unsigned());
                offset = offset
                    .checked_add(rdelta)
                    .and_then(|o| o.checked_add(T::Unsigned::ONE))
                    .ok_or(EntropicError::Internal)?;
            }
        }

        // Invariant: method will reach value at some point while iterating ranges
        Err(EntropicError::Internal) // Unreachable
    }

    #[inline]
    fn get_bounded_len<'a, I: Iterator<Item = &'a u8>>(
        &mut self,
        entropy: &mut I,
        range: RangeInclusive<usize>,
    ) -> Result<usize, EntropicError> {
        let delta = range.end().saturating_sub(*range.start());
        if delta == 0 {
            return Ok(*range.start());
        }

        let b = self.get_byte(entropy)?;
        let offset = if b & 0b0000_0011 != 0b0000_0011 || delta <= 2 {
            Ok((b & 0b0000_0011) as usize) // 1/4 chance each for 0, 1, 2
        } else if b & 0b0000_1100 != 0b0000_1100 || delta <= 5 {
            Ok(((b & 0b0000_1100) >> 2) as usize + 3) // 1/16 chance each for 3, 4, 5
        } else if b & 0b0001_0000 != 0b0001_0000 || delta <= 25 {
            self.get_uniform_range(entropy, 6usize..=25usize) // 1/32 chance for 6-25
        } else if b & 0b0010_0000 != 0b0010_0000 || delta <= 150 {
            self.get_uniform_range(entropy, 26usize..=150usize) // 1/64 chance for 26-150
        } else if b & 0b0100_0000 != 0b0100_0000 || delta <= 1500 {
            self.get_uniform_range(entropy, 151usize..=1500usize) // 1/128 chance for 151-1500
        } else if b & 0b1000_0000 == 0b1000_0000 || delta <= 7500 {
            self.get_uniform_range(entropy, 1501usize..=7500usize) // 1/256 chance for 1500-7500
        } else {
            self.get_uniform_range(entropy, 7500usize..=usize::MAX) // 1/256 chance for >7500
        }?;

        let value = range.start().saturating_add(offset);
        Ok(cmp::min(value, *range.end()))
    }

    #[inline]
    fn put_bounded_len<'a, I: Iterator<Item = &'a mut u8>>(
        &mut self,
        entropy: &mut I,
        range: RangeInclusive<usize>,
        value: usize,
    ) -> Result<usize, EntropicError> {
        if value < *range.start() || value > *range.end() {
            return Err(EntropicError::ValueOutOfRange);
        }

        if range.end().checked_sub(*range.start()) == Some(0) {
            return Ok(0);
        }

        // let delta = range.end().saturating_sub(*range.start());
        let offset = value.saturating_sub(*range.start());

        Ok(match offset {
            0 => self.put_byte(entropy, 0b0000_0000)?,
            1 => self.put_byte(entropy, 0b0000_0001)?,
            2 => self.put_byte(entropy, 0b0000_0010)?,
            3 => self.put_byte(entropy, 0b0000_0011)?,
            4 => self.put_byte(entropy, 0b0000_0111)?,
            5 => self.put_byte(entropy, 0b0000_1011)?,
            6..=25 => {
                let length = self.put_byte(entropy, 0b0000_1111)?;
                length
                    .checked_add(self.put_uniform_range(entropy, 6usize..=25usize, offset)?)
                    .ok_or(EntropicError::Internal)?
            }
            26..=150 => {
                let length = self.put_byte(entropy, 0b0001_1111)?;
                length
                    .checked_add(self.put_uniform_range(entropy, 26usize..=150usize, offset)?)
                    .ok_or(EntropicError::Internal)?
            }
            151..=1500 => {
                let length = self.put_byte(entropy, 0b0011_1111)?;
                length
                    .checked_add(self.put_uniform_range(entropy, 151usize..=1500usize, offset)?)
                    .ok_or(EntropicError::Internal)?
            }
            1501..=7500 => {
                let length = self.put_byte(entropy, 0b1111_1111)?;
                length
                    .checked_add(self.put_uniform_range(entropy, 1501usize..=7500usize, offset)?)
                    .ok_or(EntropicError::Internal)?
            }
            _ => {
                let length = self.put_byte(entropy, 0b0111_1111)?;
                length
                    .checked_add(self.put_uniform_range(entropy, 7501usize..=usize::MAX, offset)?)
                    .ok_or(EntropicError::Internal)?
            }
        })
    }

    #[inline]
    fn get_unbounded_len<'a, I: Iterator<Item = &'a u8>>(
        &mut self,
        entropy: &mut I,
    ) -> Result<usize, EntropicError> {
        self.get_bounded_len(entropy, 0..=1_000_000)
    }

    #[inline]
    fn put_unbounded_len<'a, I: Iterator<Item = &'a mut u8>>(
        &mut self,
        entropy: &mut I,
        value: usize,
    ) -> Result<usize, EntropicError> {
        self.put_bounded_len(entropy, 0..=1_000_000, value)
    }

    #[inline]
    fn get_optional<'a, I: Iterator<Item = &'a u8>>(
        &mut self,
        entropy: &mut I,
    ) -> Result<bool, EntropicError> {
        if self.get_byte(entropy)? & 1 == 1 {
            Ok(true)
        } else {
            Ok(false)
        }
    }

    #[inline]
    fn put_optional<'a, I: Iterator<Item = &'a mut u8>>(
        &mut self,
        entropy: &mut I,
        is_some: bool,
    ) -> Result<usize, EntropicError> {
        self.put_byte(entropy, if is_some { 1 } else { 0 })
    }
}