Make spatial grids entities

assets/shaders/plexus_link.wesl

@@ -1,11 +1,4 @@
-struct GridParams {
-    grid_min: vec3<f32>,
-    cell_size: f32,
-    dims_x: u32,
-    dims_y: u32,
-    dims_z: u32,
-    _pad: u32,
-}
+import processing::particles::{Grid, cell_coords, cell_index};
 
 struct Params {
     connection_radius: f32,
@@ -15,30 +8,17 @@ struct Params {
 }
 
 // Auto-bound by apply(): names must be position/color.
-@group(0) @binding(0) var<storage, read>       position:  array<f32>;
-@group(0) @binding(1) var<storage, read>       color:     array<f32>;
-@group(0) @binding(2) var<storage, read_write> edge_pos:  array<f32>;
-@group(0) @binding(3) var<storage, read_write> edge_col:  array<f32>;
-@group(0) @binding(4) var<storage, read_write> indices:   array<u32>;
-@group(0) @binding(5) var<storage, read_write> draw_args: array<atomic<u32>>;
-// Bound by grid.bind().
-@group(0) @binding(6) var<storage, read>       offsets:   array<u32>;
-@group(0) @binding(7) var<storage, read>       sorted:    array<u32>;
-@group(0) @binding(8) var<uniform>             params:    Params;
-@group(0) @binding(9) var<uniform>             gp:        GridParams;
-
-fn cell_coords(p: vec3<f32>, grid_min: vec3<f32>, cell_size: f32, dims: vec3<u32>) -> vec3<i32> {
-    let rel = (p - grid_min) / cell_size;
-    return vec3<i32>(
-        clamp(i32(floor(rel.x)), 0, i32(dims.x) - 1),
-        clamp(i32(floor(rel.y)), 0, i32(dims.y) - 1),
-        clamp(i32(floor(rel.z)), 0, i32(dims.z) - 1),
-    );
-}
-
-fn cell_index(c: vec3<u32>, dims: vec3<u32>) -> u32 {
-    return c.x + c.y * dims.x + c.z * dims.x * dims.y;
-}
+@group(0) @binding(0) var<storage, read>       position:     array<f32>;
+@group(0) @binding(1) var<storage, read>       color:        array<f32>;
+@group(0) @binding(2) var<storage, read_write> edge_pos:     array<f32>;
+@group(0) @binding(3) var<storage, read_write> edge_col:     array<f32>;
+@group(0) @binding(4) var<storage, read_write> indices:      array<u32>;
+@group(0) @binding(5) var<storage, read_write> draw_args:    array<atomic<u32>>;
+// set with link.set(grid=grid)
+@group(0) @binding(6) var<storage, read>       grid_offsets: array<u32>;
+@group(0) @binding(7) var<storage, read>       grid_sorted:  array<u32>;
+@group(0) @binding(8) var<uniform>             params:       Params;
+@group(0) @binding(9) var<uniform>             grid:         Grid;
 
 fn load_pos(i: u32) -> vec3<f32> {
     return vec3<f32>(position[i * 3u], position[i * 3u + 1u], position[i * 3u + 2u]);
@@ -70,26 +50,26 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     let rgb = load_rgb(i);
     let radius = params.connection_radius;
     let r2 = radius * radius;
-    let dims = vec3<u32>(gp.dims_x, gp.dims_y, gp.dims_z);
-    let base = cell_coords(pos, gp.grid_min, gp.cell_size, dims);
+    let dims = grid.dims;
+    let base = cell_coords(pos, grid.origin, grid.cell_size, dims);
 
-    let reach = max(1, i32(ceil(radius / gp.cell_size)));
+    let reach = max(1, i32(ceil(radius / grid.cell_size)));
     let z0 = max(base.z - reach, 0);
-    let z1 = min(base.z + reach, i32(gp.dims_z) - 1);
+    let z1 = min(base.z + reach, i32(grid.dims.z) - 1);
     let y0 = max(base.y - reach, 0);
-    let y1 = min(base.y + reach, i32(gp.dims_y) - 1);
+    let y1 = min(base.y + reach, i32(grid.dims.y) - 1);
     let x0 = max(base.x - reach, 0);
-    let x1 = min(base.x + reach, i32(gp.dims_x) - 1);
+    let x1 = min(base.x + reach, i32(grid.dims.x) - 1);
 
     var emitted = 0u;
     for (var cz = z0; cz <= z1; cz++) {
         for (var cy = y0; cy <= y1; cy++) {
             for (var cx = x0; cx <= x1; cx++) {
                 let cell = cell_index(vec3<u32>(u32(cx), u32(cy), u32(cz)), dims);
-                let start = offsets[cell];
-                let end = offsets[cell + 1u];
+                let start = grid_offsets[cell];
+                let end = grid_offsets[cell + 1u];
                 for (var s = start; s < end; s++) {
-                    let j = sorted[s];
+                    let j = grid_sorted[s];
                     if j <= i { continue; }              // one direction per edge
                     let pj = load_pos(j);
                     let diff = pos - pj;

crates/processing_core/src/error.rs

@@ -60,6 +60,8 @@ pub enum ProcessingError {
     PipelineNotReady(u32),
     #[error("Particles not found")]
     ParticlesNotFound,
+    #[error("Grid not found")]
+    GridNotFound,
     #[error("Font not found")]
     FontNotFound,
     #[error("Font load error: {0}")]

crates/processing_ffi/src/lib.rs

@@ -3442,6 +3442,27 @@ pub unsafe extern "C" fn processing_shader_set_buffer(
     });
 }
 
+/// Binds `{name}_offsets`, `{name}_sorted` and the `{name}` uniform.
+///
+/// # Safety
+/// - `name` must be non-null.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn processing_shader_set_grid(
+    entity: u64,
+    name: *const std::ffi::c_char,
+    grid_id: u64,
+) {
+    error::clear_error();
+    error::check(|| {
+        let name = unsafe { cstr_to_str(name) }?;
+        shader_set(
+            Entity::from_bits(entity),
+            name,
+            ShaderValue::Grid(Entity::from_bits(grid_id)),
+        )
+    });
+}
+
 #[unsafe(no_mangle)]
 pub extern "C" fn processing_compute_dispatch(compute_id: u64, x: u32, y: u32, z: u32) {
     error::clear_error();
@@ -3826,6 +3847,51 @@ pub extern "C" fn processing_particles_apply(particles_id: u64, compute_id: u64)
     });
 }
 
+/// Returns the grid id, or 0 on error.
+#[unsafe(no_mangle)]
+pub extern "C" fn processing_particles_grid_create(
+    particles_id: u64,
+    min_x: f32,
+    min_y: f32,
+    min_z: f32,
+    cell_size: f32,
+    dims_x: u32,
+    dims_y: u32,
+    dims_z: u32,
+) -> u64 {
+    error::clear_error();
+    error::check(|| {
+        let capacity = particles_capacity(Entity::from_bits(particles_id))?;
+        grid_create(
+            GridParams {
+                min: [min_x, min_y, min_z],
+                cell_size,
+                dims: [dims_x, dims_y, dims_z],
+            },
+            capacity,
+        )
+    })
+    .map(|e| e.to_bits())
+    .unwrap_or(0)
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn processing_particles_grid_build(grid_id: u64, position_buf_id: u64) {
+    error::clear_error();
+    error::check(|| {
+        grid_build(
+            Entity::from_bits(grid_id),
+            Entity::from_bits(position_buf_id),
+        )
+    });
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn processing_particles_grid_destroy(grid_id: u64) {
+    error::clear_error();
+    error::check(|| grid_destroy(Entity::from_bits(grid_id)));
+}
+
 #[unsafe(no_mangle)]
 pub extern "C" fn processing_particles_draw(graphics_id: u64, particles_id: u64, geometry_id: u64) {
     error::clear_error();

crates/processing_pyo3/examples/particles_lissajous.py

@@ -73,7 +73,7 @@ def draw():
         line_alpha=LINE_ALPHA,
         max_links=MAX_LINKS,
     )
-    grid.bind(link)
+    link.set(grid=grid)
     p.apply(link)
 
     blend_mode(ALPHA_OVER)

crates/processing_pyo3/examples/particles_plexus.py

@@ -65,7 +65,7 @@ def draw():
     grid.build(p.buffer("position"))
     p.reset_indices()
     link.set(indices=idx, draw_args=args, link_distance=LINK_DIST, max_links=MAX_LINKS)
-    grid.bind(link)
+    link.set(grid=grid)
     p.apply(link)
 
     particles(p, topology="lines")

crates/processing_pyo3/src/compute.rs

@@ -291,8 +291,9 @@ fn shader_value_to_py<'py>(py: Python<'py>, sv: &ShaderValue) -> PyResult<Bound<
         ShaderValue::Texture(_)
         | ShaderValue::Buffer(_)
         | ShaderValue::MeshAttribute(..)
-        | ShaderValue::MeshIndex(_) => Err(PyRuntimeError::new_err(
-            "cannot convert Texture/Buffer/Mesh* to Python value",
+        | ShaderValue::MeshIndex(_)
+        | ShaderValue::Grid(_) => Err(PyRuntimeError::new_err(
+            "cannot convert Texture/Buffer/Mesh*/Grid to Python value",
         )),
     }
 }

crates/processing_pyo3/src/material.rs

@@ -43,6 +43,9 @@ pub(crate) fn py_to_shader_value(value: &Bound<'_, PyAny>) -> PyResult<shader_va
     if let Ok(buf) = value.extract::<PyRef<Buffer>>() {
         return Ok(shader_value::ShaderValue::Buffer(buf.entity));
     }
+    if let Ok(grid) = value.extract::<PyRef<crate::particles::Grid>>() {
+        return Ok(shader_value::ShaderValue::Grid(grid.entity));
+    }
 
     if let Ok(v) = value.extract::<[f32; 4]>() {
         return Ok(shader_value::ShaderValue::Float4(v));

crates/processing_pyo3/src/particles.rs

@@ -117,23 +117,33 @@ fn parse_falloff(s: &str) -> PyResult<u32> {
 
 #[pyclass(unsendable)]
 pub struct Grid {
-    pub(crate) inner: processing_render::particles::grid::Grid,
+    pub(crate) entity: Entity,
 }
 
 #[pymethods]
 impl Grid {
+    /// Opaque id for this object.
+    pub fn id(&self) -> u64 {
+        self.entity.to_bits()
+    }
+
     pub fn build(&self, position: &Buffer) -> PyResult<()> {
-        grid_build(&self.inner, position.entity)
+        grid_build(self.entity, position.entity)
             .map_err(|e| PyRuntimeError::new_err(format!("{e}")))
     }
 
-    pub fn bind(&self, compute: &Compute) -> PyResult<()> {
-        grid_bind(&self.inner, compute.entity).map_err(|e| PyRuntimeError::new_err(format!("{e}")))
+    #[getter]
+    pub fn cell_size(&self) -> PyResult<f32> {
+        Ok(grid_get(self.entity)
+            .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?
+            .params
+            .cell_size)
     }
+}
 
-    #[getter]
-    pub fn cell_size(&self) -> f32 {
-        self.inner.params.cell_size
+impl Drop for Grid {
+    fn drop(&mut self) {
+        let _ = grid_destroy(self.entity);
     }
 }
 
@@ -673,7 +683,7 @@ impl Particles {
                 Some(v) => parse_falloff(&v.extract::<String>()?)?,
                 None => FALLOFF_SMOOTHSTEP,
             };
-            let cell = grid.inner.params.cell_size;
+            let cell = grid.cell_size()?;
             let radius = kw_f32(kwargs, "radius", cell)?.min(cell);
 
             let (out, out_comp) = self.operand(kwargs, "out")?;
@@ -704,7 +714,7 @@ impl Particles {
             };
             particles_gather(
                 self.entity,
-                &grid.inner,
+                grid.entity,
                 a,
                 out,
                 op,
@@ -944,9 +954,9 @@ impl Particles {
             cell_size,
             dims,
         };
-        let inner =
+        let entity =
             grid_create(params, capacity).map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
-        Ok(Grid { inner })
+        Ok(Grid { entity })
     }
 
     #[pyo3(signature = (grid, **kwargs))]
@@ -955,15 +965,15 @@ impl Particles {
         if let Some(kwargs) = kwargs {
             crate::compute::set_compute_kwargs(flock, kwargs)?;
         }
-        let cell = grid.inner.params.cell_size;
+        let cell = grid.cell_size()?;
         let neighbor_distance = kw_f32(kwargs, "neighbor_distance", cell)?.min(cell);
         compute_set(
             flock,
             "neighbor_distance",
             shader_value::ShaderValue::Float(neighbor_distance),
         )
         .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?;
-        particles_flock(self.entity, flock, &grid.inner)
+        particles_flock(self.entity, flock, grid.entity)
             .map_err(|e| PyRuntimeError::new_err(format!("{e}")))
     }
 

crates/processing_render/shaders/processing/particles.wesl

@@ -10,6 +10,13 @@ fn falloff(d: f32, radius: f32, mode: u32) -> f32 {
     }
 }
 
+// a spatial hash grid, bound by `set(grid=...)` as `grid_offsets`, `grid_sorted` and `grid`
+struct Grid {
+    origin: vec3<f32>,
+    cell_size: f32,
+    dims: vec3<u32>,
+}
+
 fn cell_coords(p: vec3<f32>, grid_min: vec3<f32>, cell_size: f32, dims: vec3<u32>) -> vec3<i32> {
     let rel = (p - grid_min) / cell_size;
     return vec3<i32>(

crates/processing_render/src/lib.rs

@@ -28,7 +28,7 @@ pub use particles::algebra::{
     generate, lookup, map, mix, pack, reduce_components,
 };
 pub use particles::compact::compact;
-pub use particles::grid::{Grid, GridParams, grid_bind, grid_build, grid_create};
+pub use particles::grid::{Grid, GridParams, grid_build, grid_create, grid_destroy, grid_get};
 pub use particles::reduce::{REDUCE_OP_MAX, REDUCE_OP_MIN, REDUCE_OP_SUM, reduce};
 pub use particles::sort::bitonic_sort_by_key;
 pub use particles::{
@@ -2443,7 +2443,7 @@ fn buffer_write_range(
             .get::<compute::Buffer>(entity)
             .ok_or(error::ProcessingError::BufferNotFound)?
             .synced;
-        // next read will refresh 
+        // next read will refresh
         if on_gpu && !synced {
             return Ok(());
         }

crates/processing_render/src/material/custom.rs

@@ -335,7 +335,7 @@ pub(crate) fn apply_reflect_field(
     Err(ProcessingError::UnknownShaderProperty(name.to_string()))
 }
 
-fn apply_field_coerced(field: &mut dyn PartialReflect, value: &dyn PartialReflect) {
+pub(crate) fn apply_field_coerced(field: &mut dyn PartialReflect, value: &dyn PartialReflect) {
     if let Some(coerced) = coerce_numeric(field, value) {
         field.apply(coerced.as_ref());
     } else {
@@ -495,9 +495,10 @@ pub(crate) fn shader_value_to_reflect(value: &ShaderValue) -> Result<Box<dyn Par
         ShaderValue::Texture(_)
         | ShaderValue::Buffer(_)
         | ShaderValue::MeshAttribute(..)
-        | ShaderValue::MeshIndex(_) => {
+        | ShaderValue::MeshIndex(_)
+        | ShaderValue::Grid(_) => {
             return Err(ProcessingError::InvalidArgument(
-                "Texture/Buffer/Mesh* must be bound via set_property, not as a uniform value"
+                "Texture/Buffer/Mesh*/Grid must be bound via set_property, not as a uniform value"
                     .to_string(),
             ));
         }