entity_buffer_manager.rs

entity_buffer_manager.rs

use std::ops::Range;
use log::warn;
use wgpu::naga::{FastIndexMap, FastIndexSet};

#[derive(Debug, PartialOrd, Hash, PartialEq, Eq, Clone)]
pub struct DirtyIndexRange {
  start: usize,
  until: usize,
}

impl DirtyIndexRange {
  pub fn one(index: usize) -> Self {
    Self {
      start: index,
      until: index + 1,
    }
  }

  pub fn range(start: usize, until: usize) -> Self {
    Self { start, until }
  }
}

#[derive(Debug, PartialOrd, Hash, PartialEq, Eq, Copy, Clone)]
pub struct EntityId(u32);

struct EntityDataEntry {
  entity_id: EntityId,
  data: Vec<u8>,
  cumulative_data_len: usize,
}

impl EntityDataEntry {
  fn offset(&self) -> usize {
    self.cumulative_data_len - self.data.len()
  }
}

struct EntityBufferManager {
  data_map: FastIndexMap<EntityId, EntityDataEntry>,
  device: wgpu::Device,
  queue: wgpu::Queue,
  buffer: wgpu::Buffer,
  dirty_indices: FastIndexSet<DirtyIndexRange>,
}

impl EntityBufferManager {
  // ... (initialization code, now using HashMap for data_map)

  fn last_cumulative_data_len(&self) -> usize {
    self.data_map.last().map(|(_, entry)| entry.cumulative_data_len).unwrap_or(0)
  }

  fn recalculate_cumulative_data_len_from_index(&mut self, start_index: usize) {
    let mut cumulative_data_len = self.data_map.get_index(start_index).map(|(_, entry)| entry.cumulative_data_len).unwrap_or(0);
    for (_, entry) in self.data_map.iter_mut().skip(start_index) {
      entry.cumulative_data_len = cumulative_data_len;
      cumulative_data_len += entry.data.len();
    }
  }

  fn create_entity(&mut self, entity_id: EntityId, data: Vec<u8>) -> &EntityId {
    let cumulative_data_len = self.last_cumulative_data_len() + data.len();
    self.data_map.insert(entity_id, EntityDataEntry { entity_id, data, cumulative_data_len });
    self.dirty_indices.insert(DirtyIndexRange::one(self.data_map.len() - 1));  // mark the newly inserted entity
    &self.data_map[&entity_id].entity_id
  }

  fn update_entity(&mut self, entity_id: &EntityId, data: Vec<u8>) {
    let next_entity_offset = self.data_map.len();
    let Some((index, _, entry)) = self.data_map.get_full_mut(entity_id) else {
      warn!("Entity {:?} not found in data_map", entity_id);
      return;
    };
    let is_length_different = entry.data.len() != data.len();
    if is_length_different {
      // update the dirty indices to entire range since the current.
      self.dirty_indices.insert(DirtyIndexRange::range(index, next_entity_offset));
    } else {
      // update the dirty indices to the single index.
      self.dirty_indices.insert(DirtyIndexRange::one(index));
    }

    entry.data = data;

    if is_length_different {
      self.recalculate_cumulative_data_len_from_index(index);
    }
  }

  fn delete_entity(&mut self, entity_id: &EntityId) {
    let Some((index, _, entity)) = self.data_map.swap_remove_full(entity_id) else {
      warn!("Entity {:?} not found in data_map", entity_id);
      return;
    };

    let Some((_, replaced_at_element)) = self.data_map.get_index(index) else { return; };
    let is_length_different = entity.data.len() != replaced_at_element.data.len();
    if is_length_different {
      // update the dirty indices to entire range since the current.
      self.dirty_indices.insert(DirtyIndexRange::range(index, self.data_map.len()));
    } else {
      // update the dirty indices to the single index.
      self.dirty_indices.insert(DirtyIndexRange::one(index));
    }

    if is_length_different {
      self.recalculate_cumulative_data_len_from_index(index);
    }
  }

  fn write_entity_data_to_contiguous_byte_buffer(&self, range: Range<usize>, buffer: &mut Vec<u8>) {
    for (_, entry) in self.data_map[range].iter() {
      buffer.extend(&entry.data);
    }
  }

  fn sync(&mut self) {
    let mut staging_buffer_data: Vec<u8> = vec![];
    if !self.dirty_indices.is_empty() {
      // Populate staging buffer with updated data and write to GPU:
      for &DirtyIndexRange { start, until } in &self.dirty_indices {
        self.write_entity_data_to_contiguous_byte_buffer(start..until, &mut staging_buffer_data);
        let entity = &self.data_map[start];
        self.queue.write_buffer(&self.buffer, entity.offset() as wgpu::BufferAddress, &staging_buffer_data);
        staging_buffer_data.clear();
      }
    }

    self.dirty_indices.clear(); // Clear dirty flags after syncing
  }
}