path: root/engine/src/util.rs

use std::fmt::{Debug, Display};

use crate::ecs::util::VecExt;

#[derive(Debug, Clone)]
pub struct MapVec<Key: Ord, Value>
{
    inner: Vec<(Key, Value)>,
}

impl<Key: Ord, Value> MapVec<Key, Value>
{
    pub fn with_capacity(capacity: usize) -> Self
    {
        Self { inner: Vec::with_capacity(capacity) }
    }

    pub fn insert(&mut self, key: Key, value: Value)
    {
        self.inner
            .insert_at_part_pt_by_key((key, value), |(a_key, _)| a_key);
    }

    pub fn insert_mut(&mut self, key: Key, value: Value) -> &mut Value
    {
        let insert_index = self
            .inner
            .partition_point(|(other_key, _other_value)| other_key <= &key);

        &mut self.inner.insert_mut(insert_index, (key, value)).1
    }

    pub fn remove(&mut self, key: Key) -> Option<Value>
    {
        let index = self
            .inner
            .binary_search_by_key(&&key, |(a_key, _)| a_key)
            .ok()?;

        let (_, value) = self.inner.remove(index);

        Some(value)
    }

    pub fn get(&self, key: &Key) -> Option<&Value>
    {
        let index = self
            .inner
            .binary_search_by_key(&key, |(a_key, _)| a_key)
            .ok()?;

        let Some((_, value)) = self.inner.get(index) else {
            unreachable!(); // Reason: Index from binary search cannot be OOB
        };

        Some(value)
    }

    pub fn get_mut(&mut self, key: &Key) -> Option<&mut Value>
    {
        let index = self
            .inner
            .binary_search_by_key(&key, |(a_key, _)| a_key)
            .ok()?;

        let Some((_, value)) = self.inner.get_mut(index) else {
            unreachable!(); // Reason: Index from binary search cannot be OOB
        };

        Some(value)
    }

    pub fn entry(&mut self, key: Key) -> MapVecEntry<'_, Key, Value>
    {
        let index = self
            .inner
            .binary_search_by_key(&&key, |(a_key, _)| a_key)
            .ok();

        MapVecEntry { map: self, key, index }
    }

    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&mut Key, &mut Value)>
    {
        self.inner.iter_mut().map(|(key, value)| (key, value))
    }

    pub fn values(&self) -> impl Iterator<Item = &Value>
    {
        self.inner.iter().map(|(_, value)| value)
    }
}

impl<Key: Ord, Value> FromIterator<(Key, Value)> for MapVec<Key, Value>
{
    fn from_iter<Iter: IntoIterator<Item = (Key, Value)>>(iter: Iter) -> Self
    {
        let mut items = iter.into_iter().collect::<Vec<_>>();

        if !items.is_sorted_by_key(|(key, _)| key) {
            items.sort_by(|(key_a, _), (key_b, _)| key_a.cmp(key_b));
        }

        Self { inner: items }
    }
}

impl<Key: Ord, Value> Default for MapVec<Key, Value>
{
    fn default() -> Self
    {
        Self { inner: Vec::new() }
    }
}

pub struct MapVecEntry<'map, Key: Ord, Value>
{
    map: &'map mut MapVec<Key, Value>,
    key: Key,
    index: Option<usize>,
}

impl<'map, Key: Ord, Value> MapVecEntry<'map, Key, Value>
{
    pub fn or_insert_with(self, func: impl FnOnce() -> Value) -> &'map mut Value
    {
        match self.index {
            Some(index) => {
                let (_, value) = unsafe { self.map.inner.get_unchecked_mut(index) };

                value
            }
            None => self.map.insert_mut(self.key, func()),
        }
    }
}

pub trait OptionExt<T>
{
    /// Substitute for the currently experimental function
    /// [`Option::get_or_try_insert_with`].
    /// See https://github.com/rust-lang/rust/issues/143648
    fn get_or_try_insert_with_fn<Err>(
        &mut self,
        func: impl Fn() -> Result<T, Err>,
    ) -> Result<&mut T, Err>;
}

impl<T> OptionExt<T> for Option<T>
{
    fn get_or_try_insert_with_fn<Err>(
        &mut self,
        func: impl FnOnce() -> Result<T, Err>,
    ) -> Result<&mut T, Err>
    {
        if let None = self {
            *self = Some(func()?);
        }

        Ok(unsafe { self.as_mut().unwrap_unchecked() })
    }
}

pub struct DisplaySlice<'a, Item>
{
    slice: &'a [Item],
}

impl<'a, Item> DisplaySlice<'a, Item>
{
    pub fn new(slice: &'a [Item]) -> Self
    {
        Self { slice }
    }
}

impl<Item> Display for DisplaySlice<'_, Item>
where
    Item: Display,
{
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result
    {
        let Some(first_item) = self.slice.first() else {
            return formatter.write_str("<empty>");
        };

        write!(formatter, "{first_item}")?;

        if self.slice.len() == 1 {
            return Ok(());
        }

        for item in &self.slice[1..self.slice.len() - 1] {
            write!(formatter, ", {item}")?;
        }

        let Some(last_item) = self.slice.last() else {
            unreachable!();
        };

        write!(formatter, " & {last_item}")?;

        Ok(())
    }
}

#[derive(Debug)]
pub struct BitArray<const SIZE: usize, const BITS_PER_ITEM: usize>
{
    inner: [u8; SIZE],
}

impl<const SIZE: usize, const BITS_PER_ITEM: usize> BitArray<SIZE, BITS_PER_ITEM>
{
    const ITEM_MASK: u8 = !(u8::MAX << BITS_PER_ITEM);

    pub fn new() -> Self
    {
        assert!(BITS_PER_ITEM > 1);
        assert!(BITS_PER_ITEM <= 8);
        assert_eq!(BITS_PER_ITEM % 2, 0);

        Self { inner: [0; SIZE] }
    }

    pub fn get(&self, item_index: usize) -> u8
    {
        let bit_index = item_index * BITS_PER_ITEM;

        let byte_index = bit_index / 8;

        let bit_index_in_byte = bit_index - (byte_index * 8);

        (self.inner[byte_index] >> (bit_index_in_byte)) & Self::ITEM_MASK
    }

    #[tracing::instrument(skip_all)]
    pub fn set(&mut self, item_index: usize, item_bits: u8)
    {
        let item_bits = item_bits & Self::ITEM_MASK;

        let bit_index = item_index * BITS_PER_ITEM;

        let byte_index = bit_index / 8;

        let bit_index_in_byte = bit_index - (byte_index * 8);

        tracing::trace!(
            item_bits,
            bit_index,
            byte_index,
            bit_index_in_byte,
            "Setting item bits"
        );

        self.inner[byte_index] &= !(Self::ITEM_MASK << bit_index_in_byte);

        self.inner[byte_index] |= item_bits << bit_index_in_byte;
    }

    pub fn clear(&mut self)
    {
        self.inner.fill(0);
    }

    pub fn iter_occupied(&self) -> BitArrayOccupiedIter<'_, BITS_PER_ITEM>
    {
        BitArrayOccupiedIter {
            inner: self.inner.iter().copied().enumerate(),
            byte: None,
            mask: u8::MAX,
        }
    }

    pub fn bytes_mut(&mut self) -> &mut [u8]
    {
        &mut self.inner
    }
}

impl<const SIZE: usize, const BITS_PER_ITEM: usize> Default
    for BitArray<SIZE, BITS_PER_ITEM>
{
    fn default() -> Self
    {
        Self::new()
    }
}

pub struct BitArrayOccupiedIter<'a, const BITS_PER_ITEM: usize>
{
    inner: std::iter::Enumerate<std::iter::Copied<std::slice::Iter<'a, u8>>>,
    byte: Option<(u8, usize)>,
    mask: u8,
}

impl<const BITS_PER_ITEM: usize> BitArrayOccupiedIter<'_, BITS_PER_ITEM>
{
    const ITEM_MASK: u8 = !(u8::MAX << BITS_PER_ITEM);
}

impl<const BITS_PER_ITEM: usize> Iterator for BitArrayOccupiedIter<'_, BITS_PER_ITEM>
{
    type Item = (usize, u8);

    fn next(&mut self) -> Option<Self::Item>
    {
        let (byte_masked, byte_index, item_bit_index_in_byte) = loop {
            let current = self.byte.and_then(|(byte, byte_index)| {
                let byte_masked = byte & self.mask;

                let lowest_one = byte_masked.lowest_one()?;

                let item_bit_index_in_byte = match lowest_one as usize % BITS_PER_ITEM {
                    0 => lowest_one as usize,
                    remainder => lowest_one as usize - remainder,
                };

                Some((byte_masked, byte_index, item_bit_index_in_byte))
            });

            let Some((byte_masked, byte_index, item_bit_index_in_byte)) = current else {
                let (next_byte_index, next_byte) = self.inner.next()?;

                self.byte = Some((next_byte, next_byte_index));
                self.mask = u8::MAX;

                continue;
            };

            break (byte_masked, byte_index, item_bit_index_in_byte);
        };

        let item_bits = (byte_masked >> (item_bit_index_in_byte)) & Self::ITEM_MASK;

        self.mask &= (!Self::ITEM_MASK) << item_bit_index_in_byte;

        let item_index_in_byte = item_bit_index_in_byte as usize / BITS_PER_ITEM;

        let prev_bytes_item_cnt = (byte_index * 8) / BITS_PER_ITEM;

        let index = prev_bytes_item_cnt + item_index_in_byte;

        Some((index, item_bits))
    }
}

macro_rules! try_option {
    ($expr: expr) => {
        match $expr {
            Ok(value) => value,
            Err(err) => {
                return Some(Err(err.into()));
            }
        }
    };
}

pub(crate) use try_option;

macro_rules! or {
    (($($tt: tt)+) else ($($else_tt: tt)*)) => {
        $($tt)+
    };

    (() else ($($else_tt: tt)*)) => {
        $($else_tt)*
    };
}

pub(crate) use or;

#[macro_export]
macro_rules! expand_map_opt {
    ($in: tt, no_occurance=($($no_occurance: tt)*), occurance=($($occurance: tt)*)) => {
        $($occurance)*
    };

    (, no_occurance=($($no_occurance: tt)*), occurance=($($occurance: tt)*)) => {
        $($no_occurance)*
    };
}

#[macro_export]
macro_rules! builder {
    (
        $(#[doc = $doc: literal])*
        #[builder(
            name = $builder_name: ident
            $(, derives = ($($builder_derive: ident),+))?
        )]
        $(#[$attr: meta])*
        $visibility: vis struct $name: ident
        {
            $(
                $(#[doc = $field_doc: literal])*
                $(#[builder(skip_generate_fn$($field_skip_generate_fn: tt)?)])?
                $field_visibility: vis $field: ident: $field_type: ty,
            )*
        }
    ) => {
        $(#[doc = $doc])*
        $(#[$attr])*
        $visibility struct $name
        {
            $(
                $(#[doc = $field_doc])*
                $field_visibility $field: $field_type,
            )*
        }

        $(#[derive($($builder_derive),+)])?
        $visibility struct $builder_name
        {
            $(
                $field: $field_type,
            )*
        }

        impl $builder_name
        {
            $(
                $crate::expand_map_opt!(
                    $(true $($field_skip_generate_fn)?)?,
                    no_occurance=(
                        #[must_use]
                        $visibility fn $field(mut self, $field: $field_type) -> Self
                        {
                            self.$field = $field;
                            self
                        }
                    ),
                    occurance=()
                );
            )*

            #[must_use]
            $visibility const fn build(self) -> $name
            {
                $name {
                    $(
                        $field: self.$field,
                    )*
                }
            }
        }

        impl From<$name> for $builder_name
        {
            #[allow(unused_variables)]
            fn from(built: $name) -> Self
            {
                Self {
                    $(
                        $field: built.$field,
                    )*
                }
            }
        }
    };
}