spatial_hash.rs

use glam::DVec2;
use rustc_hash::{FxHashMap};




#[derive(Clone, Copy, Debug)]
pub struct SpatialHash {
    cell_size: u32,
    conv_factor: f32,
    width: u32
}

impl SpatialHash {
    const MAX: u32 = 20480;

    pub fn create(cell_size: u32) -> Self {
        SpatialHash {
            cell_size,
            conv_factor: 1.0 / cell_size as f32,
            width: Self::MAX / cell_size
        }
    }

    pub fn hash(&self, x: f64, y: f64) -> u32 {
        (x * self.conv_factor as f64) as u32 + (y * self.conv_factor as f64) as u32 * self.width
    }

    pub fn cells_within_bounds(&self, x_in: f64, y_in: f64) -> ([Option<u32>;9], usize) {
        let mut cells: [Option<u32>;9] = [None;9];

        let x = x_in as i32;
        let y = y_in as i32;
        let cell_size = self.cell_size as i32;

        let start_x = x - cell_size;
        let start_y = y - cell_size;
        let end_x = x + cell_size;
        let end_y = y + cell_size;

        let mut index = 0;

        for row_num in (start_x..=end_x).step_by(cell_size as usize) {
            for col_num in (start_y..=end_y).step_by(cell_size as usize) {

                if row_num >= 0 && col_num >= 0 {
                    let hash = self.hash(row_num as f64, col_num as f64);
                    cells[index] = Some(hash);
                    index += 1;
                }
            }
        }
        (cells, index)
    }

}

#[derive(Clone, Copy, Debug)]
pub struct SpatialGridEntry {
    pub position: DVec2,
    pub id: u32
}

impl SpatialGridEntry {
    pub fn create(x: f64, y: f64, id: u32) -> Self {
        SpatialGridEntry {
            position: DVec2 { x, y },
            id
        }
    }
}

impl SpatialGridItem for SpatialGridEntry {
    fn position(&self) -> DVec2 {
        self.position
    }

}

pub trait SpatialGridItem {
    fn position(&self) -> DVec2;
}

pub struct SpatialGrid<T: SpatialGridItem> {
    version: usize,
    pub cells: FxHashMap<u32, GridVec<T>>,
    pub spatial_hash: SpatialHash
}

impl<T: SpatialGridItem> SpatialGrid<T> {
    pub fn create(cell_size: u32) -> Self {
        let spatial_hash = SpatialHash::create(cell_size);
        SpatialGrid {
            version: 0,
            cells: FxHashMap::default(),
            spatial_hash
        }
    }

    pub fn increment_version(&mut self) {
        self.version += 1;
    }

    pub fn insert(&mut self, value: T) {
        let position = value.position();
        let hash = self.spatial_hash.hash(position.x, position.y);
        if let Some(grid_vec) = self.cells.get_mut(&hash) {
            if grid_vec.version != self.version {
                grid_vec.reset(self.version);
            }
            grid_vec.values.push(value);
        } else {
            let mut grid_vec: GridVec<T> = GridVec::new(self.version);
            grid_vec.values.push(value);
            self.cells.insert(hash, grid_vec);
        }
    }

    pub fn values_in_range_with<F>(&self, x: f64, y: f64, max_distance: f64, func: &mut F) -> usize
    where
    F: FnMut(&T)
    {
        let mut visited = 0;
        let (cells, index) = self.spatial_hash.cells_within_bounds(x, y);
        let position = DVec2 {x, y};

        for key in cells.iter().take(index).flatten() {
            if let Some(grid_vec) = self.cells.get(key) {
                if grid_vec.version != self.version {
                    continue;
                }

                for value in &grid_vec.values {
                    visited += 1;
                    if value.position().distance(position) <= max_distance {
                        func(value);
                    }
                }
            }
        }

        visited
    }

    // allocates
    pub fn values_in_range(&self, x: f64, y: f64, max_distance: f64) -> impl Iterator<Item = &T> {
        let (cells, index) = self.spatial_hash.cells_within_bounds(x, y);
        let mut results: Vec<&Vec<T>> = Vec::new();
        let position = DVec2 {x, y};
        
        for key in cells.iter().take(index).flatten() {
            if let Some(grid_vec) = self.cells.get(key) {
                if grid_vec.version != self.version {
                    continue;
                }
                results.push(&grid_vec.values);
            }
        }
       
       results.into_iter().flatten().filter(move |value| value.position().distance(position) <= max_distance)
    }

    pub fn count(&self) -> usize {
        let mut count = 0;
        for cell in self.cells.values() {
            count += cell.values.len();
        }
        count
    }
    
}

pub struct GridVec<T> {
    pub version: usize,
    pub values: Vec<T>
}

impl<T> GridVec<T> {
    pub fn new(version: usize) -> Self {
        GridVec {
            version,
            values: Vec::new()
        }
    }

    pub fn reset(&mut self, version: usize) {
        self.version = version;
        self.values.clear();
    }
}