#version 450
// ============================================================================
// Stage 9: Vertex Polish (Clamped Gradient Descent)
//
// Refines QEF-solved vertex positions toward the true isosurface using
// clamped gradient descent. The step is clamped to 0.5 * voxel_size to
// prevent divergence at CSG min/max boundaries where the gradient flips
// direction at C0 discontinuities.
//
// Input:
// density[17³] — density grid
// voxel_flags[16³] — active voxel flags
// voxel_scan[16³] — exclusive prefix sum of voxel_flags
// vertices[active] — QEF-solved vertex positions (12 floats per vertex)
//
// Output:
// vertices[active].pos — refined positions (floats 0-2 only)
//
// Dispatch: ceil(VOXEL_COUNT / 256), 1, 1
// ============================================================================
#define THREADS 256
#define POLISH_ITERATIONS 3
layout(local_size_x = THREADS, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0, std430) readonly buffer DensityBuf { float density[]; };
layout(set = 0, binding = 4, std430) readonly buffer VoxelFlagBuf { uint voxel_flags[]; };
layout(set = 0, binding = 5, std430) readonly buffer VoxelScanBuf { uint voxel_scan[]; };
layout(set = 0, binding = 6, std430) buffer VertexBuf { float vertices[]; };
// Phase E: Compacted active voxel index list + counts
layout(set = 0, binding = 7, std430) readonly buffer ActiveVoxelsBuf { uint active_voxel_indices[]; };
layout(set = 0, binding = 8, std430) readonly buffer CountBuf { uint counts[]; };
layout(push_constant) uniform PushConstants {
vec3 cell_origin;
float voxel_size;
ivec3 grid_dims; // (16, 16, 16)
uint total_elements;
uint mesh_quality_flags;
uint _pc_pad0;
uint _pc_pad1;
uint _pc_pad2;
};
// Trilinear interpolation from density grid (matches surface_normals.glsl)
float sample_density(vec3 local_pos) {
int dim = grid_dims.x + 1; // 17
// Continuous coordinates in grid space
vec3 gc = local_pos / voxel_size;
gc = clamp(gc, vec3(0.0), vec3(float(dim - 2))); // stay in bounds
ivec3 i0 = ivec3(floor(gc));
vec3 f = gc - vec3(i0);
// 8-corner trilinear
float c000 = density[i0.z * dim * dim + i0.y * dim + i0.x];
float c100 = density[i0.z * dim * dim + i0.y * dim + (i0.x + 1)];
float c010 = density[i0.z * dim * dim + (i0.y + 1) * dim + i0.x];
float c110 = density[i0.z * dim * dim + (i0.y + 1) * dim + (i0.x + 1)];
float c001 = density[(i0.z + 1) * dim * dim + i0.y * dim + i0.x];
float c101 = density[(i0.z + 1) * dim * dim + i0.y * dim + (i0.x + 1)];
float c011 = density[(i0.z + 1) * dim * dim + (i0.y + 1) * dim + i0.x];
float c111 = density[(i0.z + 1) * dim * dim + (i0.y + 1) * dim + (i0.x + 1)];
float c00 = mix(c000, c100, f.x);
float c10 = mix(c010, c110, f.x);
float c01 = mix(c001, c101, f.x);
float c11 = mix(c011, c111, f.x);
float c0 = mix(c00, c10, f.y);
float c1 = mix(c01, c11, f.y);
return mix(c0, c1, f.z);
}
// Central differences gradient (matches surface_normals.glsl)
vec3 gradient(vec3 local_pos) {
float eps = voxel_size * 0.5;
vec3 g;
g.x = sample_density(local_pos + vec3(eps, 0, 0)) - sample_density(local_pos - vec3(eps, 0, 0));
g.y = sample_density(local_pos + vec3(0, eps, 0)) - sample_density(local_pos - vec3(0, eps, 0));
g.z = sample_density(local_pos + vec3(0, 0, eps)) - sample_density(local_pos - vec3(0, 0, eps));
return g;
}
void main() {
uint tid = gl_GlobalInvocationID.x;
if (tid >= counts[2]) return; // active voxel count from prepare_dispatch (counts[1] = quad_emit index_count)
uint gid = active_voxel_indices[tid]; // original voxel grid index
// Get vertex index from prefix sum
uint vi = voxel_scan[gid];
// 12 floats per vertex: pos(3) + normal(3) + corner_signs + seam_flag + material_id + voxel_idx + pad(2)
uint base = vi * 12u;
// Read world-space position
vec3 world_pos = vec3(vertices[base + 0u], vertices[base + 1u], vertices[base + 2u]);
// Convert to local coordinates (relative to cell_origin)
vec3 local = world_pos - cell_origin;
// Compute voxel cell bounds for clamping
// gid maps to (x, y, z) in the voxel grid
uint gz = gid / uint(grid_dims.x * grid_dims.y);
uint gy = (gid / uint(grid_dims.x)) % uint(grid_dims.y);
uint gx = gid % uint(grid_dims.x);
vec3 cell_min = vec3(float(gx), float(gy), float(gz)) * voxel_size;
vec3 cell_max = cell_min + vec3(voxel_size);
// Maximum step size: 0.5 * voxel_size prevents divergence at CSG boundaries
float max_step = 0.5 * voxel_size;
// Clamped gradient descent: 3 iterations
for (int iter = 0; iter < POLISH_ITERATIONS; iter++) {
float d = sample_density(local);
vec3 grad = gradient(local);
float grad_sq = dot(grad, grad);
if (grad_sq < 1e-12) break; // flat region, no improvement possible
// Newton step: d / dot(grad, grad), clamped for safety
float step = clamp(d / grad_sq, -max_step, max_step);
local -= step * grad;
}
// Clamp result to voxel cell bounds
local = clamp(local, cell_min, cell_max);
// Write refined position back in world coordinates
vertices[base + 0u] = local.x + cell_origin.x;
vertices[base + 1u] = local.y + cell_origin.y;
vertices[base + 2u] = local.z + cell_origin.z;
}