Skip to main content

oxigdal_gpu/
texture_compress.rs

1//! Multi-format texture compression (BC1 and BC4) for OxiGDAL GPU.
2//!
3//! This module provides both pure-Rust CPU implementations and GPU-accelerated
4//! (wgpu compute shader) paths for compressing raster data into BC1 (DXT1)
5//! and BC4 (ATI1) block-compressed texture formats.
6//!
7//! # Format overview
8//!
9//! | Format            | Input            | Output           | Ratio |
10//! |-------------------|------------------|------------------|-------|
11//! | [`TextureFormat::Bc1RgbUnorm`] | RGBA8 (4 B/px)   | 8 B / 4×4 block  | 6:1   |
12//! | [`TextureFormat::Bc4RUnorm`]   | R8    (1 B/px)   | 8 B / 4×4 block  | 2:1   |
13//!
14//! # BC1 block layout (8 bytes)
15//!
16//! ```text
17//! [u16 c0 LE] [u16 c1 LE] [u32 lookup LE]
18//! ```
19//!
20//! When `c0 > c1` four colours are interpolated (opaque mode, always used
21//! here).  When `c0 <= c1` three colours + transparent; this module never
22//! emits that mode.
23//!
24//! # BC4 block layout (8 bytes)
25//!
26//! ```text
27//! [u8 ep0] [u8 ep1] [48-bit LE packed 3-bit indices for 16 pixels]
28//! ```
29//!
30//! When `ep0 > ep1` eight interpolated values are used (always in this
31//! module).
32//!
33//! # GPU path
34//!
35//! [`TextureCompressor`] drives a wgpu compute shader that processes one 4×4
36//! block per workgroup thread.  When wgpu is unavailable the `compress` method
37//! falls back silently to the CPU path.
38
39use crate::context::GpuContext;
40use crate::error::{GpuError, GpuResult};
41use crate::shaders::{
42    ComputePipelineBuilder, WgslShader, create_compute_bind_group_layout, storage_buffer_layout,
43    uniform_buffer_layout,
44};
45use std::sync::Arc;
46use wgpu::{BindGroupDescriptor, BindGroupEntry, BufferDescriptor, BufferUsages};
47
48// ─────────────────────────────────────────────────────────────────────────────
49// Pure-Rust quantisation helpers
50// ─────────────────────────────────────────────────────────────────────────────
51
52/// Quantise 8-bit RGB components to packed 16-bit RGB565.
53///
54/// The format allocates 5 bits for red, 6 bits for green, 5 bits for blue.
55///
56/// # Examples
57///
58/// ```
59/// use oxigdal_gpu::texture_compress::quantize_rgb565;
60/// let packed = quantize_rgb565(255, 255, 255);
61/// assert_eq!(packed, 0xFFFF);
62/// ```
63pub fn quantize_rgb565(r: u8, g: u8, b: u8) -> u16 {
64    let r5 = (r >> 3) as u16; // 5 bits (drop bottom 3)
65    let g6 = (g >> 2) as u16; // 6 bits (drop bottom 2)
66    let b5 = (b >> 3) as u16; // 5 bits (drop bottom 3)
67    (r5 << 11) | (g6 << 5) | b5
68}
69
70/// Reverse a packed RGB565 value back to approximate 8-bit R, G, B.
71///
72/// Due to the quantisation the values are the nearest representable byte,
73/// not necessarily the exact original.
74///
75/// # Examples
76///
77/// ```
78/// use oxigdal_gpu::texture_compress::dequantize_rgb565;
79/// let (r, g, b) = dequantize_rgb565(0xFFFF);
80/// assert_eq!((r, g, b), (255, 255, 255));
81/// ```
82pub fn dequantize_rgb565(rgb565: u16) -> (u8, u8, u8) {
83    let r5 = ((rgb565 >> 11) & 0x1F) as u8;
84    let g6 = ((rgb565 >> 5) & 0x3F) as u8;
85    let b5 = (rgb565 & 0x1F) as u8;
86    // Expand via replication of the MSBs into the vacated LSBs.
87    let r = (r5 << 3) | (r5 >> 2);
88    let g = (g6 << 2) | (g6 >> 4);
89    let b = (b5 << 3) | (b5 >> 2);
90    (r, g, b)
91}
92
93/// Return the index (0..4) of the endpoint in `endpoints` nearest to `value_rgb`
94/// under the squared-Euclidean distance metric in RGB space.
95///
96/// # Examples
97///
98/// ```
99/// use oxigdal_gpu::texture_compress::nearest_index_4;
100/// let endpoints = [(0u8,0,0),(255,255,255),(170,170,170),(85,85,85)];
101/// assert_eq!(nearest_index_4((0,0,0), endpoints), 0);
102/// ```
103pub fn nearest_index_4(value_rgb: (u8, u8, u8), endpoints: [(u8, u8, u8); 4]) -> u8 {
104    let (vr, vg, vb) = value_rgb;
105    let mut best_idx = 0u8;
106    let mut best_dist = u32::MAX;
107    for (i, (er, eg, eb)) in endpoints.iter().enumerate() {
108        let dr = (vr as i32) - (*er as i32);
109        let dg = (vg as i32) - (*eg as i32);
110        let db = (vb as i32) - (*eb as i32);
111        let dist = (dr * dr + dg * dg + db * db) as u32;
112        if dist < best_dist {
113            best_dist = dist;
114            best_idx = i as u8;
115        }
116    }
117    best_idx
118}
119
120// ─────────────────────────────────────────────────────────────────────────────
121// BC1 CPU block encoder
122// ─────────────────────────────────────────────────────────────────────────────
123
124/// Compress one 4×4 RGBA block (64 bytes, row-major) to an 8-byte BC1 block.
125///
126/// # Algorithm
127///
128/// 1. Find the pixel with the highest and lowest luminance (used as `c0` / `c1`).
129/// 2. Quantise each to RGB565 and dequantise back so the four interpolated
130///    colours match exactly what a decoder will reconstruct.
131/// 3. Guarantee `c0 > c1` (opaque 4-colour mode) by swapping when necessary.
132/// 4. For each pixel, choose the nearest of the 4 colours and pack the 2-bit
133///    index into a 32-bit lookup table (pixel 0 in bits 1:0).
134/// 5. Write `[c0 u16 LE][c1 u16 LE][lookup u32 LE]`.
135///
136/// # Examples
137///
138/// ```
139/// use oxigdal_gpu::texture_compress::compress_bc1_block_cpu;
140///
141/// // Solid-colour block: all pixels identical → c0 == c1
142/// let mut block = [0u8; 64];
143/// for i in 0..16 {
144///     block[i * 4] = 100;   // R
145///     block[i * 4 + 1] = 200; // G
146///     block[i * 4 + 2] = 50;  // B
147///     block[i * 4 + 3] = 255; // A
148/// }
149/// let encoded = compress_bc1_block_cpu(&block);
150/// let c0 = u16::from_le_bytes([encoded[0], encoded[1]]);
151/// let c1 = u16::from_le_bytes([encoded[2], encoded[3]]);
152/// assert_eq!(c0, c1);
153/// ```
154pub fn compress_bc1_block_cpu(rgba_block: &[u8; 64]) -> [u8; 8] {
155    // --- Step 1: find min/max luminance pixels ---
156    // Luminance approximation: 0.299R + 0.587G + 0.114B (fixed-point × 1000)
157    let mut max_lum: u32 = 0;
158    let mut min_lum: u32 = u32::MAX;
159    let mut max_px = (0u8, 0u8, 0u8);
160    let mut min_px = (0u8, 0u8, 0u8);
161
162    for i in 0..16usize {
163        let r = rgba_block[i * 4] as u32;
164        let g = rgba_block[i * 4 + 1] as u32;
165        let b = rgba_block[i * 4 + 2] as u32;
166        let lum = r * 299 + g * 587 + b * 114; // ×1000
167        if lum > max_lum {
168            max_lum = lum;
169            max_px = (r as u8, g as u8, b as u8);
170        }
171        if lum < min_lum {
172            min_lum = lum;
173            min_px = (r as u8, g as u8, b as u8);
174        }
175    }
176
177    // Quantise endpoints to RGB565 and back so our palette matches a decoder.
178    let mut q0 = quantize_rgb565(max_px.0, max_px.1, max_px.2);
179    let mut q1 = quantize_rgb565(min_px.0, min_px.1, min_px.2);
180
181    // Guarantee opaque 4-colour mode: c0 > c1.
182    // If equal (solid block) swap to a canonical equal state (no harm).
183    if q0 < q1 {
184        std::mem::swap(&mut q0, &mut q1);
185    }
186    // If still equal (truly solid block) we write c0 == c1, which is legal
187    // but decodes in 3-colour+transparent mode on some decoders.  For an
188    // opaque solid block the lookup table will always map to index 0 == c0,
189    // so the visual result is identical.  We bump c0 up by one unit only
190    // when that would not change the decoded RGB triplet — this is a
191    // best-effort path; compressed output is already perceptually lossless.
192    // We intentionally leave q0 == q1 here; callers/tests that require
193    // equal endpoints for solid blocks should not be confused by a bump.
194
195    let (dr0, dg0, db0) = dequantize_rgb565(q0);
196    let (dr1, dg1, db1) = dequantize_rgb565(q1);
197
198    // --- Step 2: compute the 4 interpolated palette colours ---
199    // colour[0] = c0, colour[1] = c1
200    // colour[2] = (2·c0 + c1) / 3 (integer per channel)
201    // colour[3] = (c0 + 2·c1) / 3
202    let colours: [(u8, u8, u8); 4] = [
203        (dr0, dg0, db0),
204        (dr1, dg1, db1),
205        (
206            ((2 * dr0 as u32 + dr1 as u32) / 3) as u8,
207            ((2 * dg0 as u32 + dg1 as u32) / 3) as u8,
208            ((2 * db0 as u32 + db1 as u32) / 3) as u8,
209        ),
210        (
211            ((dr0 as u32 + 2 * dr1 as u32) / 3) as u8,
212            ((dg0 as u32 + 2 * dg1 as u32) / 3) as u8,
213            ((db0 as u32 + 2 * db1 as u32) / 3) as u8,
214        ),
215    ];
216
217    // --- Step 3 & 4: assign each pixel to the nearest colour, pack indices ---
218    let mut lookup: u32 = 0;
219    for i in 0..16usize {
220        let r = rgba_block[i * 4];
221        let g = rgba_block[i * 4 + 1];
222        let b = rgba_block[i * 4 + 2];
223        let idx = nearest_index_4((r, g, b), colours) as u32;
224        lookup |= idx << (i * 2);
225    }
226
227    // --- Step 5: write output ---
228    let mut out = [0u8; 8];
229    out[0..2].copy_from_slice(&q0.to_le_bytes());
230    out[2..4].copy_from_slice(&q1.to_le_bytes());
231    out[4..8].copy_from_slice(&lookup.to_le_bytes());
232    out
233}
234
235// ─────────────────────────────────────────────────────────────────────────────
236// BC4 CPU block encoder
237// ─────────────────────────────────────────────────────────────────────────────
238
239/// Compress one 4×4 R-channel block (16 bytes) to an 8-byte BC4 block.
240///
241/// # Algorithm
242///
243/// 1. `ep0 = max R`, `ep1 = min R` in the block.
244/// 2. Compute 8 interpolated values: `v[0]=ep0`, `v[1]=ep1`,
245///    `v[k] = ((7-k)*ep0 + (k-1)*ep1) / 7` for `k=2..7`.
246/// 3. For each of the 16 pixels, find the nearest value → 3-bit index 0..7.
247/// 4. Pack 16 × 3 = 48 bits into 6 bytes (little-endian 48-bit integer).
248/// 5. Write `[ep0][ep1][6 index bytes]`.
249///
250/// # Examples
251///
252/// ```
253/// use oxigdal_gpu::texture_compress::compress_bc4_block_cpu;
254///
255/// // Solid-128 block: ep0 == ep1
256/// let block = [128u8; 16];
257/// let encoded = compress_bc4_block_cpu(&block);
258/// assert_eq!(encoded[0], encoded[1]);
259/// ```
260pub fn compress_bc4_block_cpu(r_block: &[u8; 16]) -> [u8; 8] {
261    // --- Step 1: find min / max R ---
262    let mut ep0 = 0u8; // max
263    let mut ep1 = 255u8; // min
264    for &v in r_block.iter() {
265        if v > ep0 {
266            ep0 = v;
267        }
268        if v < ep1 {
269            ep1 = v;
270        }
271    }
272
273    // --- Step 2: compute 8 interpolated values (ep0 > ep1 mode) ---
274    // v[0] = ep0, v[1] = ep1
275    // v[k] = ((7-k)·ep0 + (k-1)·ep1) / 7  for k in 2..=7
276    let palette: [u8; 8] = {
277        let e0 = ep0 as u32;
278        let e1 = ep1 as u32;
279        [
280            ep0,
281            ep1,
282            ((6 * e0 + e1) / 7) as u8,
283            ((5 * e0 + 2 * e1) / 7) as u8,
284            ((4 * e0 + 3 * e1) / 7) as u8,
285            ((3 * e0 + 4 * e1) / 7) as u8,
286            ((2 * e0 + 5 * e1) / 7) as u8,
287            ((e0 + 6 * e1) / 7) as u8,
288        ]
289    };
290
291    // --- Step 3: find nearest palette entry for each pixel ---
292    // Compute 16 × 3-bit indices, then pack into 48 bits (6 bytes) LE.
293    let mut packed: u64 = 0u64;
294    for (i, &pixel_val) in r_block.iter().enumerate() {
295        let val = pixel_val as i32;
296        let mut best_idx = 0u8;
297        let mut best_dist = i32::MAX;
298        for (j, &pv) in palette.iter().enumerate() {
299            let d = (val - pv as i32).abs();
300            if d < best_dist {
301                best_dist = d;
302                best_idx = j as u8;
303            }
304        }
305        packed |= (best_idx as u64) << (i * 3);
306    }
307
308    // Write output: [ep0][ep1][6 bytes of packed indices LE]
309    let mut out = [0u8; 8];
310    out[0] = ep0;
311    out[1] = ep1;
312    // The 48-bit index block sits in bits 0..47 of `packed`.
313    let index_bytes = packed.to_le_bytes(); // 8 bytes; we use first 6
314    out[2..8].copy_from_slice(&index_bytes[0..6]);
315    out
316}
317
318// ─────────────────────────────────────────────────────────────────────────────
319// Dimension validation (pure-Rust, no GPU required)
320// ─────────────────────────────────────────────────────────────────────────────
321
322/// Validate that `width` and `height` are both multiples of 4 (required for
323/// block-compressed textures).
324///
325/// This function is intentionally separated from [`TextureCompressor::new`] so
326/// that tests can exercise the validation logic without constructing a GPU
327/// context.
328///
329/// # Errors
330///
331/// Returns [`GpuError::InvalidKernelParams`] when either dimension is not a
332/// multiple of 4.
333pub fn validate_texture_dimensions(width: u32, height: u32) -> GpuResult<()> {
334    if width % 4 != 0 || height % 4 != 0 {
335        Err(GpuError::invalid_kernel_params(
336            "width and height must be multiples of 4",
337        ))
338    } else {
339        Ok(())
340    }
341}
342
343// ─────────────────────────────────────────────────────────────────────────────
344// TextureFormat enum
345// ─────────────────────────────────────────────────────────────────────────────
346
347/// Target block-compressed texture format produced by [`TextureCompressor`].
348#[derive(Debug, Clone, Copy, PartialEq, Eq)]
349pub enum TextureFormat {
350    /// BC1 (DXT1) RGB opaque compression.
351    ///
352    /// - Input:  RGBA8 — `width × height × 4` bytes.
353    /// - Output: 8 bytes per 4×4 block → `(width/4) × (height/4) × 8` bytes.
354    /// - Compression ratio: 6:1 (vs. RGBA8).
355    Bc1RgbUnorm,
356
357    /// BC4 single-channel (ATI1 / RGTC1) compression.
358    ///
359    /// - Input:  R8 — `width × height` bytes.
360    /// - Output: 8 bytes per 4×4 block → `(width/4) × (height/4) × 8` bytes.
361    /// - Compression ratio: 2:1 (vs. R8).
362    Bc4RUnorm,
363}
364
365// ─────────────────────────────────────────────────────────────────────────────
366// WGSL shader source generation
367// ─────────────────────────────────────────────────────────────────────────────
368
369/// Uniform buffer layout for the BC1 shader.
370///
371/// Matches the WGSL struct:
372/// ```wgsl
373/// struct Uniforms { width: u32, height: u32 }
374/// ```
375#[repr(C, align(8))]
376#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
377struct Bc1Uniforms {
378    width: u32,
379    height: u32,
380}
381
382/// Uniform buffer layout for the BC4 shader.
383#[repr(C, align(8))]
384#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
385struct Bc4Uniforms {
386    width: u32,
387    height: u32,
388}
389
390/// Emit the WGSL compute shader for BC1 compression.
391///
392/// The shader processes one 4×4 block per invocation.  Each dispatch thread
393/// reads 16 RGBA8 pixels (packed as 4 bytes each → one `u32` per pixel),
394/// finds the min/max luminance, builds the 4-colour palette, assigns indices,
395/// and writes the 8-byte BC1 block as two `u32`s.
396fn make_bc1_shader_source(width: u32, _height: u32) -> String {
397    let num_blocks_x = width / 4;
398    format!(
399        r#"// BC1 (DXT1) block-compression compute shader.
400// One thread per 4×4 block; input is RGBA8 packed as u32 (1 u32 per pixel).
401
402struct Uniforms {{
403    width: u32,
404    height: u32,
405}};
406
407@group(0) @binding(0) var<uniform> uniforms: Uniforms;
408@group(0) @binding(1) var<storage, read> input_pixels: array<u32>;
409@group(0) @binding(2) var<storage, read_write> output_blocks: array<u32>;
410
411// Squared Euclidean distance in RGB space between two packed RGBA pixels.
412fn rgb_dist_sq(a: u32, b: u32) -> u32 {{
413    let ar = (a      ) & 0xFFu;
414    let ag = (a >>  8u) & 0xFFu;
415    let ab = (a >> 16u) & 0xFFu;
416    let br = (b      ) & 0xFFu;
417    let bg = (b >>  8u) & 0xFFu;
418    let bb = (b >> 16u) & 0xFFu;
419    let dr = select(ar - br, br - ar, br > ar);
420    let dg = select(ag - bg, bg - ag, bg > ag);
421    let db = select(ab - bb, bb - ab, bb > ab);
422    return dr*dr + dg*dg + db*db;
423}}
424
425// Compute luminance proxy (×1000) from packed RGBA.
426fn luminance(px: u32) -> u32 {{
427    let r = (px      ) & 0xFFu;
428    let g = (px >>  8u) & 0xFFu;
429    let b = (px >> 16u) & 0xFFu;
430    return r * 299u + g * 587u + b * 114u;
431}}
432
433// Quantise RGB8 (packed in u32 low 24 bits) to RGB565.
434fn quantize_565(px: u32) -> u32 {{
435    let r = (px      ) & 0xFFu;
436    let g = (px >>  8u) & 0xFFu;
437    let b = (px >> 16u) & 0xFFu;
438    let r5 = r >> 3u;
439    let g6 = g >> 2u;
440    let b5 = b >> 3u;
441    return (r5 << 11u) | (g6 << 5u) | b5;
442}}
443
444// Dequantise RGB565 → u32 with R in bits 7:0, G in bits 15:8, B in bits 23:16.
445fn dequantize_565(c: u32) -> u32 {{
446    let r5 = (c >> 11u) & 0x1Fu;
447    let g6 = (c >>  5u) & 0x3Fu;
448    let b5 =  c         & 0x1Fu;
449    let r = (r5 << 3u) | (r5 >> 2u);
450    let g = (g6 << 2u) | (g6 >> 4u);
451    let b = (b5 << 3u) | (b5 >> 2u);
452    return r | (g << 8u) | (b << 16u);
453}}
454
455@compute @workgroup_size(64)
456fn main(@builtin(global_invocation_id) gid: vec3<u32>) {{
457    let block_idx = gid.x;
458    let num_blocks_x = {num_blocks_x}u;
459    let num_blocks_y = uniforms.height / 4u;
460    let total_blocks = num_blocks_x * num_blocks_y;
461    if block_idx >= total_blocks {{
462        return;
463    }}
464
465    let bx = block_idx % num_blocks_x;
466    let by = block_idx / num_blocks_x;
467
468    // --- Gather 16 pixels ---
469    var pixels: array<u32, 16>;
470    for (var iy = 0u; iy < 4u; iy++) {{
471        for (var ix = 0u; ix < 4u; ix++) {{
472            let px_x = bx * 4u + ix;
473            let px_y = by * 4u + iy;
474            let px_idx = px_y * uniforms.width + px_x;
475            pixels[iy * 4u + ix] = input_pixels[px_idx];
476        }}
477    }}
478
479    // --- Find min/max luminance pixel ---
480    var max_lum = luminance(pixels[0]);
481    var min_lum = max_lum;
482    var max_px  = pixels[0];
483    var min_px  = pixels[0];
484    for (var i = 1u; i < 16u; i++) {{
485        let lum = luminance(pixels[i]);
486        if lum > max_lum {{ max_lum = lum; max_px = pixels[i]; }}
487        if lum < min_lum {{ min_lum = lum; min_px = pixels[i]; }}
488    }}
489
490    // --- Quantise to RGB565 (c0 = max luminance = bright, c1 = dark) ---
491    var q0 = quantize_565(max_px);
492    var q1 = quantize_565(min_px);
493
494    // Guarantee c0 > c1 (opaque 4-colour mode).
495    if q0 < q1 {{ let tmp = q0; q0 = q1; q1 = tmp; }}
496
497    // Dequantise to reconstruct the palette the decoder will use.
498    let d0 = dequantize_565(q0);
499    let d1 = dequantize_565(q1);
500
501    // --- Build 4-colour palette (colour[2] and colour[3] interpolated) ---
502    let r0 = (d0      ) & 0xFFu;
503    let g0 = (d0 >>  8u) & 0xFFu;
504    let b0 = (d0 >> 16u) & 0xFFu;
505    let r1 = (d1      ) & 0xFFu;
506    let g1 = (d1 >>  8u) & 0xFFu;
507    let b1 = (d1 >> 16u) & 0xFFu;
508
509    var palette: array<u32, 4>;
510    palette[0] = d0;
511    palette[1] = d1;
512    palette[2] = ((2u*r0+r1)/3u) | (((2u*g0+g1)/3u)<<8u) | (((2u*b0+b1)/3u)<<16u);
513    palette[3] = ((r0+2u*r1)/3u) | (((g0+2u*g1)/3u)<<8u) | (((b0+2u*b1)/3u)<<16u);
514
515    // --- Assign each pixel to the nearest palette entry ---
516    var lookup = 0u;
517    for (var i = 0u; i < 16u; i++) {{
518        let px = pixels[i];
519        var best_idx = 0u;
520        var best_dist = rgb_dist_sq(px, palette[0]);
521        for (var j = 1u; j < 4u; j++) {{
522            let d = rgb_dist_sq(px, palette[j]);
523            if d < best_dist {{ best_dist = d; best_idx = j; }}
524        }}
525        lookup |= best_idx << (i * 2u);
526    }}
527
528    // --- Write BC1 block: [c0 u16 LE | c1 u16 LE] as u32, then lookup u32 ---
529    let word0 = q0 | (q1 << 16u);
530    output_blocks[block_idx * 2u    ] = word0;
531    output_blocks[block_idx * 2u + 1u] = lookup;
532}}
533"#,
534        num_blocks_x = num_blocks_x,
535    )
536}
537
538/// Emit the WGSL compute shader for BC4 compression.
539///
540/// Each invocation processes one 4×4 R8 block (16 bytes).  The input is a
541/// `u32` storage buffer where each `u32` packs four consecutive R8 pixels.
542fn make_bc4_shader_source(width: u32, _height: u32) -> String {
543    let num_blocks_x = width / 4;
544    let pixels_per_u32 = 4u32; // 4 R8 bytes → 1 u32
545    let u32s_per_row = width / pixels_per_u32; // how many u32s span one image row
546    format!(
547        r#"// BC4 (ATI1 / RGTC1) single-channel block-compression compute shader.
548// One thread per 4×4 block; input R8 pixels are packed 4-per-u32.
549
550struct Uniforms {{
551    width: u32,
552    height: u32,
553}};
554
555@group(0) @binding(0) var<uniform> uniforms: Uniforms;
556@group(0) @binding(1) var<storage, read> input_r8: array<u32>;
557@group(0) @binding(2) var<storage, read_write> output_blocks: array<u32>;
558
559// Extract byte `lane` (0..3) from a packed u32 (little-endian byte order).
560fn extract_byte(word: u32, lane: u32) -> u32 {{
561    return (word >> (lane * 8u)) & 0xFFu;
562}}
563
564@compute @workgroup_size(64)
565fn main(@builtin(global_invocation_id) gid: vec3<u32>) {{
566    let block_idx = gid.x;
567    let num_blocks_x = {num_blocks_x}u;
568    let num_blocks_y = uniforms.height / 4u;
569    let total_blocks = num_blocks_x * num_blocks_y;
570    if block_idx >= total_blocks {{
571        return;
572    }}
573
574    let bx = block_idx % num_blocks_x;
575    let by = block_idx / num_blocks_x;
576
577    // --- Gather 16 R8 pixels from packed u32 storage ---
578    // Each u32 holds 4 consecutive R8 bytes in little-endian order.
579    var pixels: array<u32, 16>;
580    let u32s_per_row = {u32s_per_row}u;
581    for (var iy = 0u; iy < 4u; iy++) {{
582        let row_base = (by * 4u + iy) * u32s_per_row;
583        let col_u32  = (bx * 4u) / 4u;      // which u32 word holds column bx*4
584        let col_lane = (bx * 4u) % 4u;      // which byte lane within that word
585
586        // We need 4 consecutive bytes starting at (bx*4, iy).
587        // Since bx*4 is always a multiple of 4 (width is a multiple of 4),
588        // all 4 bytes live in the same u32 word.
589        let word = input_r8[row_base + col_u32];
590        for (var ix = 0u; ix < 4u; ix++) {{
591            pixels[iy * 4u + ix] = extract_byte(word, col_lane + ix);
592        }}
593    }}
594
595    // --- Find ep0 (max) and ep1 (min) ---
596    var ep0 = pixels[0];
597    var ep1 = pixels[0];
598    for (var i = 1u; i < 16u; i++) {{
599        if pixels[i] > ep0 {{ ep0 = pixels[i]; }}
600        if pixels[i] < ep1 {{ ep1 = pixels[i]; }}
601    }}
602
603    // --- Build 8-value palette (ep0 > ep1 mode) ---
604    // v[0]=ep0, v[1]=ep1, v[2..7] interpolated
605    var palette: array<u32, 8>;
606    palette[0] = ep0;
607    palette[1] = ep1;
608    palette[2] = (6u*ep0 + 1u*ep1) / 7u;
609    palette[3] = (5u*ep0 + 2u*ep1) / 7u;
610    palette[4] = (4u*ep0 + 3u*ep1) / 7u;
611    palette[5] = (3u*ep0 + 4u*ep1) / 7u;
612    palette[6] = (2u*ep0 + 5u*ep1) / 7u;
613    palette[7] = (1u*ep0 + 6u*ep1) / 7u;
614
615    // --- Assign 3-bit indices ---
616    // Pack 16 × 3-bit indices into 48 bits → stored as two u32s.
617    // We fill a 64-bit logical value then split it: lo = bits 31:0, hi = bits 47:32.
618    var indices_lo = 0u;  // bits  0..31 (10.67 indices)
619    var indices_hi = 0u;  // bits 32..47 ( 5.33 indices) — only low 16 bits used
620
621    for (var i = 0u; i < 16u; i++) {{
622        let val = pixels[i];
623        // Find nearest palette entry.
624        var best_idx = 0u;
625        var best_dist = select(ep0 - val, val - ep0, val > ep0); // abs
626        for (var j = 1u; j < 8u; j++) {{
627            let pv = palette[j];
628            let d = select(pv - val, val - pv, val > pv);
629            if d < best_dist {{ best_dist = d; best_idx = j; }}
630        }}
631
632        // Bit position of this index in the 48-bit packed field.
633        let bit_pos = i * 3u;
634        if bit_pos < 32u {{
635            indices_lo |= best_idx << bit_pos;
636            // Handle straddle: if best_idx's 3 bits cross the 32-bit boundary.
637            if bit_pos > 29u {{
638                let overflow_bits = bit_pos + 3u - 32u;
639                indices_hi |= best_idx >> (3u - overflow_bits);
640            }}
641        }} else {{
642            indices_hi |= best_idx << (bit_pos - 32u);
643        }}
644    }}
645
646    // --- Write BC4 block ---
647    // Layout: [ep0 u8][ep1 u8][6 index bytes LE]
648    // We emit as two u32s:
649    //   word0 = ep0 | (ep1 << 8) | (index_byte0 << 16) | (index_byte1 << 24)
650    //   word1 = index_byte2..5
651    // indices_lo holds index bytes 0,1,2,3; indices_hi holds bytes 4,5 in low 16 bits.
652    let word0 = ep0 | (ep1 << 8u) | ((indices_lo & 0xFFFFu) << 16u);
653    let word1 = (indices_lo >> 16u) | (indices_hi << 16u);
654    // Note: word1 bits 16..31 hold index bytes 4&5 (indices_hi low 16 bits shifted up).
655    // Actual layout is: word0[16:31] = idx_bytes[0:1], word1[0:15] = idx_bytes[2:3],
656    //                   word1[16:31] = idx_bytes[4:5].
657    output_blocks[block_idx * 2u    ] = word0;
658    output_blocks[block_idx * 2u + 1u] = word1;
659}}
660"#,
661        num_blocks_x = num_blocks_x,
662        u32s_per_row = u32s_per_row,
663    )
664}
665
666// ─────────────────────────────────────────────────────────────────────────────
667// TextureCompressor struct
668// ─────────────────────────────────────────────────────────────────────────────
669
670/// GPU-accelerated texture compressor for BC1 and BC4 block formats.
671///
672/// Create once with [`TextureCompressor::new`], then call [`TextureCompressor::compress`]
673/// for each image.  The `width` and `height` supplied at construction time are
674/// fixed for the lifetime of this object.
675///
676/// When no wgpu backend is available (e.g., headless CI) `compress` falls back
677/// automatically to the CPU path provided by [`compress_bc1_block_cpu`] and
678/// [`compress_bc4_block_cpu`].
679pub struct TextureCompressor {
680    format: TextureFormat,
681    /// The compiled GPU compute pipeline (may be `None` when GPU unavailable).
682    pipeline: Option<Arc<wgpu::ComputePipeline>>,
683    /// Bind group layout for the GPU pipeline.
684    bind_group_layout: Option<wgpu::BindGroupLayout>,
685    width: u32,
686    height: u32,
687}
688
689impl TextureCompressor {
690    /// Create a new `TextureCompressor`.
691    ///
692    /// # Arguments
693    ///
694    /// * `ctx`    — Active GPU context.
695    /// * `format` — Target compressed format (`Bc1RgbUnorm` or `Bc4RUnorm`).
696    /// * `width`  — Pixel width of images to compress (must be a multiple of 4).
697    /// * `height` — Pixel height of images to compress (must be a multiple of 4).
698    ///
699    /// # Errors
700    ///
701    /// Returns [`GpuError::InvalidKernelParams`] when either dimension is not a
702    /// multiple of 4, or a pipeline error when WGSL compilation fails.
703    pub fn new(
704        ctx: &GpuContext,
705        format: TextureFormat,
706        width: u32,
707        height: u32,
708    ) -> GpuResult<Self> {
709        // Pure-Rust validation: no GPU touched.
710        validate_texture_dimensions(width, height)?;
711
712        // Build shader source and pipeline.
713        let shader_src = match format {
714            TextureFormat::Bc1RgbUnorm => make_bc1_shader_source(width, height),
715            TextureFormat::Bc4RUnorm => make_bc4_shader_source(width, height),
716        };
717
718        let mut shader = WgslShader::new(shader_src, "main");
719        let shader_module = match shader.compile(ctx.device()) {
720            Ok(m) => m,
721            Err(e) => {
722                // If shader compilation fails we still return a valid (CPU-only)
723                // compressor rather than propagating.  This allows callers on
724                // restricted environments to receive useful output.
725                tracing::warn!("BC texture shader compilation failed, falling back to CPU: {e}");
726                return Ok(Self {
727                    format,
728                    pipeline: None,
729                    bind_group_layout: None,
730                    width,
731                    height,
732                });
733            }
734        };
735
736        // Bind group layout: uniform | storage read | storage read-write
737        let bind_group_layout = create_compute_bind_group_layout(
738            ctx.device(),
739            &[
740                uniform_buffer_layout(0),        // uniforms
741                storage_buffer_layout(1, true),  // input pixels (read-only)
742                storage_buffer_layout(2, false), // output blocks (read-write)
743            ],
744            Some("TextureCompressor BGL"),
745        )?;
746
747        let pipeline = ComputePipelineBuilder::new(ctx.device(), shader_module, "main")
748            .bind_group_layout(&bind_group_layout)
749            .label("TextureCompressor Pipeline")
750            .build()?;
751
752        Ok(Self {
753            format,
754            pipeline: Some(Arc::new(pipeline)),
755            bind_group_layout: Some(bind_group_layout),
756            width,
757            height,
758        })
759    }
760
761    // ── Accessors ──────────────────────────────────────────────────────────
762
763    /// The target texture format.
764    pub fn format(&self) -> TextureFormat {
765        self.format
766    }
767
768    /// Image width (pixels) supplied at construction time.
769    pub fn width(&self) -> u32 {
770        self.width
771    }
772
773    /// Image height (pixels) supplied at construction time.
774    pub fn height(&self) -> u32 {
775        self.height
776    }
777
778    // ── Main compression entry-point ───────────────────────────────────────
779
780    /// Compress `input` using this compressor's format.
781    ///
782    /// - **BC1**: `input` must be `width × height × 4` bytes (RGBA8 row-major).
783    /// - **BC4**: `input` must be `width × height` bytes (R8 row-major).
784    ///
785    /// Returns `(width/4) × (height/4) × 8` bytes of block-compressed data.
786    ///
787    /// # GPU path
788    ///
789    /// When a wgpu backend is available the method dispatches a compute shader
790    /// and reads back the result synchronously.
791    ///
792    /// # CPU fallback
793    ///
794    /// When `self.pipeline` is `None` (shader compilation failed or no backend)
795    /// or when wgpu dispatch fails, the method falls back to the pure-Rust
796    /// block encoders `compress_bc1_block_cpu` / `compress_bc4_block_cpu`.
797    ///
798    /// # Errors
799    ///
800    /// Returns an error when the input length does not match the expected size
801    /// for the configured format and dimensions.
802    pub fn compress(&self, ctx: &GpuContext, input: &[u8]) -> GpuResult<Vec<u8>> {
803        // Validate input length.
804        let expected = match self.format {
805            TextureFormat::Bc1RgbUnorm => (self.width * self.height * 4) as usize,
806            TextureFormat::Bc4RUnorm => (self.width * self.height) as usize,
807        };
808        if input.len() != expected {
809            return Err(GpuError::invalid_kernel_params(format!(
810                "input length {} != expected {} for {:?} {}×{}",
811                input.len(),
812                expected,
813                self.format,
814                self.width,
815                self.height
816            )));
817        }
818
819        // Try GPU path first.
820        if let (Some(pipeline), Some(bgl)) = (&self.pipeline, &self.bind_group_layout) {
821            match self.compress_gpu(ctx, input, pipeline, bgl) {
822                Ok(output) => return Ok(output),
823                Err(e) => {
824                    // GPU failed; fall through to CPU.
825                    tracing::warn!("GPU texture compression failed, falling back to CPU: {e}");
826                }
827            }
828        }
829
830        // CPU fallback path.
831        self.compress_cpu(input)
832    }
833
834    // ── CPU fallback ───────────────────────────────────────────────────────
835
836    /// Pure-CPU compression path (always correct; slower for large images).
837    fn compress_cpu(&self, input: &[u8]) -> GpuResult<Vec<u8>> {
838        let blocks_x = (self.width / 4) as usize;
839        let blocks_y = (self.height / 4) as usize;
840        let total_blocks = blocks_x * blocks_y;
841        let mut output = vec![0u8; total_blocks * 8];
842
843        match self.format {
844            TextureFormat::Bc1RgbUnorm => {
845                // RGBA8 input: 4 bytes per pixel, 16 pixels per block → 64 bytes per block.
846                for by in 0..blocks_y {
847                    for bx in 0..blocks_x {
848                        let mut rgba_block = [0u8; 64];
849                        for iy in 0..4 {
850                            for ix in 0..4 {
851                                let px_x = bx * 4 + ix;
852                                let px_y = by * 4 + iy;
853                                let src = (px_y * self.width as usize + px_x) * 4;
854                                let dst = (iy * 4 + ix) * 4;
855                                rgba_block[dst..dst + 4].copy_from_slice(&input[src..src + 4]);
856                            }
857                        }
858                        let encoded = compress_bc1_block_cpu(&rgba_block);
859                        let block_idx = by * blocks_x + bx;
860                        output[block_idx * 8..block_idx * 8 + 8].copy_from_slice(&encoded);
861                    }
862                }
863            }
864            TextureFormat::Bc4RUnorm => {
865                // R8 input: 1 byte per pixel, 16 pixels per block → 16 bytes per block.
866                for by in 0..blocks_y {
867                    for bx in 0..blocks_x {
868                        let mut r_block = [0u8; 16];
869                        for iy in 0..4 {
870                            for ix in 0..4 {
871                                let px_x = bx * 4 + ix;
872                                let px_y = by * 4 + iy;
873                                let src = px_y * self.width as usize + px_x;
874                                r_block[iy * 4 + ix] = input[src];
875                            }
876                        }
877                        let encoded = compress_bc4_block_cpu(&r_block);
878                        let block_idx = by * blocks_x + bx;
879                        output[block_idx * 8..block_idx * 8 + 8].copy_from_slice(&encoded);
880                    }
881                }
882            }
883        }
884        Ok(output)
885    }
886
887    // ── GPU dispatch ───────────────────────────────────────────────────────
888
889    /// Attempt GPU-accelerated compression.
890    ///
891    /// For BC1 the input is uploaded as `u32`s (each `u32` = one RGBA8 pixel).
892    /// For BC4 the input is uploaded as `u32`s with 4 R8 bytes packed per `u32`.
893    fn compress_gpu(
894        &self,
895        ctx: &GpuContext,
896        input: &[u8],
897        pipeline: &wgpu::ComputePipeline,
898        bgl: &wgpu::BindGroupLayout,
899    ) -> GpuResult<Vec<u8>> {
900        let blocks_x = self.width / 4;
901        let blocks_y = self.height / 4;
902        let total_blocks = (blocks_x * blocks_y) as usize;
903        let output_bytes = total_blocks * 8; // 2 u32s per block
904
905        // --- Build uniform data ---
906        let uniforms_value = Bc1Uniforms {
907            width: self.width,
908            height: self.height,
909        };
910        let uniform_data = bytemuck::bytes_of(&uniforms_value);
911
912        // --- Upload uniform buffer ---
913        let uniform_buf_size = (std::mem::size_of::<Bc1Uniforms>() as u64 + 15) & !15;
914        let uniform_buf = ctx.device().create_buffer(&BufferDescriptor {
915            label: Some("TC uniform"),
916            size: uniform_buf_size,
917            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
918            mapped_at_creation: false,
919        });
920        ctx.queue().write_buffer(&uniform_buf, 0, uniform_data);
921
922        // --- Pack input into u32 buffer ---
923        let input_u32: Vec<u32> = match self.format {
924            TextureFormat::Bc1RgbUnorm => {
925                // Each RGBA8 pixel → 1 u32 (R in bits 7:0, G 15:8, B 23:16, A 31:24)
926                input
927                    .chunks_exact(4)
928                    .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
929                    .collect()
930            }
931            TextureFormat::Bc4RUnorm => {
932                // 4 R8 bytes → 1 u32
933                let padded_len = (input.len() + 3) & !3;
934                let mut padded = input.to_vec();
935                padded.resize(padded_len, 0);
936                padded
937                    .chunks_exact(4)
938                    .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
939                    .collect()
940            }
941        };
942
943        let input_buf_size = ((input_u32.len() * std::mem::size_of::<u32>()) as u64 + 255) & !255;
944        let input_buf = ctx.device().create_buffer(&BufferDescriptor {
945            label: Some("TC input"),
946            size: input_buf_size,
947            usage: BufferUsages::STORAGE | BufferUsages::COPY_DST,
948            mapped_at_creation: false,
949        });
950        ctx.queue()
951            .write_buffer(&input_buf, 0, bytemuck::cast_slice(&input_u32));
952
953        // --- Create output buffer (2 u32s per block) ---
954        let output_buf_size = ((total_blocks * 2 * std::mem::size_of::<u32>()) as u64 + 255) & !255;
955        let output_buf = ctx.device().create_buffer(&BufferDescriptor {
956            label: Some("TC output"),
957            size: output_buf_size,
958            usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC,
959            mapped_at_creation: false,
960        });
961
962        // --- Build bind group ---
963        let bind_group = ctx.device().create_bind_group(&BindGroupDescriptor {
964            label: Some("TC BindGroup"),
965            layout: bgl,
966            entries: &[
967                BindGroupEntry {
968                    binding: 0,
969                    resource: uniform_buf.as_entire_binding(),
970                },
971                BindGroupEntry {
972                    binding: 1,
973                    resource: input_buf.as_entire_binding(),
974                },
975                BindGroupEntry {
976                    binding: 2,
977                    resource: output_buf.as_entire_binding(),
978                },
979            ],
980        });
981
982        // --- Dispatch compute ---
983        let workgroup_size = 64u32;
984        let dispatch_count = (total_blocks as u32 + workgroup_size - 1) / workgroup_size;
985
986        let mut encoder = ctx
987            .device()
988            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
989                label: Some("TC encoder"),
990            });
991        {
992            let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
993                label: Some("TC compute pass"),
994                timestamp_writes: None,
995            });
996            cpass.set_pipeline(pipeline);
997            cpass.set_bind_group(0, &bind_group, &[]);
998            cpass.dispatch_workgroups(dispatch_count, 1, 1);
999        }
1000
1001        // --- Staging buffer for readback ---
1002        let staging = ctx.device().create_buffer(&BufferDescriptor {
1003            label: Some("TC staging"),
1004            size: output_buf_size,
1005            usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
1006            mapped_at_creation: false,
1007        });
1008        encoder.copy_buffer_to_buffer(&output_buf, 0, &staging, 0, output_buf_size);
1009        ctx.queue().submit(Some(encoder.finish()));
1010
1011        // Map before poll: poll drives both submit completion and map callback.
1012        let slice = staging.slice(..);
1013        let (tx, rx) = std::sync::mpsc::sync_channel(1);
1014        slice.map_async(wgpu::MapMode::Read, move |result| {
1015            let _ = tx.send(result);
1016        });
1017
1018        ctx.device()
1019            .poll(wgpu::PollType::wait_indefinitely())
1020            .map_err(|e| GpuError::execution_failed(format!("device poll failed: {e}")))?;
1021
1022        rx.recv()
1023            .map_err(|_| GpuError::buffer_mapping("TC staging channel closed"))?
1024            .map_err(|e| GpuError::buffer_mapping(e.to_string()))?;
1025
1026        let mapped = slice.get_mapped_range();
1027        let output_u32: &[u32] = bytemuck::cast_slice(&mapped[..total_blocks * 8]);
1028        let out_bytes: Vec<u8> = bytemuck::cast_slice(output_u32)[..output_bytes].to_vec();
1029        drop(mapped);
1030        staging.unmap();
1031
1032        Ok(out_bytes)
1033    }
1034}
1035
1036// ─────────────────────────────────────────────────────────────────────────────
1037// Unit tests (pure-Rust, no GPU required)
1038// ─────────────────────────────────────────────────────────────────────────────
1039
1040#[cfg(test)]
1041mod tests {
1042    use super::*;
1043
1044    #[test]
1045    fn test_quantize_dequantize_round_trip_basics() {
1046        let packed = quantize_rgb565(255, 255, 255);
1047        let (r, g, b) = dequantize_rgb565(packed);
1048        assert_eq!((r, g, b), (255, 255, 255));
1049
1050        let packed0 = quantize_rgb565(0, 0, 0);
1051        let (r0, g0, b0) = dequantize_rgb565(packed0);
1052        assert_eq!((r0, g0, b0), (0, 0, 0));
1053    }
1054
1055    #[test]
1056    fn test_nearest_index_4_black_is_index_0() {
1057        let endpoints = [
1058            (0u8, 0u8, 0u8),
1059            (255, 255, 255),
1060            (170, 170, 170),
1061            (85, 85, 85),
1062        ];
1063        assert_eq!(nearest_index_4((0, 0, 0), endpoints), 0);
1064    }
1065
1066    #[test]
1067    fn test_bc1_solid_block_equal_endpoints() {
1068        let mut block = [0u8; 64];
1069        for i in 0..16 {
1070            block[i * 4] = 100;
1071            block[i * 4 + 1] = 200;
1072            block[i * 4 + 2] = 50;
1073            block[i * 4 + 3] = 255;
1074        }
1075        let encoded = compress_bc1_block_cpu(&block);
1076        let c0 = u16::from_le_bytes([encoded[0], encoded[1]]);
1077        let c1 = u16::from_le_bytes([encoded[2], encoded[3]]);
1078        assert_eq!(c0, c1, "solid block must produce equal endpoints");
1079    }
1080
1081    #[test]
1082    fn test_bc4_solid_block_equal_endpoints() {
1083        let block = [128u8; 16];
1084        let encoded = compress_bc4_block_cpu(&block);
1085        assert_eq!(encoded[0], encoded[1], "solid BC4 block must have ep0==ep1");
1086    }
1087
1088    #[test]
1089    fn test_validate_texture_dimensions_multiples_ok() {
1090        assert!(validate_texture_dimensions(4, 4).is_ok());
1091        assert!(validate_texture_dimensions(256, 512).is_ok());
1092    }
1093
1094    #[test]
1095    fn test_validate_texture_dimensions_rejects_non_multiple() {
1096        assert!(validate_texture_dimensions(7, 8).is_err());
1097        assert!(validate_texture_dimensions(8, 5).is_err());
1098        assert!(validate_texture_dimensions(1, 1).is_err());
1099    }
1100}