onepass-seed 0.3.2

use core::cmp::max;
use std::io::{Result, Write};

use crypto_bigint::{NonZero, U256};
use secrecy::ExposeSecretMut;

use super::{Eval, util::u256_to_word};

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Chars(pub Box<[CharRange]>);

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct CharRange {
    pub start: char,
    pub end: char,
}

impl From<(char, char)> for CharRange {
    fn from((start, end): (char, char)) -> Self {
        CharRange { start, end }
    }
}

impl Chars {
    /// # Safety
    /// This function is only safe for evaluation if the ranges are non-overlapping.
    pub unsafe fn from_ranges_unchecked<T, V>(ranges: V) -> Self
    where
        V: IntoIterator<Item = T>,
        T: Into<CharRange>,
    {
        Chars(ranges.into_iter().map(Into::into).collect())
    }

    pub fn from_ranges<T, V>(ranges: V) -> Self
    where
        V: IntoIterator<Item = T>,
        T: Into<CharRange>,
    {
        let mut ranges = ranges.into_iter().map(Into::into).collect::<Vec<_>>();
        ranges.sort_unstable_by_key(|a| a.start);
        let mut i = 0;
        let mut j = 1;
        while j < ranges.len() {
            let Some(next) = next_char(ranges[i].end) else {
                break;
            };
            if next >= ranges[j].start {
                ranges[i].end = max(ranges[i].end, ranges[j].end);
                j += 1;
                continue;
            }
            if j != i + 1 {
                ranges.swap(i + 1, j);
            }
            i += 1;
            j += 1;
        }
        ranges.drain(i + 1..);
        Chars(ranges.into())
    }

    fn size(&self) -> u32 {
        self.0.iter().map(|range| range.size()).sum()
    }

    fn nth(&self, mut n: u32) -> char {
        for range in &self.0 {
            let sz = range.size();
            if n < sz {
                return range.nth(n);
            }
            n -= sz;
        }
        unreachable!()
    }
}

pub(super) fn next_char(c: char) -> Option<char> {
    match c {
        '\u{d7ff}' => Some(0xe000),
        _ => u32::from(c).checked_add(1),
    }
    .and_then(char::from_u32)
}

impl CharRange {
    // TODO(someday): replace these with `Step` methods once those are stabilized.

    fn size(&self) -> u32 {
        let start = self.start as u32;
        let end = self.end as u32;
        assert!(start <= end, "{:?} > {:?}", self.start, self.end);
        let count = end - start + 1;
        if start < 0xD800 && 0xE000 <= end {
            count - 0x800
        } else {
            count
        }
    }

    fn nth(&self, n: u32) -> char {
        let start = self.start as u32;
        let res = start + n;
        let res = if start < 0xD800 && res >= 0xD800 {
            res + 0x800
        } else {
            res
        };
        let res = char::from_u32(res).unwrap();
        assert!(res <= self.end);
        res
    }
}

impl Eval for Chars {
    fn size(&self) -> NonZero<U256> {
        NonZero::new(Chars::size(self).into()).unwrap()
    }

    fn write_to(&self, w: &mut dyn Write, index: &mut dyn ExposeSecretMut<U256>) -> Result<()> {
        let c = self.nth(u256_to_word(index.expose_secret_mut()).try_into().unwrap());
        write!(w, "{}", c)
    }
}

#[cfg(test)]
mod tests {
    use std::io::BufWriter;

    use secrecy::SecretBox;

    use super::*;

    #[test]
    fn test_from_ranges() {
        let rs = Chars::from_ranges([('b', 'e'), ('a', 'c'), ('z', 'z')]);
        assert_eq!(6, rs.size());
        assert_eq!('a', rs.nth(0));
        assert_eq!('z', rs.nth(5));

        let rs = Chars::from_ranges(vec![('a', 'z')]);
        assert_eq!(26, rs.size());
    }

    #[test]
    fn test_eval_boundary() {
        let cs = Chars::from_ranges([('\u{d7ff}', char::MAX)]);
        let mut buf = BufWriter::new(Vec::new());
        cs.write_to(&mut buf, &mut SecretBox::new(Box::new(U256::ONE)))
            .unwrap();
        assert_eq!(
            "\u{e000}",
            String::from_utf8(buf.into_inner().unwrap()).unwrap()
        );
        assert_eq!(char::MAX as u32 - 0xe000 + 2, cs.size());
    }

    #[test]
    fn test_next_char_boundary() {
        assert_eq!(Some('\u{e000}'), next_char('\u{d7ff}'));
    }
}