malware-modeler 0.0.5

// SPDX-License-Identifier: Apache-2.0

//! The malware modeler uses the presence of n-grams as features for the malware detection models
//! it creates. But ultimately each feature is the presence or absense of an n-gram. So instead of
//! wasting memory in integers for just one or zero, use each bit as a true or false, since the
//! presence or absense of a feature n-gram is a binary situation.
//!
//! The only question is: "what size integer to use?"
//! If many n-grams are expected, then a larger integer size makes sense. But for smaller feature
//! sets a smaller integer size might be better. So this [`BitArray`] type allows for using any
//! unsigned integer type as a backend, since the user knows their data best. Unsigned is best since
//! we'll never need a negative representation, especially since we don't use the integer variable
//! as an integer anyway.
//!
//! Imagine this bit array as a vector if booleans, but compressed to use the least amount of memory
//! as possible since for our purposes, we're likely to have a lot of n-gram features (possibly in
//! the millions).

use std::fmt::Display;
use std::ops::{
    Add, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div, Mul, Not, Rem, Shl,
    Shr, Sub,
};
use std::str::FromStr;

use serde::{Deserialize, Deserializer, Serialize, Serializer};

/// Trait for integer types which can be used at bit storage
pub trait BitStorage<Rhs = Self, Output = Self>:
    Copy
    + Default
    + PartialEq
    + Eq
    + Display
    + std::fmt::Debug
    + std::hash::Hash
    + Add<Rhs, Output = Output>
    + Sub<Rhs, Output = Output>
    + Mul<Rhs, Output = Output>
    + Div<Rhs, Output = Output>
    + Rem<Rhs, Output = Output>
    + std::fmt::Binary
    + Shl<usize, Output = Output>
    + Shr<usize, Output = Output>
    + BitAnd
    + BitAnd<Output = Self>
    + BitAndAssign
    + BitOr
    + BitOrAssign
    + BitXor
    + BitXorAssign
    + Not<Output = Self>
    + Sized
    + Send
    + Sync
{
    /// Number of bits used for this type
    const BITS: usize;

    /// One as this type as a convenience
    const ONE: Self;
}

macro_rules! impl_bit_storage {
    ($($t:ty),*) => {
        $(impl BitStorage for $t {
            const BITS: usize = <$t>::BITS as usize;
            const ONE: Self = 1 as Self;
        })*
    };
}

impl_bit_storage!(u8, u16, u32, u64, u128, usize);

/// Array of booleans as a large bit vector
#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
pub struct BitArray<T: BitStorage> {
    /// Vector of integers holding the bits as integers.
    data: Vec<T>,

    /// Bits remaining in the last integer.
    remaining: usize,
}

impl<T: BitStorage> Display for BitArray<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.data.is_empty() {
            write!(f, "")?;
            return Ok(());
        }

        for bit in self {
            write!(f, "{}", u8::from(bit))?;
        }
        Ok(())
    }
}

impl<T: BitStorage> Serialize for BitArray<T> {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let string = format!("{self}");
        serializer.serialize_str(&string)
    }
}

impl<'de, T: BitStorage> Deserialize<'de> for BitArray<T> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        use serde::de::Error;

        let string = String::deserialize(deserializer)?;
        BitArray::<T>::from_str(&string).map_err(D::Error::custom)
    }
}

impl<T: BitStorage> IntoIterator for BitArray<T> {
    type Item = bool;
    type IntoIter = BitArrayOwnedIterator<T>;

    fn into_iter(self) -> Self::IntoIter {
        Self::IntoIter {
            array: self,
            pos: 0,
        }
    }
}

impl<'a, T: BitStorage + 'a> IntoIterator for &'a BitArray<T> {
    type Item = bool;
    type IntoIter = BitArrayIterator<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        Self::IntoIter {
            array: self,
            pos: 0,
        }
    }
}

impl<T: BitStorage> BitArray<T> {
    /// Create a new bit array where the size is the expected number of bits.
    #[must_use]
    pub fn new(size: usize) -> Self {
        let elements = size.div_ceil(T::BITS);
        let remaining = if size < T::BITS {
            T::BITS - size
        } else {
            size % T::BITS
        };

        Self {
            data: vec![T::default(); elements],
            remaining,
        }
    }

    /// Empty vector with some bytes allocated, where the size of the memory allocated
    /// depends on the underlying integer size used.
    #[must_use]
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            data: Vec::with_capacity(capacity),
            remaining: 0,
        }
    }

    // TODO: add `from_bytes(bytes: &[u8])`

    /// Number of bits in use
    #[inline]
    #[must_use]
    pub const fn len(&self) -> usize {
        self.data.len() * T::BITS - self.remaining
    }

    /// Returns true if the array contains no elements.
    #[inline]
    #[must_use]
    pub const fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Indicates if none of the bits are set
    #[inline]
    #[must_use]
    pub fn all_zeroes(&self) -> bool {
        self.data.iter().all(|x| *x == T::default())
    }

    /// Number of bits available
    #[inline]
    #[must_use]
    pub const fn capacity(&self) -> usize {
        self.data.len() * T::BITS
    }

    /// Iterate over the bits in the array
    #[must_use]
    pub fn iter(&self) -> BitArrayIterator<'_, T> {
        BitArrayIterator {
            array: self,
            pos: 0,
        }
    }

    /// Get the bit value from a given index
    ///
    /// # Panic
    /// This will panic if `index` points to a value beyond the number of bits stored in this array.
    #[must_use]
    pub fn get(&self, index: usize) -> bool
    where
        <T as BitAnd>::Output: PartialEq<T>,
    {
        let (block, bit) = Self::bit_pos(index);
        ((self.data[block] >> bit) & T::ONE) == T::ONE
    }

    /// Set the bit value at a given index
    ///
    /// # Panic
    /// This will panic if `index` points to a value beyond the number of bits stored in this array.
    pub fn set(&mut self, index: usize, value: bool) {
        let (block, bit) = Self::bit_pos(index);
        if value {
            self.data[block] |= T::ONE << bit;
        } else {
            self.data[block] &= !(T::ONE << bit);
        }
    }

    /// Unset the bit at a given index
    ///
    /// # Panic
    /// This will panic if `index` points to a value beyond the number of bits stored in this array.
    pub fn unset(&mut self, index: usize) {
        self.set(index, false);
    }

    /// Push a new bit onto the end of the array.
    pub fn push(&mut self, value: bool) {
        // All integers are used, so add another
        if self.remaining == 0 || self.data.is_empty() {
            self.data.push(T::default());
            self.remaining = T::BITS;
        }

        let index = self.len();
        self.remaining -= 1;
        self.set(index, value);
    }

    /// Clear all bits
    #[inline]
    pub fn clear(&mut self) {
        self.data.iter_mut().for_each(|x| *x = T::default());
    }

    /// Remove all data and set the size to zero
    #[inline]
    pub fn reset(&mut self) {
        self.data.clear();
        self.remaining = 0;
    }

    #[inline]
    const fn bit_pos(idx: usize) -> (usize, usize) {
        let block = idx / T::BITS;
        let bit = idx % T::BITS;
        (block, bit)
    }
}

impl<T: BitStorage> FromStr for BitArray<T> {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut array = BitArray::<T>::with_capacity(s.len());
        for (index, bit) in s.chars().enumerate() {
            match bit {
                '0' => array.push(false),
                '1' => array.push(true),
                _ => return Err(format!("Invalid bit value {bit} at index {index}")),
            }
        }
        Ok(array)
    }
}

/// Iterator over the bits in a bit array holding a reference to the array
#[derive(Clone)]
pub struct BitArrayIterator<'a, T: BitStorage> {
    array: &'a BitArray<T>,
    pos: usize,
}

impl<T: BitStorage> Iterator for BitArrayIterator<'_, T> {
    type Item = bool;

    fn next(&mut self) -> Option<Self::Item> {
        if self.pos >= self.array.len() {
            None
        } else {
            let value = self.array.get(self.pos);
            self.pos += 1;
            Some(value)
        }
    }
}

impl<T: BitStorage> BitArrayIterator<'_, T> {
    /// Reset the iterator to the beginning
    pub fn reset(&mut self) {
        self.pos = 0;
    }
}

/// Iterator over the bits in a bit array owning the array
pub struct BitArrayOwnedIterator<T: BitStorage> {
    array: BitArray<T>,
    pos: usize,
}

impl<T: BitStorage> Iterator for BitArrayOwnedIterator<T> {
    type Item = bool;

    fn next(&mut self) -> Option<Self::Item> {
        if self.pos >= self.array.len() {
            None
        } else {
            let value = self.array.get(self.pos);
            self.pos += 1;
            Some(value)
        }
    }
}

impl<T: BitStorage> BitArrayOwnedIterator<T> {
    /// Reset the iterator to the beginning
    pub fn reset(&mut self) {
        self.pos = 0;
    }
}

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

    #[test]
    fn bit_positioning() {
        assert_eq!(BitArray::<u8>::bit_pos(0), (0, 0));
        assert_eq!(BitArray::<u8>::bit_pos(1), (0, 1));
        assert_eq!(BitArray::<u8>::bit_pos(2), (0, 2));
        assert_eq!(BitArray::<u8>::bit_pos(3), (0, 3));
        assert_eq!(BitArray::<u8>::bit_pos(4), (0, 4));
        assert_eq!(BitArray::<u8>::bit_pos(12), (1, 4));
    }

    #[test]
    fn building() {
        let mut array = BitArray::<u16>::default();
        eprintln!("Array default(): {array:?}");
        assert_eq!("", format!("{array}"));
        array.push(true);
        assert!(!array.is_empty());
        eprintln!("Array pushed true: {array:?}");
        assert_eq!("1", format!("{array}"));
        array.push(false);
        assert_eq!("10", format!("{array}"));
        array.push(true);
        assert_eq!("101", format!("{array}"));
        array.clear();
        assert_eq!("000", format!("{array}"));

        array.reset();
        assert_eq!("", format!("{array}"));
    }

    #[test]
    fn empty() {
        let default = BitArray::<u64>::default();
        assert!(default.is_empty());
        assert_eq!(default.len(), 0);
        assert_eq!(default.capacity(), 0);
        assert_eq!("", format!("{default}"));
    }

    #[test]
    fn with_one_integer_u64() {
        let mut array = BitArray::<u64>::new(10);
        assert_eq!(array.data.len(), 1); // One integer
        assert_eq!(array.remaining, 54); // 54 bits available
        assert!(array.all_zeroes());
        assert_eq!(format!("{array}"), "0000000000");
        let original_len = array.len();
        assert_eq!(original_len, 10);

        array.set(5, true);
        assert_eq!(array.len(), original_len);
        assert!(array.get(5));
        assert_eq!(format!("{array}"), "0000010000");

        array.set(6, true);
        assert_eq!(array.len(), original_len);
        assert!(array.get(5));
        assert_eq!(format!("{array}"), "0000011000");

        array.set(5, false);
        assert!(!array.get(5));

        array.set(6, true);
        assert_eq!(array.len(), original_len);
        assert!(array.get(6));
        assert_eq!(format!("{array}"), "0000001000");

        array.push(true);
        assert_eq!(format!("{array}"), "00000010001");

        let array_string = format!("{array}");
        let array_from_string = BitArray::<u64>::from_str(&array_string).unwrap();
        assert_eq!(array, array_from_string);

        array.clear();
        assert!(array.all_zeroes());
        assert_eq!(array.len(), original_len + 1);
        assert_eq!(format!("{array}"), "00000000000");
    }

    #[test]
    fn with_one_integer_u32() {
        let mut array = BitArray::<u32>::new(10);
        assert_eq!(array.data.len(), 1); // One integer
        assert_eq!(array.remaining, 22); // 22 bits available
        assert!(array.all_zeroes());
        let original_len = array.len();

        array.set(5, true);
        assert_eq!(format!("{array}"), "0000010000");
        assert_eq!(array.len(), original_len);
        assert!(array.get(5));

        let array_string = format!("{array}");
        let array_from_string = BitArray::<u32>::from_str(&array_string).unwrap();
        assert_eq!(array, array_from_string);

        array.set(5, false);
        assert!(!array.get(5));

        array.clear();
        assert!(array.all_zeroes());
        assert_eq!(array.len(), original_len);
    }

    #[test]
    fn with_one_integer_u16() {
        let mut array = BitArray::<u16>::new(10);
        assert_eq!(array.data.len(), 1); // One integer
        assert_eq!(array.remaining, 6); // Six bits available
        assert!(array.all_zeroes());
        let original_len = array.len();

        array.set(5, true);
        assert_eq!(format!("{array}"), "0000010000");
        assert_eq!(array.len(), original_len);
        assert!(array.get(5));

        let array_string = format!("{array}");
        let array_from_string = BitArray::<u16>::from_str(&array_string).unwrap();
        assert_eq!(array, array_from_string);

        array.set(5, false);
        assert!(!array.get(5));

        array.clear();
        assert!(array.all_zeroes());
        assert_eq!(array.len(), original_len);
    }

    #[test]
    fn with_one_integer_u8() {
        let mut array = BitArray::<u8>::new(6);
        assert_eq!(array.data.len(), 1); // One integer
        assert_eq!(array.remaining, 2); // Two bits on the second byte available
        assert!(array.all_zeroes());
        let original_len = array.len();

        array.set(5, true);
        assert_eq!(format!("{array}"), "000001");
        assert_eq!(array.len(), original_len);
        assert!(array.get(5));

        let array_string = format!("{array}");
        let array_from_string = BitArray::<u8>::from_str(&array_string).unwrap();
        assert_eq!(array, array_from_string);

        array.set(5, false);
        assert!(!array.get(5));

        array.clear();
        assert!(array.all_zeroes());
        assert_eq!(array.len(), original_len);
    }

    #[test]
    fn with_several_integers_u64() {
        let mut array = BitArray::<u64>::new(2001);
        assert_eq!(array.data.len(), 32);
        assert_eq!(array.remaining, 17);
        println!("{array}");
        let original_len = array.len();

        array.set(5, true);
        assert_eq!(array.len(), original_len);
        assert!(array.get(5));

        array.set(6, true);
        assert_eq!(array.len(), original_len);
        assert!(array.get(5));

        array.set(5, false);
        assert!(!array.get(5));

        array.set(6, true);
        assert_eq!(array.len(), original_len);
        assert!(array.get(6));

        let array_string = format!("{array}");
        let array_from_string = BitArray::<u64>::from_str(&array_string).unwrap();
        assert_eq!(array, array_from_string);

        array.clear();
        assert_eq!(array.len(), original_len);
    }
}