use alice_codec::quant::Quantizer;
use alice_codec::rans::{FrequencyTable, RansDecoder, RansEncoder};
use alice_codec::wavelet::Wavelet3D;
use crate::eval::eval_batch_parallel;
use crate::types::SdfNode;
use glam::Vec3;
#[derive(Clone, Debug)]
pub struct SdfVolume {
pub data: Vec<f32>,
pub width: usize,
pub height: usize,
pub depth: usize,
pub origin: Vec3,
pub voxel_size: f32,
}
impl SdfVolume {
#[inline]
pub const fn len(&self) -> usize {
self.width * self.height * self.depth
}
#[inline]
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
#[inline]
pub fn get(&self, x: usize, y: usize, z: usize) -> f32 {
self.data[z * self.height * self.width + y * self.width + x]
}
#[inline]
pub fn world_pos(&self, x: usize, y: usize, z: usize) -> Vec3 {
self.origin
+ Vec3::new(
x as f32 * self.voxel_size,
y as f32 * self.voxel_size,
z as f32 * self.voxel_size,
)
}
}
#[derive(Clone, Debug)]
pub struct VolumeStats {
pub min_distance: f32,
pub max_distance: f32,
pub zero_crossings: usize,
pub total_voxels: usize,
pub surface_ratio: f64,
pub rms: f64,
}
#[derive(Clone, Debug)]
pub struct EncodeConfig {
pub quality: u8,
pub fixed_point_scale: f32,
pub lossless_wavelet: bool,
}
impl Default for EncodeConfig {
fn default() -> Self {
Self {
quality: 75,
fixed_point_scale: 1024.0,
lossless_wavelet: false,
}
}
}
impl EncodeConfig {
pub const fn high_quality() -> Self {
Self {
quality: 95,
fixed_point_scale: 4096.0,
lossless_wavelet: true,
}
}
pub const fn fast() -> Self {
Self {
quality: 50,
fixed_point_scale: 512.0,
lossless_wavelet: false,
}
}
pub const fn lossless() -> Self {
Self {
quality: 100,
fixed_point_scale: 4096.0,
lossless_wavelet: true,
}
}
}
#[derive(Clone, Debug)]
struct EncodedHeader {
width: u32,
height: u32,
depth: u32,
origin: [f32; 3],
voxel_size: f32,
fixed_point_scale: f32,
quality: u8,
flags: u8,
symbol_count: u32,
}
impl EncodedHeader {
const SIZE: usize = 4 * 3 + 4 * 3 + 4 + 4 + 1 + 1 + 4;
const FLAG_LOSSLESS_WAVELET: u8 = 0x01;
const FLAG_RANS_COMPRESSED: u8 = 0x02;
fn to_bytes(&self) -> Vec<u8> {
let mut buf = Vec::with_capacity(Self::SIZE);
buf.extend_from_slice(&self.width.to_le_bytes());
buf.extend_from_slice(&self.height.to_le_bytes());
buf.extend_from_slice(&self.depth.to_le_bytes());
buf.extend_from_slice(&self.origin[0].to_le_bytes());
buf.extend_from_slice(&self.origin[1].to_le_bytes());
buf.extend_from_slice(&self.origin[2].to_le_bytes());
buf.extend_from_slice(&self.voxel_size.to_le_bytes());
buf.extend_from_slice(&self.fixed_point_scale.to_le_bytes());
buf.push(self.quality);
buf.push(self.flags);
buf.extend_from_slice(&self.symbol_count.to_le_bytes());
buf
}
fn from_bytes(data: &[u8]) -> Self {
let read_u32 = |offset: usize| -> u32 {
u32::from_le_bytes([
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3],
])
};
let read_f32 = |offset: usize| -> f32 {
f32::from_le_bytes([
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3],
])
};
Self {
width: read_u32(0),
height: read_u32(4),
depth: read_u32(8),
origin: [read_f32(12), read_f32(16), read_f32(20)],
voxel_size: read_f32(24),
fixed_point_scale: read_f32(28),
quality: data[32],
flags: data[33],
symbol_count: read_u32(34),
}
}
const fn lossless_wavelet(&self) -> bool {
self.flags & Self::FLAG_LOSSLESS_WAVELET != 0
}
const fn rans_compressed(&self) -> bool {
self.flags & Self::FLAG_RANS_COMPRESSED != 0
}
}
pub fn voxelize_sdf(sdf: &SdfNode, origin: Vec3, extent: Vec3, dims: [usize; 3]) -> SdfVolume {
let [w, h, d] = dims;
let size = extent - origin;
let voxel_size_vec = Vec3::new(
if w > 1 {
size.x / (w - 1) as f32
} else {
size.x
},
if h > 1 {
size.y / (h - 1) as f32
} else {
size.y
},
if d > 1 {
size.z / (d - 1) as f32
} else {
size.z
},
);
let voxel_size = voxel_size_vec.x.min(voxel_size_vec.y).min(voxel_size_vec.z);
let total = w * h * d;
let mut points = Vec::with_capacity(total);
for z in 0..d {
for y in 0..h {
for x in 0..w {
points.push(Vec3::new(
(x as f32).mul_add(voxel_size_vec.x, origin.x),
(y as f32).mul_add(voxel_size_vec.y, origin.y),
(z as f32).mul_add(voxel_size_vec.z, origin.z),
));
}
}
}
let data = eval_batch_parallel(sdf, &points);
SdfVolume {
data,
width: w,
height: h,
depth: d,
origin,
voxel_size,
}
}
pub fn voxelize_sdf_uniform(
sdf: &SdfNode,
origin: Vec3,
extent: Vec3,
resolution: usize,
) -> SdfVolume {
voxelize_sdf(sdf, origin, extent, [resolution, resolution, resolution])
}
pub fn encode_sdf_volume(volume: &SdfVolume, config: &EncodeConfig) -> Vec<u8> {
let total = volume.len();
let scale = config.fixed_point_scale;
let mut coeffs: Vec<i32> = volume
.data
.iter()
.map(|&d| (d * scale).round() as i32)
.collect();
let wavelet = if config.lossless_wavelet {
Wavelet3D::cdf53()
} else {
Wavelet3D::cdf97()
};
wavelet.forward(&mut coeffs, volume.width, volume.height, volume.depth);
let quantizer = {
if config.quality >= 95 {
Quantizer::with_dead_zone(1, 0)
} else {
let max_abs = coeffs.iter().map(|c| c.abs()).max().unwrap_or(1).max(1) as f32;
let inv_quality = (100 - config.quality.min(100)) as f32 / 100.0;
let quality_factor = inv_quality * inv_quality * inv_quality;
let step = (1.0 + quality_factor * (max_abs / 4.0)).round() as i32;
let step = step.max(1);
let dead_zone = step / 2;
Quantizer::with_dead_zone(step, dead_zone)
}
};
let mut quantized = vec![0i32; total];
quantizer
.quantize_buffer(&coeffs, &mut quantized)
.expect("quantized buffer is sized to the coefficient buffer");
let raw_bytes: Vec<u8> = quantized
.iter()
.flat_map(|&v| {
let clamped = v.max(i16::MIN as i32).min(i16::MAX as i32) as i16;
clamped.to_le_bytes()
})
.collect();
let mut histogram = [0u32; 256];
for &b in &raw_bytes {
histogram[b as usize] += 1;
}
let max_freq = *histogram.iter().max().unwrap_or(&0);
let use_rans = raw_bytes.len() >= 512 && max_freq < (raw_bytes.len() as u32 / 2);
let payload: Vec<u8> = if use_rans {
let freq_table = FrequencyTable::from_histogram(&histogram);
let mut encoder = RansEncoder::with_capacity(raw_bytes.len());
encoder.encode_symbols(&raw_bytes, &freq_table);
encoder.finish()
} else {
raw_bytes
};
let mut flags = 0u8;
if config.lossless_wavelet {
flags |= EncodedHeader::FLAG_LOSSLESS_WAVELET;
}
if use_rans {
flags |= EncodedHeader::FLAG_RANS_COMPRESSED;
}
let header = EncodedHeader {
width: volume.width as u32,
height: volume.height as u32,
depth: volume.depth as u32,
origin: [volume.origin.x, volume.origin.y, volume.origin.z],
voxel_size: volume.voxel_size,
fixed_point_scale: scale,
quality: config.quality,
flags,
symbol_count: total as u32,
};
let mut output = header.to_bytes();
output.extend_from_slice(&quantizer.step.to_le_bytes());
output.extend_from_slice(&quantizer.dead_zone.to_le_bytes());
if use_rans {
for &count in &histogram {
output.extend_from_slice(&count.to_le_bytes());
}
}
let payload_len = payload.len() as u32;
output.extend_from_slice(&payload_len.to_le_bytes());
output.extend(payload);
output
}
pub fn decode_sdf_volume(data: &[u8]) -> SdfVolume {
let header = EncodedHeader::from_bytes(data);
let (w, h, d) = (
header.width as usize,
header.height as usize,
header.depth as usize,
);
let total = w * h * d;
let scale = header.fixed_point_scale;
let mut offset = EncodedHeader::SIZE;
let step = i32::from_le_bytes([
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3],
]);
offset += 4;
let dead_zone = i32::from_le_bytes([
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3],
]);
offset += 4;
let quantizer = Quantizer::with_dead_zone(step, dead_zone);
let use_rans = header.rans_compressed();
let mut histogram = [0u32; 256];
if use_rans {
for h in &mut histogram {
*h = u32::from_le_bytes([
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3],
]);
offset += 4;
}
}
let payload_len = u32::from_le_bytes([
data[offset],
data[offset + 1],
data[offset + 2],
data[offset + 3],
]) as usize;
offset += 4;
let payload_bytes = &data[offset..offset + payload_len];
let raw_byte_count = total * 2;
let raw_bytes: Vec<u8> = if use_rans {
let freq_table = FrequencyTable::from_histogram(&histogram);
let mut decoder = RansDecoder::new(payload_bytes);
decoder.decode_n(raw_byte_count, &freq_table)
} else {
payload_bytes[..raw_byte_count].to_vec()
};
let quantized: Vec<i32> = raw_bytes
.chunks_exact(2)
.map(|chunk| i16::from_le_bytes([chunk[0], chunk[1]]) as i32)
.collect();
let mut coeffs = vec![0i32; total];
quantizer
.dequantize_buffer(&quantized, &mut coeffs)
.expect("coefficient buffer is sized to the quantized buffer");
let wavelet = if header.lossless_wavelet() {
Wavelet3D::cdf53()
} else {
Wavelet3D::cdf97()
};
wavelet.inverse(&mut coeffs, w, h, d);
let inv_scale = 1.0 / scale;
let float_data: Vec<f32> = coeffs.iter().map(|&c| c as f32 * inv_scale).collect();
SdfVolume {
data: float_data,
width: w,
height: h,
depth: d,
origin: Vec3::new(header.origin[0], header.origin[1], header.origin[2]),
voxel_size: header.voxel_size,
}
}
pub fn volume_stats(volume: &SdfVolume) -> VolumeStats {
let mut min = f32::MAX;
let mut max = f32::MIN;
let mut sum_sq: f64 = 0.0;
for &v in &volume.data {
if v < min {
min = v;
}
if v > max {
max = v;
}
sum_sq = (v as f64).mul_add(v as f64, sum_sq);
}
let mut zero_crossing_count = 0usize;
for z in 0..volume.depth {
for y in 0..volume.height {
for x in 1..volume.width {
let idx = z * volume.height * volume.width + y * volume.width + x;
if volume.data[idx] * volume.data[idx - 1] < 0.0 {
zero_crossing_count += 1;
}
}
}
}
let n = volume.data.len();
let rms = if n > 0 {
(sum_sq / n as f64).sqrt()
} else {
0.0
};
VolumeStats {
min_distance: min,
max_distance: max,
zero_crossings: zero_crossing_count,
total_voxels: n,
surface_ratio: if n > 0 {
zero_crossing_count as f64 / n as f64
} else {
0.0
},
rms,
}
}
pub fn compression_ratio(volume: &SdfVolume, encoded: &[u8]) -> f64 {
let uncompressed = volume.len() * std::mem::size_of::<f32>();
uncompressed as f64 / encoded.len() as f64
}
#[derive(Clone, Debug)]
pub struct CompressResult {
pub encoded: Vec<u8>,
pub stats: VolumeStats,
pub ratio: f64,
pub resolution: [usize; 3],
}
pub fn compress_sdf(
sdf: &SdfNode,
origin: Vec3,
extent: Vec3,
resolution: usize,
config: &EncodeConfig,
) -> CompressResult {
let volume = voxelize_sdf_uniform(sdf, origin, extent, resolution);
let stats = volume_stats(&volume);
let encoded = encode_sdf_volume(&volume, config);
let ratio = compression_ratio(&volume, &encoded);
CompressResult {
encoded,
stats,
ratio,
resolution: [resolution, resolution, resolution],
}
}
pub fn decompress_sdf(data: &[u8]) -> SdfVolume {
decode_sdf_volume(data)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::SdfNode;
#[test]
fn test_voxelize_sphere() {
let sphere = SdfNode::sphere(1.0);
let volume = voxelize_sdf(&sphere, Vec3::splat(-2.0), Vec3::splat(2.0), [8, 8, 8]);
assert_eq!(volume.len(), 512);
assert_eq!(volume.width, 8);
assert_eq!(volume.height, 8);
assert_eq!(volume.depth, 8);
let center_d = volume.get(3, 3, 3);
assert!(
center_d < 0.0,
"Center should be inside sphere, got {}",
center_d
);
let corner_d = volume.get(0, 0, 0);
assert!(
corner_d > 0.0,
"Corner should be outside sphere, got {}",
corner_d
);
}
#[test]
fn test_voxelize_uniform() {
let box_sdf = SdfNode::box3d(1.0, 1.0, 1.0);
let volume = voxelize_sdf_uniform(&box_sdf, Vec3::splat(-2.0), Vec3::splat(2.0), 4);
assert_eq!(volume.len(), 64);
}
#[test]
fn test_volume_stats() {
let sphere = SdfNode::sphere(1.0);
let volume = voxelize_sdf_uniform(&sphere, Vec3::splat(-2.0), Vec3::splat(2.0), 8);
let stats = volume_stats(&volume);
assert!(
stats.min_distance < 0.0,
"Sphere should have negative interior"
);
assert!(
stats.max_distance > 0.0,
"Sphere should have positive exterior"
);
assert!(
stats.zero_crossings > 0,
"Sphere should have surface crossings"
);
assert!(stats.surface_ratio > 0.0 && stats.surface_ratio < 1.0);
assert!(stats.rms > 0.0);
assert_eq!(stats.total_voxels, 512);
}
#[test]
fn test_encode_decode_roundtrip() {
let sphere = SdfNode::sphere(1.0);
let volume = voxelize_sdf_uniform(&sphere, Vec3::splat(-2.0), Vec3::splat(2.0), 8);
let config = EncodeConfig {
quality: 99,
fixed_point_scale: 4096.0,
lossless_wavelet: true, };
let encoded = encode_sdf_volume(&volume, &config);
assert!(!encoded.is_empty(), "Encoded data should not be empty");
let decoded = decode_sdf_volume(&encoded);
assert_eq!(decoded.width, volume.width);
assert_eq!(decoded.height, volume.height);
assert_eq!(decoded.depth, volume.depth);
assert_eq!(decoded.data.len(), volume.data.len());
assert!((decoded.origin - volume.origin).length() < 1e-6);
let mut max_error: f32 = 0.0;
for (orig, recon) in volume.data.iter().zip(decoded.data.iter()) {
let err = (orig - recon).abs();
if err > max_error {
max_error = err;
}
}
assert!(
max_error < 0.5,
"Max reconstruction error {} is too large",
max_error
);
}
#[test]
fn test_encode_preserves_sign() {
let sphere = SdfNode::sphere(1.0);
let volume = voxelize_sdf_uniform(&sphere, Vec3::splat(-2.0), Vec3::splat(2.0), 8);
let config = EncodeConfig::high_quality();
let encoded = encode_sdf_volume(&volume, &config);
let decoded = decode_sdf_volume(&encoded);
let threshold = 0.1;
let mut sign_mismatches = 0;
let mut checked = 0;
for (&orig, &recon) in volume.data.iter().zip(decoded.data.iter()) {
if orig.abs() > threshold {
checked += 1;
if orig.signum() != recon.signum() {
sign_mismatches += 1;
}
}
}
let mismatch_rate = if checked > 0 {
sign_mismatches as f64 / checked as f64
} else {
0.0
};
assert!(
mismatch_rate < 0.05,
"Sign mismatch rate {} is too high ({}/{} voxels)",
mismatch_rate,
sign_mismatches,
checked
);
}
#[test]
fn test_compression_ratio_positive() {
let sphere = SdfNode::sphere(1.0);
let volume = voxelize_sdf_uniform(&sphere, Vec3::splat(-2.0), Vec3::splat(2.0), 8);
let encoded = encode_sdf_volume(&volume, &EncodeConfig::default());
let ratio = compression_ratio(&volume, &encoded);
assert!(ratio > 0.0, "Compression ratio should be positive");
}
#[test]
fn test_world_pos() {
let volume = SdfVolume {
data: vec![0.0; 8],
width: 2,
height: 2,
depth: 2,
origin: Vec3::new(1.0, 2.0, 3.0),
voxel_size: 0.5,
};
let p = volume.world_pos(1, 1, 1);
assert!((p.x - 1.5).abs() < 1e-6);
assert!((p.y - 2.5).abs() < 1e-6);
assert!((p.z - 3.5).abs() < 1e-6);
}
#[test]
fn test_compress_sdf_pipeline() {
let sphere = SdfNode::sphere(1.0);
let result = compress_sdf(
&sphere,
Vec3::splat(-2.0),
Vec3::splat(2.0),
8,
&EncodeConfig::default(),
);
assert!(!result.encoded.is_empty());
assert!(result.ratio > 0.0);
assert_eq!(result.resolution, [8, 8, 8]);
assert_eq!(result.stats.total_voxels, 512);
assert!(result.stats.zero_crossings > 0);
}
#[test]
fn test_compress_decompress_roundtrip() {
let sphere = SdfNode::sphere(1.0);
let result = compress_sdf(
&sphere,
Vec3::splat(-2.0),
Vec3::splat(2.0),
8,
&EncodeConfig::high_quality(),
);
let volume = decompress_sdf(&result.encoded);
assert_eq!(volume.width, 8);
assert_eq!(volume.height, 8);
assert_eq!(volume.depth, 8);
assert_eq!(volume.data.len(), 512);
}
#[test]
fn test_compress_result_fields() {
let sphere = SdfNode::sphere(1.0);
let result = compress_sdf(
&sphere,
Vec3::splat(-2.0),
Vec3::splat(2.0),
4,
&EncodeConfig::fast(),
);
assert_eq!(result.resolution, [4, 4, 4]);
assert_eq!(result.stats.total_voxels, 64);
}
#[test]
fn test_encode_config_presets() {
let fast = EncodeConfig::fast();
assert_eq!(fast.quality, 50);
let hq = EncodeConfig::high_quality();
assert_eq!(hq.quality, 95);
let lossless = EncodeConfig::lossless();
assert!(lossless.lossless_wavelet);
assert_eq!(lossless.quality, 100);
}
}