fix(engine): prevent ignored keyboard key input when main loop is slowHEADmaster

1 files changed, 104 insertions, 13 deletions

diff --git a/engine/src/util.rs b/engine/src/util.rs
index 9d1c887..84ce726 100644
--- a/engine/src/util.rs
+++ b/engine/src/util.rs
@@ -1,6 +1,6 @@
 use std::fmt::{Debug, Display};
 
-use crate::ecs::util::VecExt;
+use crate::ecs::util::{StreamingIterator, VecExt};
 
 #[derive(Debug, Clone)]
 pub struct MapVec<Key: Ord, Value>
@@ -259,16 +259,18 @@ impl<const SIZE: usize, const BITS_PER_ITEM: usize> BitArray<SIZE, BITS_PER_ITEM
         self.inner[byte_index] |= item_bits << bit_index_in_byte;
     }
 
-    pub fn clear(&mut self)
+    pub fn iter_occupied(&self) -> BitArrayOccupiedIter<'_, BITS_PER_ITEM>
     {
-        self.inner.fill(0);
+        BitArrayOccupiedIter {
+            inner: self.inner.iter().enumerate().peekable(),
+            mask: u8::MAX,
+        }
     }
 
-    pub fn iter_occupied(&self) -> BitArrayOccupiedIter<'_, BITS_PER_ITEM>
+    pub fn iter_occupied_mut(&mut self) -> BitArrayOccupiedMutIter<'_, BITS_PER_ITEM>
     {
-        BitArrayOccupiedIter {
-            inner: self.inner.iter().copied().enumerate(),
-            byte: None,
+        BitArrayOccupiedMutIter {
+            inner: self.inner.iter_mut().enumerate().peekable(),
             mask: u8::MAX,
         }
     }
@@ -290,8 +292,7 @@ impl<const SIZE: usize, const BITS_PER_ITEM: usize> Default
 
 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)>,
+    inner: std::iter::Peekable<std::iter::Enumerate<std::slice::Iter<'a, u8>>>,
     mask: u8,
 }
 
@@ -307,8 +308,8 @@ impl<const BITS_PER_ITEM: usize> Iterator for BitArrayOccupiedIter<'_, BITS_PER_
     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 current = self.inner.peek().and_then(|(byte_index, byte)| {
+                let byte_masked = **byte & self.mask;
 
                 let lowest_one = byte_masked.lowest_one()?;
 
@@ -321,9 +322,8 @@ impl<const BITS_PER_ITEM: usize> Iterator for BitArrayOccupiedIter<'_, BITS_PER_
             });
 
             let Some((byte_masked, byte_index, item_bit_index_in_byte)) = current else {
-                let (next_byte_index, next_byte) = self.inner.next()?;
+                let _ = self.inner.next()?;
 
-                self.byte = Some((next_byte, next_byte_index));
                 self.mask = u8::MAX;
 
                 continue;
@@ -346,6 +346,97 @@ impl<const BITS_PER_ITEM: usize> Iterator for BitArrayOccupiedIter<'_, BITS_PER_
     }
 }
 
+pub struct BitArrayOccupiedMutIter<'a, const BITS_PER_ITEM: usize>
+{
+    inner: std::iter::Peekable<std::iter::Enumerate<std::slice::IterMut<'a, u8>>>,
+    mask: u8,
+}
+
+impl<const BITS_PER_ITEM: usize> BitArrayOccupiedMutIter<'_, BITS_PER_ITEM>
+{
+    const ITEM_MASK: u8 = !(u8::MAX << BITS_PER_ITEM);
+}
+
+impl<const BITS_PER_ITEM: usize> StreamingIterator
+    for BitArrayOccupiedMutIter<'_, BITS_PER_ITEM>
+{
+    type Item<'a>
+        = BitArrayItemMut<'a, BITS_PER_ITEM>
+    where
+        Self: 'a;
+
+    fn streaming_next(&mut self) -> Option<Self::Item<'_>>
+    {
+        let (byte_masked, byte_index, item_bit_index_in_byte) = loop {
+            let current = self.inner.peek().and_then(|(byte_index, byte)| {
+                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 _ = self.inner.next()?;
+
+                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(BitArrayItemMut {
+            index,
+            bits: item_bits,
+            byte: self.inner.peek_mut().unwrap().1,
+            bit_index_in_byte: item_bit_index_in_byte,
+        })
+    }
+}
+
+pub struct BitArrayItemMut<'a, const BITS_PER_ITEM: usize>
+{
+    pub index: usize,
+    pub bits: u8,
+    byte: &'a mut u8,
+    bit_index_in_byte: usize,
+}
+
+impl<'a, const BITS_PER_ITEM: usize> BitArrayItemMut<'a, BITS_PER_ITEM>
+{
+    const ITEM_MASK: u8 = !(u8::MAX << BITS_PER_ITEM);
+
+    pub fn clear_and_set(&mut self, new_bits: u8, clear_mask: u8)
+    {
+        let new_bits = new_bits & Self::ITEM_MASK;
+
+        let clear_mask = clear_mask & Self::ITEM_MASK;
+
+        *self.byte &= !(clear_mask << self.bit_index_in_byte);
+
+        *self.byte |= new_bits << self.bit_index_in_byte;
+
+        self.bits = new_bits;
+    }
+}
+
 macro_rules! try_option {
     ($expr: expr) => {
         match $expr {