axsys-noun 0.1.2

use std::{
    collections::hash_map::DefaultHasher,
    ffi::OsStr,
    fmt::{Display, Error, Formatter},
    hash::Hasher,
    str::{self, Utf8Error},
};

use ibig::UBig;

use crate::{sept::{date, Aura, FormatError, ParseError}, Noun};

/// Returns the length in bits of a sequence of bytes.
fn bit_len(bytes: &[u8]) -> usize {
    if let Some(last_byte) = bytes.last() {
        let byte_len = u32::try_from(bytes.len()).expect("usize to u32");
        let bit_len = u8::BITS * (byte_len - 1) + (u8::BITS - last_byte.leading_zeros());
        usize::try_from(bit_len).expect("u32 to usize")
    } else {
        0
    }
}

/// A bitwise [`Atom`] builder.
pub struct Builder {
    bytes: Vec<u8>,
    bit_idx: usize,
}

impl Builder {
    /// Creates an empty atom builder.
    pub fn new() -> Self {
        Self {
            bytes: Vec::new(),
            bit_idx: 0,
        }
    }

    /// Returns the current bitwise position of the end of this builder.
    pub fn pos(&self) -> usize {
        self.bit_idx
    }

    /// Pushes a bit onto the end of this builder.
    pub fn push_bit(&mut self, bit: bool) {
        let u8_bits = usize::try_from(u8::BITS).expect("u32 to usize");
        let byte_idx = self.bit_idx / u8_bits;
        if byte_idx == self.bytes.len() {
            self.bytes.push(0);
        }
        let byte = &mut self.bytes[byte_idx];
        let shift = self.bit_idx % u8_bits;
        if bit {
            *byte |= 1 << shift;
        } else {
            *byte &= !(1 << shift);
        }
        self.bit_idx += 1;
    }

    /// Converts this builder into an `Atom`, consuming the builder.
    pub fn into_atom(self) -> Atom {
        let bytes = self.bytes;
        let bit_len = bit_len(&bytes[..]);
        Atom { bytes, bit_len }
    }
}

impl Default for Builder {
    fn default() -> Self {
        Self::new()
    }
}

/// An arbitrarily large unsigned integer.
///
/// An [atom] is an arbitrarily large unsigned integer represented as a little-endian contiguous
/// sequence of bytes. An atom can be:
/// - created a single bit at a time or from other types that can be easily converted into atoms
///   like primitive unsigned integers, strings, and string slices;
/// - iterated over a single bit at a time;
/// - compared to other atoms and other atom-like types;
/// - pretty-printed as a hexadecimal number;
/// - converted into a noun, a primitive unsigned integer type, or a string slice.
///
/// [atom]: https://developers.urbit.org/reference/glossary/atom
///
/// # Examples
///
/// To create a new atom, use one of the `From<T>` implementations:
///
/// ```
/// # use axsys_noun::atom::Atom;
/// let atom = Atom::from("hello");
/// assert_eq!(atom, "hello");
/// ```
///
/// ```
/// # use axsys_noun::atom::Atom;
/// let atom = Atom::from(0u8);
/// assert_eq!(atom, 0u8);
/// ```
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Atom {
    bytes: Vec<u8>,
    bit_len: usize,
}

/// Converts an atom into an unsigned integer, returning `None` if the byte width of the atom
/// exceeds the byte width of the target unsigned integer type.
macro_rules! atom_as_uint {
    ($atom:expr, $uint:ty) => {{
        let atom = $atom.as_bytes();
        const N: usize = std::mem::size_of::<$uint>();
        let len = atom.len();
        if len <= N {
            let mut bytes: [u8; N] = [0; N];
            let _ = &mut bytes[..len].copy_from_slice(atom);
            Some(<$uint>::from_le_bytes(bytes))
        } else {
            None
        }
    }};
}

impl Atom {
    /// Creates an empty atom builder.
    ///
    /// This method is equivalent to `Builder::new()`.
    pub fn builder() -> Builder {
        Builder::new()
    }

    /// Creates the atom `0`.
    pub const fn null() -> Self {
        Self {
            bytes: Vec::new(),
            bit_len: 0,
        }
    }

    /// Returns `true` if this atom is null (i.e. the atom `0`).
    pub const fn is_null(&self) -> bool {
        self.bit_len() == 0
    }

    /// Returns the length in bits of this atom.
    pub const fn bit_len(&self) -> usize {
        self.bit_len
    }

    /// Computes the hash of this atom.
    pub fn hash(&self) -> u64 {
        let mut hasher = DefaultHasher::new();
        hasher.write(self.as_bytes());
        hasher.finish()
    }

    /// Converts this atom into a byte slice.
    pub fn as_bytes(&self) -> &[u8] {
        &self.bytes
    }

    /// Converts this atom into a string slice, returning an error if the atom is not composed of
    /// valid UTF-8 bytes.
    pub fn as_str(&self) -> Result<&str, Utf8Error> {
        str::from_utf8(self.as_bytes())
    }

    /// Converts this atom into an 8-bit unsigned integer, returning `None` if the atom is greater
    /// than `u8::MAX`.
    ///
    /// # Examples
    /// ```
    /// # use axsys_noun::atom::Atom;
    /// let uint = u8::MAX;
    /// let atom = Atom::from(uint);
    /// assert_eq!(atom.as_u8().unwrap(), uint);
    /// ```
    pub fn as_u8(&self) -> Option<u8> {
        atom_as_uint!(self, u8)
    }

    /// Converts this atom into an 16-bit unsigned integer, returning `None` if the atom is greater
    /// than `u16::MAX`.
    ///
    /// # Examples
    /// ```
    /// # use axsys_noun::atom::Atom;
    /// let uint = u16::MAX;
    /// let atom = Atom::from(uint);
    /// assert_eq!(atom.as_u16().unwrap(), uint);
    /// ```
    pub fn as_u16(&self) -> Option<u16> {
        atom_as_uint!(self, u16)
    }

    /// Converts this atom into an 32-bit unsigned integer, returning `None` if the atom is greater
    /// than `u32::MAX`.
    ///
    /// # Examples
    /// ```
    /// # use axsys_noun::atom::Atom;
    /// let uint = u32::MAX;
    /// let atom = Atom::from(uint);
    /// assert_eq!(atom.as_u32().unwrap(), uint);
    /// ```
    pub fn as_u32(&self) -> Option<u32> {
        atom_as_uint!(self, u32)
    }

    /// Converts this atom into an 64-bit unsigned integer, returning `None` if the atom is greater
    /// than `u64::MAX`.
    ///
    /// # Examples
    /// ```
    /// # use axsys_noun::atom::Atom;
    /// let uint = u64::MAX;
    /// let atom = Atom::from(uint);
    /// assert_eq!(atom.as_u64().unwrap(), uint);
    /// ```
    pub fn as_u64(&self) -> Option<u64> {
        atom_as_uint!(self, u64)
    }

    /// Converts this atom into an 128-bit unsigned integer, returning `None` if the atom is greater
    /// than `u128::MAX`.
    ///
    /// # Examples
    /// ```
    /// # use axsys_noun::atom::Atom;
    /// let uint = u128::MAX;
    /// let atom = Atom::from(uint);
    /// assert_eq!(atom.as_u128().unwrap(), uint);
    /// ```
    pub fn as_u128(&self) -> Option<u128> {
        atom_as_uint!(self, u128)
    }

    /// Converts this atom into a pointer-sized unsigned integer, returning `None` if the atom is
    /// greater than `usize::MAX`.
    ///
    /// # Examples
    /// ```
    /// # use axsys_noun::atom::Atom;
    /// let uint = usize::MAX;
    /// let atom = Atom::from(uint);
    /// assert_eq!(atom.as_usize().unwrap(), uint);
    /// ```
    pub fn as_usize(&self) -> Option<usize> {
        atom_as_uint!(self, usize)
    }

    /// Copies this atom into a byte vector.
    pub fn to_vec(&self) -> Vec<u8> {
        Vec::from(self.as_bytes())
    }

    /// Converts this atom into a byte vector, consuming the atom.
    ///
    /// This method does not allocate on the heap.
    pub fn into_vec(self) -> Vec<u8> {
        self.bytes
    }

    pub fn as_noun(&self) -> Noun {
        Noun::Atom(self.clone())
    }

    pub fn into_noun(self) -> Noun {
        Noun::Atom(self)
    }

    /// Returns a bitwise iterator over this atom.
    pub fn iter(&self) -> Iter {
        Iter {
            atom: self,
            bit_idx: 0,
            bit_mask: 0b1,
        }
    }

    /// Returns the atom as a big integer
    pub fn as_big(&self) -> UBig {
        UBig::from_le_bytes(&self.bytes)
    }

    fn format_decimal(&self) -> String {
        let big = self.as_big();
        let s = big.to_string();
        let mut chs = s.chars().rev().collect::<Vec<char>>();
        let len = chs.len();
        let gap = len % 3;
        if gap != 0 {
            chs.resize(len + (3 - gap), ' ');
        }
        let mut res: Vec<String> = vec![];
        for chunk in chs.chunks_exact(3) {
            let bloq = chunk.iter().rev().collect::<String>();
            res.push(bloq);
        }
        res.reverse();
        res.join(".").trim().to_string()
    }

    /// Formats an atom into a string using the given aura
    ///
    /// # Examples
    /// ```
    /// # use axsys_noun::atom::Atom;
    /// # use axsys_noun::sept::Aura;
    /// let atom = Atom::from(123434910u64);
    /// let s = atom.format_aura(Aura::U).unwrap();
    /// assert_eq!(s, "123.434.910");
    /// ```
    pub fn format_aura(&self, aura: Aura) -> Result<String, FormatError> {
        match aura {
            Aura::T => self
                .as_str()
                .map(|s| s.to_string())
                .map_err(FormatError::Utf8),
            // TODO: sanity checks
            Aura::Ta => self
                .as_str()
                .map(|s| s.to_string())
                .map_err(FormatError::Utf8),
            Aura::Tas => self
                .as_str()
                .map(|s| s.to_string())
                .map_err(FormatError::Utf8),

            Aura::U => Ok(self.format_decimal()),
            Aura::Ud => Ok(self.format_decimal()),
            Aura::Ux => unimplemented!(),
            Aura::Uv => unimplemented!(),
            Aura::Uw => unimplemented!(),
            Aura::Da => self
                .as_u128()
                .map(date::format_da)
                .ok_or(FormatError::TooLarge),
            Aura::Dr => unimplemented!(),
        }
    }

    fn parse_decimal(s: &str) -> Result<Self, ParseError> {
        let bloqs = s.split('.').collect::<Vec<&str>>().join("");
        let big = UBig::from_str_radix(&bloqs, 10).map_err(|_| ParseError::Invalid)?;
        Ok(Atom::from(big.to_le_bytes()))
    }

    /// Parses a string into an atom using the given aura
    ///
    /// # Examples
    /// ```
    /// # use axsys_noun::atom::Atom;
    /// # use axsys_noun::sept::Aura;
    /// let atom = Atom::parse_aura(Aura::U, "123.434.910").unwrap();
    /// assert_eq!(atom, Atom::from(123434910u64));
    /// ```
    pub fn parse_aura(aura: Aura, s: &str) -> Result<Self, ParseError> {
        match aura {
            Aura::T | Aura::Ta | Aura::Tas => Ok(Atom::from(s)),
            Aura::U | Aura::Ud => Self::parse_decimal(s),
            Aura::Da => Ok(Self::from(
                date::parse_da(s).map_err(|_e| ParseError::Invalid)?,
            )),

            _ => unimplemented!(),
        }
    }
}

impl Display for Atom {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        write!(f, "0x")?;
        if self.bytes.is_empty() {
            write!(f, "0")
        } else {
            for (i, byte) in (self.bytes).iter().enumerate() {
                if i > 0 && i % 4 == 0 {
                    write!(f, ".")?;
                }
                write!(f, "{:x}", byte)?;
            }
            Ok(())
        }
    }
}

impl TryFrom<&OsStr> for Atom {
    type Error = ();

    fn try_from(string: &OsStr) -> Result<Self, Self::Error> {
        Ok(Self::from(string.to_str().ok_or(())?))
    }
}

impl From<&str> for Atom {
    fn from(string: &str) -> Self {
        let bytes = string.as_bytes().to_vec();
        let bit_len = bit_len(&bytes[..]);
        Self { bytes, bit_len }
    }
}

impl From<String> for Atom {
    fn from(string: String) -> Self {
        Self::from(string.into_bytes())
    }
}

/// Convert an unsigned integer primitive into an atom.
macro_rules! impl_from_uint_for_atom {
    ($uint:ty) => {
        impl From<$uint> for Atom {
            fn from(uint: $uint) -> Self {
                Atom::from(Vec::from(uint.to_le_bytes()))
            }
        }
    };
}

impl_from_uint_for_atom!(u8);
impl_from_uint_for_atom!(u16);
impl_from_uint_for_atom!(u32);
impl_from_uint_for_atom!(u64);
impl_from_uint_for_atom!(u128);
impl_from_uint_for_atom!(usize);

impl From<Vec<u8>> for Atom {
    fn from(mut vec: Vec<u8>) -> Self {
        let len = match vec.iter().rposition(|x| *x != 0) {
            Some(idx) => idx + 1,
            None => 0,
        };
        vec.truncate(len);
        let bit_len = bit_len(&vec[..]);
        Self {
            bytes: vec,
            bit_len,
        }
    }
}

impl PartialEq<&Self> for Atom {
    fn eq(&self, other: &&Self) -> bool {
        self.bytes == other.bytes
    }
}

impl PartialEq<str> for Atom {
    fn eq(&self, other: &str) -> bool {
        if let Ok(string) = str::from_utf8(self.as_bytes()) {
            string == other
        } else {
            false
        }
    }
}

impl PartialEq<&str> for Atom {
    fn eq(&self, other: &&str) -> bool {
        if let Ok(string) = str::from_utf8(self.as_bytes()) {
            string == *other
        } else {
            false
        }
    }
}

/// Compares an atom to an unsigned integer primitive.
macro_rules! impl_partial_eq_uint_for_atom {
    ($uint:ty, $as_uint:ident) => {
        impl PartialEq<$uint> for Atom {
            fn eq(&self, other: &$uint) -> bool {
                if let Some(uint) = self.$as_uint() {
                    uint == *other
                } else {
                    false
                }
            }
        }
    };
}

impl_partial_eq_uint_for_atom!(u8, as_u8);
impl_partial_eq_uint_for_atom!(u16, as_u16);
impl_partial_eq_uint_for_atom!(u32, as_u32);
impl_partial_eq_uint_for_atom!(u64, as_u64);
impl_partial_eq_uint_for_atom!(u128, as_u128);
impl_partial_eq_uint_for_atom!(usize, as_usize);

/// An iterator over the bits of an [`Atom`].
///
/// Iteration starts with the least significant bit of the [`Atom`] and ends with the most
/// significant bit.
pub struct Iter<'a> {
    /// Atom being interated over.
    atom: &'a Atom,
    /// Index of the current bit.
    bit_idx: usize,
    /// Mask to access current bit.
    bit_mask: u8,
}

impl Iter<'_> {
    /// Returns the current bitwise position of this iterator.
    pub fn pos(&self) -> usize {
        self.bit_idx
    }
}

impl Iterator for Iter<'_> {
    type Item = bool;

    fn next(&mut self) -> Option<Self::Item> {
        if self.bit_idx == self.atom.bit_len {
            return None;
        }
        let byte_idx = self.bit_idx / usize::try_from(u8::BITS).expect("u32 to usize");
        let bit = (self.atom.bytes[byte_idx] & self.bit_mask) != 0;
        self.bit_mask = self.bit_mask.rotate_left(1);
        self.bit_idx += 1;
        Some(bit)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn bit_len() {
        {
            let num = 0b111u8.to_le_bytes();
            assert_eq!(super::bit_len(&num[..]), 3);
        }

        {
            let num = 0b10001011u8.to_le_bytes();
            assert_eq!(super::bit_len(&num[..]), 8);
        }

        {
            let num = 0b100000000u16.to_le_bytes();
            assert_eq!(super::bit_len(&num[..]), 9);
        }

        {
            let num = [
                0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
                0x37,
            ];
            assert_eq!(super::bit_len(&num[..]), 134);
        }
    }

    #[test]
    fn is_null() {
        assert!(Atom::from(0u8).is_null());
        assert!(!Atom::from(1u8).is_null());
    }

    #[test]
    fn iter() {
        {
            let atom = Atom::from(0b0u8);
            let mut atom_iter = atom.iter();
            assert_eq!(None, atom_iter.next());
        }

        {
            let atom = Atom::from(0b10u8);
            let mut atom_iter = atom.iter();
            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(true), atom_iter.next());
            assert_eq!(None, atom_iter.next());
        }

        {
            let atom = Atom::from(0x2f004u32);
            let mut atom_iter = atom.iter();
            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(true), atom_iter.next());
            assert_eq!(Some(false), atom_iter.next());

            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(false), atom_iter.next());

            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(false), atom_iter.next());

            assert_eq!(Some(true), atom_iter.next());
            assert_eq!(Some(true), atom_iter.next());
            assert_eq!(Some(true), atom_iter.next());
            assert_eq!(Some(true), atom_iter.next());

            assert_eq!(Some(false), atom_iter.next());
            assert_eq!(Some(true), atom_iter.next());

            assert_eq!(None, atom_iter.next());
        }
    }

    #[test]
    fn partial_eq() {
        {
            let lh = Atom::from("The Importance of Being Ernest");
            let rh = Atom::from("The Importance of Being Ernest");
            assert_eq!(lh, rh);
        }

        {
            let lh = Atom::from("Oh, to be a glove");
            let rh = Atom::from("upon that hand.");
            assert_ne!(lh, rh);
        }

        {
            let string = "hello";
            let atom = Atom::from(string);
            assert_eq!(atom, string);
        }

        {
            let atom = Atom::from("hello");
            assert_ne!(atom, "goodbye");
        }

        {
            macro_rules! uint_eq_test {
                ($uint:expr) => {
                    let atom = Atom::from($uint);
                    assert_eq!(atom, $uint);
                };
            }

            uint_eq_test!(0u8);
            uint_eq_test!(107u8);
            uint_eq_test!(16_000u16);
            uint_eq_test!(949_543_111u32);
            uint_eq_test!(184_884_819u64);
            uint_eq_test!(19_595_184_881_994_188_181u128);
            uint_eq_test!(10_101_044_481_818usize);
        }

        {
            macro_rules! uint_ne_test {
                ($atom:expr, $uint:expr) => {
                    let atom = Atom::from($atom);
                    assert_ne!(atom, $uint);
                };
            }

            uint_ne_test!(97u8, 103u8);
            uint_ne_test!(98u8, 64_222u16);
            uint_ne_test!(99u8, 777_919_400u32);
            uint_ne_test!(100u8, 881_944_000_887u64);
            uint_ne_test!(881_944_000_887u64, 21_601_185_860_100_176_183u128);
            uint_ne_test!(64_222u16, 127usize);
        }
    }

    #[test]
    fn format_aura() {
        let atom = Atom::from(123434910u64);
        let aura = Aura::U;
        let s = atom.format_aura(aura).unwrap();
        assert_eq!(s, "123.434.910");
        let atom = Atom::from(1234u64);
        let aura = Aura::Ud;
        let s = atom.format_aura(aura).unwrap();
        assert_eq!(s, "1.234");
        let atom = Atom::from(1234u64);
        let aura = Aura::Ud;
        let s = atom.format_aura(aura).unwrap();
        assert_eq!(s, "1.234");
    }
}