zoe 0.0.31

//! Specification structs for generating arbitrary bytes (`u8` values).

use crate::data::arbitrary::ArbitrarySpecs;
use arbitrary::{Result, Unstructured};

/// Specifications for generating an arbitrary `u8` byte.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)]
pub struct ByteSpecs {
    /// The character set to which the `u8` byte must belong.
    pub set: ByteSet,

    /// The case of the ASCII characters.
    pub case: Case,
}

impl<'a> ArbitrarySpecs<'a> for ByteSpecs {
    type Output = u8;

    #[inline]
    fn make_arbitrary(&self, u: &mut Unstructured<'a>) -> Result<Self::Output> {
        let byte = self.set.make_arbitrary(u)?;

        let cased = match self.case {
            Case::Lowercase => byte.to_ascii_lowercase(),
            Case::Uppercase => byte.to_ascii_uppercase(),
            Case::Any => byte,
        };

        Ok(cased)
    }
}

/// The desired case of the ASCII generated byte(s).
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)]
pub enum Case {
    /// Converts all ASCII letters to lowercase.
    Lowercase,

    /// Converts all ASCII letters to uppercase.
    Uppercase,

    /// Imposes no restrictions on the case of the generated byte(s).
    #[default]
    Any,
}

/// The character set of the generated byte(s).
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)]
pub enum ByteSet {
    /// Includes any `u8` value.
    #[default]
    Any,

    /// Includes only ASCII values (range `0..=127`).
    Ascii,

    /// Includes only graphic ASCII values or spaces (range `32..=126`).
    AsciiGraphicOrSpace,

    /// Includes only graphic ASCII values (range `33..=126`).
    AsciiGraphic,

    /// Includes only graphic ASCII that is not a special character in bash.
    ///
    /// This includes the characters
    /// `%+,-./0123456789:@ABCDEFGHIJKLMNOPQRSTUVWXYZ^_abcdefghijklmnopqrstuvwxyz`.
    AsciiGraphicBashSafe,

    /// Includes the canonical DNA bases `ACGT` in both cases, with additional
    /// options configured by [`NucleotidesByteSet`].
    NucleotidesCanonical(NucleotidesByteSet),

    /// Includes the IUPAC DNA bases `ACGTURYSWKMBDHV` in both cases, with
    /// additional options configured by [`NucleotidesByteSet`].
    NucleotidesIupac(NucleotidesByteSet),

    /// Includes the 20 canonical amino acids `ACDEFGHIKLMNPQRSTVWY` in both
    /// cases, with additional options configured by [`AminoAcidsByteSet`].
    AminoAcidsCanonical(AminoAcidsByteSet),

    /// Includes a custom set of bytes.
    Custom(&'static [u8]),
}

/// Specifications for what special characters should be included when
/// generating arbitrary nucleotide bases.
///
/// All options are opt-out. By default, `N` and `U` and `-` will be included.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct NucleotidesByteSet {
    /// Whether to include `N` in both cases.
    pub include_n: bool,

    /// Whether to include `U` in both cases.
    pub include_u: bool,

    /// Whether to include the `-` gap character.
    pub with_gaps: bool,
}

impl Default for NucleotidesByteSet {
    #[inline]
    fn default() -> Self {
        Self {
            include_n: true,
            include_u: true,
            with_gaps: true,
        }
    }
}

impl NucleotidesByteSet {
    /// A helper function to generate the full character set from a base
    /// character set and the specifications.
    #[inline]
    fn add_to_base_chars(self, chars: &[u8]) -> Vec<u8> {
        let mut chars = chars.to_vec();
        if self.include_n {
            chars.extend_from_slice(b"Nn");
        }
        if self.include_u {
            chars.extend_from_slice(b"Uu");
        }
        if self.with_gaps {
            chars.push(b'-');
        }
        chars
    }
}

/// Specifications for what special characters should be included when
/// generating arbitrary amino acids.
///
/// All options are opt-out. By default, `X` and `-` will be included.
#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
pub struct AminoAcidsByteSet {
    /// Whether to include `X` in both cases.
    pub include_x: bool,

    /// Whether to include the `-` gap character.
    pub with_gaps: bool,
}

impl Default for AminoAcidsByteSet {
    fn default() -> Self {
        Self {
            include_x: true,
            with_gaps: true,
        }
    }
}

impl AminoAcidsByteSet {
    /// A helper function to generate the full character set from a base
    /// character set and the specifications.
    #[inline]
    fn add_to_base_chars(self, chars: &[u8]) -> Vec<u8> {
        let mut chars = chars.to_vec();
        if self.include_x {
            chars.extend_from_slice(b"Xx");
        }
        if self.with_gaps {
            chars.push(b'-');
        }
        chars
    }
}

impl<'a> ArbitrarySpecs<'a> for ByteSet {
    type Output = u8;

    #[inline]
    fn make_arbitrary(&self, u: &mut Unstructured<'a>) -> Result<Self::Output> {
        match self {
            ByteSet::Any => u.arbitrary(),
            ByteSet::Ascii => u.int_in_range(0..=127),
            ByteSet::AsciiGraphicOrSpace => u.int_in_range(b' '..=b'~'),
            ByteSet::AsciiGraphic => u.int_in_range(b'!'..=b'~'),
            ByteSet::AsciiGraphicBashSafe => u
                .choose(b"%+,-./0123456789:@ABCDEFGHIJKLMNOPQRSTUVWXYZ^_abcdefghijklmnopqrstuvwxyz")
                .copied(),
            ByteSet::NucleotidesCanonical(set) => {
                let chars = set.add_to_base_chars(b"ACGTacgt");
                u.choose(&chars).copied()
            }
            ByteSet::NucleotidesIupac(set) => {
                let chars = set.add_to_base_chars(b"ACGTURYSWKMBDHVacgturyswkmbdhv");
                u.choose(&chars).copied()
            }
            ByteSet::AminoAcidsCanonical(set) => {
                let chars = set.add_to_base_chars(b"ACDEFGHIKLMNPQRSTVWYacdefghiklmnpqrstvwy");
                u.choose(&chars).copied()
            }
            ByteSet::Custom(items) => u.choose(items).copied(),
        }
    }
}