1 use super::common::GearId;

     2 use std::{

     3     any::TypeId,

     4     fmt::{Debug, Error, Formatter},

     5     marker::PhantomData,

     6     mem::{align_of, size_of, MaybeUninit},

     7     num::NonZeroU16,

     8     ptr::{copy_nonoverlapping, null_mut, NonNull},

     9     slice,

    10 };

    11 

    12 pub trait TypeTuple: Sized {

    13     fn get_types(types: &mut Vec<TypeId>);

    14 }

    15 

    16 impl TypeTuple for () {

    17     fn get_types(_types: &mut Vec<TypeId>) {}

    18 }

    19 

    20 impl<T: 'static> TypeTuple for &T {

    21     fn get_types(types: &mut Vec<TypeId>) {

    22         types.push(TypeId::of::<T>());

    23     }

    24 }

    25 

    26 pub trait TypeIter: TypeTuple {

    27     unsafe fn iter<F: FnMut(GearId, Self)>(slices: &[*mut u8], count: usize, f: F);

    28 }

    29 

    30 macro_rules! type_tuple_impl {

    31     ($($n: literal: $t: ident),+) => {

    32         impl<$($t: 'static),+> TypeTuple for ($(&$t),+,) {

    33             fn get_types(types: &mut Vec<TypeId>) {

    34                 $(types.push(TypeId::of::<$t>()));+

    35             }

    36         }

    37 

    38         impl<$($t: 'static),+> TypeIter for ($(&$t),+,) {

    39             unsafe fn iter<F: FnMut(GearId, Self)>(slices: &[*mut u8], count: usize, mut f: F) {

    40                 for i in 0..count {

    41                     f(*(*slices.get_unchecked(0) as *const GearId).add(i),

    42                       ($(&*(*slices.get_unchecked($n + 1) as *mut $t).add(i)),+,));

    43                 }

    44             }

    45         }

    46 

    47         impl<$($t: 'static),+> TypeTuple for ($(&mut $t),+,) {

    48             fn get_types(types: &mut Vec<TypeId>) {

    49                 $(types.push(TypeId::of::<$t>()));+

    50             }

    51         }

    52 

    53         impl<$($t: 'static),+> TypeIter for ($(&mut $t),+,) {

    54             unsafe fn iter<F: FnMut(GearId, Self)>(slices: &[*mut u8], count: usize, mut f: F) {

    55                 for i in 0..count {

    56                     f(*(*slices.get_unchecked(0) as *const GearId).add(i),

    57                       ($(&mut *(*slices.get_unchecked($n + 1) as *mut $t).add(i)),+,));

    58                 }

    59             }

    60         }

    61     }

    62 }

    63 

    64 type_tuple_impl!(0: A);

    65 type_tuple_impl!(0: A, 1: B);

    66 type_tuple_impl!(0: A, 1: B, 2: C);

    67 type_tuple_impl!(0: A, 1: B, 2: C, 3: D);

    68 type_tuple_impl!(0: A, 1: B, 2: C, 3: D, 4: E);

    69 

    70 const BLOCK_SIZE: usize = 32768;

    71 

    72 struct DataBlock {

    73     max_elements: u16,

    74     elements_count: u16,

    75     data: Box<[u8; BLOCK_SIZE]>,

    76     component_blocks: [Option<NonNull<u8>>; 64],

    77     element_sizes: Box<[u16]>,

    78 }

    79 

    80 impl Unpin for DataBlock {}

    81 

    82 impl Debug for DataBlock {

    83     fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {

    84         write!(

    85             f,

    86             "Block ({}/{}) {{\n",

    87             self.elements_count, self.max_elements

    88         )?;

    89         write!(f, "\tIDs: [")?;

    90         let id_slice = unsafe {

    91             slice::from_raw_parts(

    92                 self.data.as_ptr() as *const GearId,

    93                 self.elements_count as usize,

    94             )

    95         };

    96         for gear_id in id_slice {

    97             write!(f, "{}, ", gear_id)?;

    98         }

    99         write!(f, "]\n")?;

   100         for type_index in 0..self.element_sizes.len() {

   101             if let Some(ptr) = self.component_blocks[type_index] {

   102                 write!(f, "\tC{}: [", type_index)?;

   103                 let slice = unsafe {

   104                     slice::from_raw_parts(

   105                         ptr.as_ptr(),

   106                         (self.elements_count * self.element_sizes[type_index]) as usize,

   107                     )

   108                 };

   109                 for byte in slice {

   110                     write!(f, "{}, ", byte)?;

   111                 }

   112                 write!(f, "]\n")?;

   113             }

   114         }

   115         write!(f, "}}\n")

   116     }

   117 }

   118 

   119 impl DataBlock {

   120     fn new(mask: u64, element_sizes: &[u16], element_alignments: &[u8]) -> Self {

   121         let total_padding: usize = element_alignments.iter().map(|x| *x as usize).sum();

   122         let total_size: u16 = element_sizes

   123             .iter()

   124             .enumerate()

   125             .filter(|(i, _)| mask & (1 << *i as u64) != 0)

   126             .map(|(_, size)| *size)

   127             .sum();

   128         let max_elements =

   129             ((BLOCK_SIZE - total_padding) / (total_size as usize + size_of::<GearId>())) as u16;

   130 

   131         //ensure the block memory is aligned to GearId

   132         let tmp_data: Box<[GearId; BLOCK_SIZE / size_of::<GearId>()]> =

   133             Box::new(unsafe { MaybeUninit::uninit().assume_init() });

   134         let mut data: Box<[u8; BLOCK_SIZE]> =

   135             unsafe { Box::from_raw(Box::into_raw(tmp_data) as *mut [u8; BLOCK_SIZE]) };

   136 

   137         let mut blocks = [None; 64];

   138         let mut address = unsafe {

   139             data.as_mut_ptr()

   140                 .add(size_of::<GearId>() * max_elements as usize)

   141         };

   142 

   143         for i in 0..element_sizes.len() {

   144             if mask & (1 << i as u64) != 0 {

   145                 unsafe {

   146                     address = address.add(address.align_offset(element_alignments[i] as usize));

   147                     blocks[i] = Some(NonNull::new_unchecked(address));

   148                     address = address.add(element_sizes[i] as usize * max_elements as usize)

   149                 };

   150             }

   151         }

   152 

   153         Self {

   154             elements_count: 0,

   155             max_elements,

   156             data,

   157             component_blocks: blocks,

   158             element_sizes: Box::from(element_sizes),

   159         }

   160     }

   161 

   162     fn gear_ids(&self) -> &[GearId] {

   163         unsafe {

   164             slice::from_raw_parts(

   165                 self.data.as_ptr() as *const GearId,

   166                 self.max_elements as usize,

   167             )

   168         }

   169     }

   170 

   171     fn gear_ids_mut(&mut self) -> &mut [GearId] {

   172         unsafe {

   173             slice::from_raw_parts_mut(

   174                 self.data.as_mut_ptr() as *mut GearId,

   175                 self.max_elements as usize,

   176             )

   177         }

   178     }

   179 

   180     fn is_full(&self) -> bool {

   181         self.elements_count == self.max_elements

   182     }

   183 }

   184 

   185 #[derive(Clone, Copy, Debug, Default)]

   186 struct LookupEntry {

   187     index: Option<NonZeroU16>,

   188     block_index: u16,

   189 }

   190 

   191 impl LookupEntry {

   192     fn new(block_index: u16, index: u16) -> Self {

   193         Self {

   194             index: unsafe { Some(NonZeroU16::new_unchecked(index + 1)) },

   195             block_index,

   196         }

   197     }

   198 }

   199 

   200 #[derive(Copy, Clone, Eq, PartialEq, Debug)]

   201 struct BlockMask {

   202     type_mask: u64,

   203     tag_mask: u64,

   204 }

   205 

   206 impl BlockMask {

   207     #[inline]

   208     fn new(type_mask: u64, tag_mask: u64) -> Self {

   209         Self {

   210             type_mask,

   211             tag_mask,

   212         }

   213     }

   214 

   215     #[inline]

   216     fn with_type(&self, type_bit: u64) -> Self {

   217         Self::new(self.type_mask | type_bit, self.tag_mask)

   218     }

   219 

   220     #[inline]

   221     fn with_tag(&self, tag_bit: u64) -> Self {

   222         Self::new(self.type_mask, self.tag_mask | tag_bit)

   223     }

   224 }

   225 

   226 pub struct GearDataManager {

   227     types: Vec<TypeId>,

   228     tags: Vec<TypeId>,

   229     blocks: Vec<DataBlock>,

   230     block_masks: Vec<BlockMask>,

   231     element_sizes: Box<[u16; 64]>,

   232     element_alignments: Box<[u8; 64]>,

   233     lookup: Box<[LookupEntry]>,

   234 }

   235 

   236 impl GearDataManager {

   237     pub fn new() -> Self {

   238         Self {

   239             types: Vec::with_capacity(64),

   240             tags: Vec::with_capacity(64),

   241             blocks: vec![],

   242             block_masks: vec![],

   243             element_sizes: Box::new([0; 64]),

   244             element_alignments: Box::new([0; 64]),

   245             lookup: vec![LookupEntry::default(); u16::max_value() as usize].into_boxed_slice(),

   246         }

   247     }

   248 

   249     #[inline]

   250     fn get_type_index<T: 'static>(&self) -> Option<usize> {

   251         let type_id = TypeId::of::<T>();

   252         self.types.iter().position(|id| *id == type_id)

   253     }

   254 

   255     #[inline]

   256     fn get_tag_index<T: 'static>(&self) -> Option<usize> {

   257         let type_id = TypeId::of::<T>();

   258         self.tags.iter().position(|id| *id == type_id)

   259     }

   260 

   261     fn move_between_blocks(

   262         &mut self,

   263         src_block_index: u16,

   264         src_index: u16,

   265         dest_block_index: u16,

   266     ) -> u16 {

   267         debug_assert!(src_block_index != dest_block_index);

   268         let src_mask = self.block_masks[src_block_index as usize];

   269         let dest_mask = self.block_masks[dest_block_index as usize];

   270         debug_assert!(src_mask.type_mask & dest_mask.type_mask == src_mask.type_mask);

   271 

   272         let src_block = &self.blocks[src_block_index as usize];

   273         let dest_block = &self.blocks[dest_block_index as usize];

   274         debug_assert!(src_index < src_block.elements_count);

   275         debug_assert!(!dest_block.is_full());

   276 

   277         let dest_index = dest_block.elements_count;

   278         for i in 0..self.types.len() {

   279             if src_mask.type_mask & (1 << i as u64) != 0 {

   280                 let size = self.element_sizes[i];

   281                 let src_ptr = src_block.component_blocks[i].unwrap().as_ptr();

   282                 let dest_ptr = dest_block.component_blocks[i].unwrap().as_ptr();

   283                 unsafe {

   284                     copy_nonoverlapping(

   285                         src_ptr.add((src_index * size) as usize),

   286                         dest_ptr.add((dest_index * size) as usize),

   287                         size as usize,

   288                     );

   289                     if src_index < src_block.elements_count - 1 {

   290                         copy_nonoverlapping(

   291                             src_ptr.add((size * (src_block.elements_count - 1)) as usize),

   292                             src_ptr.add((size * src_index) as usize),

   293                             size as usize,

   294                         );

   295                     }

   296                 }

   297             }

   298         }

   299 

   300         let src_block = &mut self.blocks[src_block_index as usize];

   301         let gear_id = src_block.gear_ids()[src_index as usize];

   302 

   303         if src_index < src_block.elements_count - 1 {

   304             let relocated_index = src_block.elements_count as usize - 1;

   305             let gear_ids = src_block.gear_ids_mut();

   306             let relocated_id = gear_ids[relocated_index];

   307 

   308             gear_ids[src_index as usize] = relocated_id;

   309             self.lookup[relocated_id.get() as usize - 1] =

   310                 LookupEntry::new(src_block_index, src_index);

   311         }

   312         src_block.elements_count -= 1;

   313 

   314         let dest_block = &mut self.blocks[dest_block_index as usize];

   315         let dest_index = dest_block.elements_count;

   316 

   317         dest_block.gear_ids_mut()[dest_index as usize] = gear_id;

   318         self.lookup[gear_id.get() as usize - 1] = LookupEntry::new(dest_block_index, dest_index);

   319         dest_block.elements_count += 1;

   320         dest_block.elements_count - 1

   321     }

   322 

   323     fn add_to_block<T: Clone>(&mut self, gear_id: GearId, block_index: u16, value: &T) {

   324         debug_assert!(

   325             self.block_masks[block_index as usize]

   326                 .type_mask

   327                 .count_ones()

   328                 == 1

   329         );

   330 

   331         let block = &mut self.blocks[block_index as usize];

   332         debug_assert!(block.elements_count < block.max_elements);

   333 

   334         unsafe {

   335             *(block.component_blocks[0].unwrap().as_ptr() as *mut T)

   336                 .add(block.elements_count as usize) = value.clone();

   337         };

   338 

   339         let index = block.elements_count;

   340         self.lookup[gear_id.get() as usize - 1] = LookupEntry::new(block_index, index);

   341         block.gear_ids_mut()[index as usize] = gear_id;

   342         block.elements_count += 1;

   343     }

   344 

   345     fn remove_from_block(&mut self, block_index: u16, index: u16) {

   346         let block = &mut self.blocks[block_index as usize];

   347         debug_assert!(index < block.elements_count);

   348 

   349         for (i, size) in self.element_sizes.iter().cloned().enumerate() {

   350             if index < block.elements_count - 1 {

   351                 if let Some(ptr) = block.component_blocks[i] {

   352                     unsafe {

   353                         copy_nonoverlapping(

   354                             ptr.as_ptr()

   355                                 .add((size * (block.elements_count - 1)) as usize),

   356                             ptr.as_ptr().add((size * index) as usize),

   357                             size as usize,

   358                         );

   359                     }

   360                 }

   361             }

   362         }

   363 

   364         self.lookup[block.gear_ids()[index as usize].get() as usize - 1] = LookupEntry::default();

   365         if index < block.elements_count - 1 {

   366             let relocated_index = block.elements_count as usize - 1;

   367             let gear_ids = block.gear_ids_mut();

   368 

   369             gear_ids[index as usize] = gear_ids[relocated_index];

   370             self.lookup[gear_ids[relocated_index].get() as usize - 1] =

   371                 LookupEntry::new(block_index, index);

   372         }

   373         block.elements_count -= 1;

   374     }

   375 

   376     fn write_component<T: Clone>(

   377         &mut self,

   378         block_index: u16,

   379         index: u16,

   380         type_index: usize,

   381         value: &T,

   382     ) {

   383         debug_assert!(type_index < self.types.len());

   384         let block = &mut self.blocks[block_index as usize];

   385         debug_assert!(index < block.elements_count);

   386 

   387         unsafe {

   388             *(block.component_blocks[type_index].unwrap().as_ptr() as *mut T).add(index as usize) =

   389                 value.clone();

   390         };

   391     }

   392 

   393     #[inline]

   394     fn ensure_block(&mut self, mask: BlockMask) -> u16 {

   395         if let Some(index) = self

   396             .block_masks

   397             .iter()

   398             .enumerate()

   399             .position(|(i, m)| *m == mask && !self.blocks[i].is_full())

   400         {

   401             index as u16

   402         } else {

   403             self.blocks.push(DataBlock::new(

   404                 mask.type_mask,

   405                 &self.element_sizes[0..self.types.len()],

   406                 &self.element_alignments[0..self.types.len()],

   407             ));

   408             self.block_masks.push(mask);

   409             (self.blocks.len() - 1) as u16

   410         }

   411     }

   412 

   413     pub fn add<T: Clone + 'static>(&mut self, gear_id: GearId, value: &T) {

   414         if let Some(type_index) = self.get_type_index::<T>() {

   415             let type_bit = 1 << type_index as u64;

   416             let entry = self.lookup[gear_id.get() as usize - 1];

   417 

   418             if let Some(index) = entry.index {

   419                 let mask = self.block_masks[entry.block_index as usize];

   420                 let new_mask = mask.with_type(type_bit);

   421 

   422                 if new_mask != mask {

   423                     let dest_block_index = self.ensure_block(new_mask);

   424                     let dest_index = self.move_between_blocks(

   425                         entry.block_index,

   426                         index.get() - 1,

   427                         dest_block_index,

   428                     );

   429                     self.write_component(dest_block_index, dest_index, type_index, value);

   430                 }

   431             } else {

   432                 let dest_block_index = self.ensure_block(BlockMask::new(type_bit, 0));

   433                 self.add_to_block(gear_id, dest_block_index, value);

   434             }

   435         } else {

   436             panic!("Unregistered type")

   437         }

   438     }

   439 

   440     pub fn add_tag<T: 'static>(&mut self, gear_id: GearId) {

   441         if let Some(tag_index) = self.get_tag_index::<T>() {

   442             let tag_bit = 1 << tag_index as u64;

   443             let entry = self.lookup[gear_id.get() as usize - 1];

   444 

   445             if let Some(index) = entry.index {

   446                 let mask = self.block_masks[entry.block_index as usize];

   447                 let new_mask = mask.with_tag(tag_bit);

   448 

   449                 if new_mask != mask {

   450                     let dest_block_index = self.ensure_block(new_mask);

   451                     self.move_between_blocks(entry.block_index, index.get() - 1, dest_block_index);

   452                 }

   453             } else {

   454                 panic!("Cannot tag a gear with no data")

   455             }

   456         } else {

   457             panic!("Unregistered tag")

   458         }

   459     }

   460 

   461     pub fn remove<T: 'static>(&mut self, gear_id: GearId) {

   462         if let Some(type_index) = self.get_type_index::<T>() {

   463             let entry = self.lookup[gear_id.get() as usize - 1];

   464             if let Some(index) = entry.index {

   465                 let mut dest_mask = self.block_masks[entry.block_index as usize];

   466                 dest_mask.type_mask &= !(1 << type_index as u64);

   467 

   468                 if dest_mask.type_mask == 0 {

   469                     self.remove_from_block(entry.block_index, index.get() - 1);

   470                 } else {

   471                     let dest_block_index = self.ensure_block(dest_mask);

   472                     self.move_between_blocks(entry.block_index, index.get() - 1, dest_block_index);

   473                 }

   474             }

   475         } else {

   476             panic!("Unregistered type")

   477         }

   478     }

   479 

   480     pub fn remove_all(&mut self, gear_id: GearId) {

   481         let entry = self.lookup[gear_id.get() as usize - 1];

   482         if let Some(index) = entry.index {

   483             self.remove_from_block(entry.block_index, index.get() - 1);

   484         }

   485     }

   486 

   487     pub fn register<T: 'static>(&mut self) {

   488         debug_assert!(!std::mem::needs_drop::<T>());

   489         debug_assert!(size_of::<T>() <= u16::max_value() as usize);

   490 

   491         let id = TypeId::of::<T>();

   492         if size_of::<T>() == 0 {

   493             if !self.tags.contains(&id) {

   494                 debug_assert!(self.tags.len() <= 64);

   495                 self.tags.push(id)

   496             }

   497         } else {

   498             if !self.types.contains(&id) {

   499                 debug_assert!(self.types.len() <= 64);

   500                 self.element_sizes[self.types.len()] = size_of::<T>() as u16;

   501                 self.element_alignments[self.types.len()] = align_of::<T>() as u8;

   502                 self.types.push(id);

   503             }

   504         }

   505     }

   506 

   507     fn run_impl<T: TypeIter + 'static, F: FnMut(GearId, T)>(

   508         &mut self,

   509         type_selector: u64,

   510         included_tags: u64,

   511         type_indices: &[i8],

   512         mut f: F,

   513     ) {

   514         let mut slices = vec![null_mut(); type_indices.len() + 1];

   515 

   516         for (block_index, mask) in self.block_masks.iter().enumerate() {

   517             if mask.type_mask & type_selector == type_selector

   518                 && mask.tag_mask & included_tags == included_tags

   519             {

   520                 let block = &mut self.blocks[block_index];

   521                 slices[0] = block.data.as_mut_ptr();

   522 

   523                 for (arg_index, type_index) in type_indices.iter().cloned().enumerate() {

   524                     slices[arg_index as usize + 1] = block.component_blocks[type_index as usize]

   525                         .unwrap()

   526                         .as_ptr()

   527                 }

   528 

   529                 unsafe {

   530                     T::iter(&slices[..], block.elements_count as usize, |id, x| f(id, x));

   531                 }

   532             }

   533         }

   534     }

   535 

   536     pub fn iter<T: TypeIter + 'static>(&mut self) -> DataIterator<T> {

   537         let mut arg_types = Vec::with_capacity(64);

   538         T::get_types(&mut arg_types);

   539         let mut type_indices = vec![-1i8; arg_types.len()];

   540         let mut selector = 0u64;

   541 

   542         for (arg_index, type_id) in arg_types.iter().enumerate() {

   543             match self.types.iter().position(|t| t == type_id) {

   544                 Some(i) if selector & (1 << i as u64) != 0 => panic!("Duplicate type"),

   545                 Some(i) => {

   546                     type_indices[arg_index] = i as i8;

   547                     selector |= 1 << i as u64;

   548                 }

   549                 None => panic!("Unregistered type"),

   550             }

   551         }

   552         DataIterator::new(self, selector, type_indices)

   553     }

   554 }

   555 

   556 pub struct DataIterator<'a, T> {

   557     data: &'a mut GearDataManager,

   558     types: u64,

   559     type_indices: Vec<i8>,

   560     tags: u64,

   561     phantom_types: PhantomData<T>,

   562 }

   563 

   564 impl<'a, T: TypeIter + 'static> DataIterator<'a, T> {

   565     fn new(

   566         data: &'a mut GearDataManager,

   567         types: u64,

   568         type_indices: Vec<i8>,

   569     ) -> DataIterator<'a, T> {

   570         Self {

   571             data,

   572             types,

   573             type_indices,

   574             tags: 0,

   575             phantom_types: PhantomData,

   576         }

   577     }

   578 

   579     pub fn with_tags<U: TypeTuple + 'static>(self) -> Self {

   580         let mut tag_types = Vec::with_capacity(64);

   581         U::get_types(&mut tag_types);

   582         let mut tags = 0;

   583 

   584         for (i, tag) in self.data.tags.iter().enumerate() {

   585             if tag_types.contains(tag) {

   586                 tags |= 1 << i as u64;

   587             }

   588         }

   589         Self { tags, ..self }

   590     }

   591 

   592     #[inline]

   593     pub fn run<F: FnMut(T)>(&mut self, mut f: F) {

   594         self.run_id(|_, x| f(x))

   595     }

   596 

   597     #[inline]

   598     pub fn run_id<F: FnMut(GearId, T)>(&mut self, f: F) {

   599         self.data

   600             .run_impl(self.types, self.tags, &self.type_indices, f);

   601     }

   602 }

   603 

   604 #[cfg(test)]

   605 mod test {

   606     use super::{super::common::GearId, GearDataManager};

   607 

   608     #[derive(Clone)]

   609     struct Datum {

   610         value: u32,

   611     }

   612 

   613     #[derive(Clone)]

   614     struct Tag;

   615 

   616     #[test]

   617     fn single_component_iteration() {

   618         let mut manager = GearDataManager::new();

   619         manager.register::<Datum>();

   620         for i in 1..=5 {

   621             manager.add(GearId::new(i as u16).unwrap(), &Datum { value: i });

   622         }

   623 

   624         let mut sum = 0;

   625         manager.iter().run(|(d,): (&Datum,)| sum += d.value);

   626         assert_eq!(sum, 15);

   627 

   628         manager.iter().run(|(d,): (&mut Datum,)| d.value += 1);

   629         manager.iter().run(|(d,): (&Datum,)| sum += d.value);

   630         assert_eq!(sum, 35);

   631     }

   632 

   633     #[test]

   634     fn tagged_component_iteration() {

   635         let mut manager = GearDataManager::new();

   636         manager.register::<Datum>();

   637         manager.register::<Tag>();

   638         for i in 1..=10 {

   639             let gear_id = GearId::new(i as u16).unwrap();

   640             manager.add(gear_id, &Datum { value: i });

   641         }

   642 

   643         for i in 1..=10 {

   644             let gear_id = GearId::new(i as u16).unwrap();

   645             if i & 1 == 0 {

   646                 manager.add_tag::<Tag>(gear_id);

   647             }

   648         }

   649 

   650         let mut sum = 0;

   651         manager.iter().run(|(d,): (&Datum,)| sum += d.value);

   652         assert_eq!(sum, 55);

   653 

   654         let mut tag_sum = 0;

   655         manager

   656             .iter()

   657             .with_tags::<&Tag>()

   658             .run(|(d,): (&Datum,)| tag_sum += d.value);

   659         assert_eq!(tag_sum, 30);

   660     }

   661 }