perfect-postcode/server-rs/src/index.rs

132 lines
3.7 KiB
Rust

/// Grid-based spatial index for fast rectangle queries over property rows.
///
/// Divides the UK bounding box into cells of ~0.01 degrees (~1km),
/// each storing indices of rows whose lat/lon falls within that cell.
pub struct GridIndex {
min_lat: f64,
min_lon: f64,
cell_size: f64,
cols: usize,
rows: usize,
/// cells[row * cols + col] = vec of row indices
cells: Vec<Vec<u32>>,
}
impl GridIndex {
pub fn build(lat: &[f64], lon: &[f64], cell_size: f64) -> Self {
let mut min_lat = f64::INFINITY;
let mut max_lat = f64::NEG_INFINITY;
let mut min_lon = f64::INFINITY;
let mut max_lon = f64::NEG_INFINITY;
for i in 0..lat.len() {
let la = lat[i];
let lo = lon[i];
if la < min_lat {
min_lat = la;
}
if la > max_lat {
max_lat = la;
}
if lo < min_lon {
min_lon = lo;
}
if lo > max_lon {
max_lon = lo;
}
}
min_lat -= cell_size;
min_lon -= cell_size;
max_lat += cell_size;
max_lon += cell_size;
let rows = ((max_lat - min_lat) / cell_size).ceil() as usize + 1;
let cols = ((max_lon - min_lon) / cell_size).ceil() as usize + 1;
tracing::debug!(
rows_grid = rows,
cols_grid = cols,
total_cells = rows * cols,
cell_size,
"Building grid index"
);
let mut cells: Vec<Vec<u32>> = vec![Vec::new(); rows * cols];
for i in 0..lat.len() {
let r = ((lat[i] - min_lat) / cell_size) as usize;
let c = ((lon[i] - min_lon) / cell_size) as usize;
let idx = r * cols + c;
cells[idx].push(i as u32);
}
tracing::debug!("Grid index built");
GridIndex {
min_lat,
min_lon,
cell_size,
cols,
rows,
cells,
}
}
pub fn query(&self, south: f64, west: f64, north: f64, east: f64) -> Vec<u32> {
let (r_min, r_max, c_min, c_max) = self.clamp_bounds(south, west, north, east);
let mut result = Vec::new();
for r in r_min..=r_max {
let row_start = r * self.cols;
for c in c_min..=c_max {
result.extend_from_slice(&self.cells[row_start + c]);
}
}
result
}
/// Iterate all row indices in bounds without allocating a Vec.
#[inline]
pub fn for_each_in_bounds(
&self,
south: f64,
west: f64,
north: f64,
east: f64,
mut f: impl FnMut(u32),
) {
let (r_min, r_max, c_min, c_max) = self.clamp_bounds(south, west, north, east);
for r in r_min..=r_max {
let row_start = r * self.cols;
for c in c_min..=c_max {
for &row_idx in &self.cells[row_start + c] {
f(row_idx);
}
}
}
}
fn clamp_bounds(
&self,
south: f64,
west: f64,
north: f64,
east: f64,
) -> (usize, usize, usize, usize) {
let r_min = ((south - self.min_lat) / self.cell_size) as isize;
let r_max = ((north - self.min_lat) / self.cell_size) as isize;
let c_min = ((west - self.min_lon) / self.cell_size) as isize;
let c_max = ((east - self.min_lon) / self.cell_size) as isize;
let r_min = r_min.max(0) as usize;
let r_max = (r_max.min(self.rows as isize - 1)).max(0) as usize;
let c_min = c_min.max(0) as usize;
let c_max = (c_max.min(self.cols as isize - 1)).max(0) as usize;
(r_min, r_max, c_min, c_max)
}
}