Cord/crates/cord-decompile/src/sparse_grid.rs

use crate::bvh::BVH;
use crate::mesh::{AABB, TriangleMesh, Vec3};
use std::collections::HashMap;

/// Adaptive sparse octree storing SDF values.
/// Only cells near the surface (where the sign changes) are refined.
pub struct SparseGrid {
    pub bounds: AABB,
    pub max_depth: u8,
    pub cells: HashMap<CellKey, CellData>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CellKey {
    pub depth: u8,
    pub x: u32,
    pub y: u32,
    pub z: u32,
}

#[derive(Debug, Clone)]
pub struct CellData {
    pub center: Vec3,
    pub size: f64,
    pub sdf_value: f64,
    pub normal: Vec3,
    pub is_surface: bool,
}

/// A surface sample extracted from the grid — point + normal on the zero isosurface.
#[derive(Debug, Clone, Copy)]
pub struct SurfaceSample {
    pub position: Vec3,
    pub normal: Vec3,
    pub cell_key: CellKey,
}

impl SparseGrid {
    pub fn from_mesh(mesh: &TriangleMesh, bvh: &BVH, max_depth: u8) -> Self {
        let padding = mesh.bounds.diagonal() * 0.05;
        let bounds = AABB {
            min: mesh.bounds.min - Vec3::new(padding, padding, padding),
            max: mesh.bounds.max + Vec3::new(padding, padding, padding),
        };

        let mut grid = SparseGrid {
            bounds,
            max_depth,
            cells: HashMap::new(),
        };

        // Seed at depth 0
        grid.subdivide_recursive(mesh, bvh, CellKey { depth: 0, x: 0, y: 0, z: 0 });
        grid
    }

    fn cell_bounds(&self, key: CellKey) -> AABB {
        let divisions = (1u32 << key.depth) as f64;
        let extent = self.bounds.max - self.bounds.min;
        let cell_size = Vec3::new(
            extent.x / divisions,
            extent.y / divisions,
            extent.z / divisions,
        );
        let min = Vec3::new(
            self.bounds.min.x + key.x as f64 * cell_size.x,
            self.bounds.min.y + key.y as f64 * cell_size.y,
            self.bounds.min.z + key.z as f64 * cell_size.z,
        );
        AABB { min, max: min + cell_size }
    }

    fn subdivide_recursive(&mut self, mesh: &TriangleMesh, bvh: &BVH, key: CellKey) {
        let cb = self.cell_bounds(key);
        let center = cb.center();
        let size = cb.diagonal();

        let sdf_value = bvh.signed_distance(mesh, center);
        let normal = sdf_gradient(mesh, bvh, center);

        let is_surface = sdf_value.abs() < size * 0.75;

        self.cells.insert(key, CellData {
            center,
            size,
            sdf_value,
            normal,
            is_surface,
        });

        if key.depth < self.max_depth && is_surface {
            // Refine: subdivide into 8 children
            let child_depth = key.depth + 1;
            for dz in 0..2u32 {
                for dy in 0..2u32 {
                    for dx in 0..2u32 {
                        let child_key = CellKey {
                            depth: child_depth,
                            x: key.x * 2 + dx,
                            y: key.y * 2 + dy,
                            z: key.z * 2 + dz,
                        };
                        self.subdivide_recursive(mesh, bvh, child_key);
                    }
                }
            }
        }
    }

    /// Extract surface samples — cells at max depth that straddle the surface.
    pub fn surface_samples(&self) -> Vec<SurfaceSample> {
        self.cells.iter()
            .filter(|(_, data)| data.is_surface && data.sdf_value.abs() < data.size * 0.5)
            .map(|(key, data)| SurfaceSample {
                position: data.center,
                normal: data.normal,
                cell_key: *key,
            })
            .collect()
    }

    /// Extract leaf cells (deepest level for each spatial region).
    pub fn leaf_cells(&self) -> Vec<(&CellKey, &CellData)> {
        self.cells.iter()
            .filter(|(key, _)| {
                // A cell is a leaf if it has no children in the map
                if key.depth >= self.max_depth {
                    return true;
                }
                let child_depth = key.depth + 1;
                let child_key = CellKey {
                    depth: child_depth,
                    x: key.x * 2,
                    y: key.y * 2,
                    z: key.z * 2,
                };
                !self.cells.contains_key(&child_key)
            })
            .collect()
    }

    pub fn surface_cell_count(&self) -> usize {
        self.cells.values().filter(|d| d.is_surface).count()
    }
}

/// Compute SDF gradient (approximate normal) via central differences.
fn sdf_gradient(mesh: &TriangleMesh, bvh: &BVH, p: Vec3) -> Vec3 {
    let eps = 0.001;
    let dx = bvh.signed_distance(mesh, Vec3::new(p.x + eps, p.y, p.z))
           - bvh.signed_distance(mesh, Vec3::new(p.x - eps, p.y, p.z));
    let dy = bvh.signed_distance(mesh, Vec3::new(p.x, p.y + eps, p.z))
           - bvh.signed_distance(mesh, Vec3::new(p.x, p.y, p.z - eps));
    let dz = bvh.signed_distance(mesh, Vec3::new(p.x, p.y, p.z + eps))
           - bvh.signed_distance(mesh, Vec3::new(p.x, p.y, p.z - eps));
    Vec3::new(dx, dy, dz).normalized()
}