Skip to main content

basisu_c_sys/extra/
encoder.rs

1use crate::common;
2use crate::encoder as enc_sys;
3use crate::extra::BuHeap;
4use crate::utils::BasisTextureFormat;
5use alloc::vec::Vec;
6use async_lock::OnceCell;
7use bytemuck::NoUninit;
8use bytemuck::PodCastError;
9use wgpu_types::{Extent3d, TextureViewDimension};
10
11#[derive(Debug, Clone)]
12pub enum SourceImageData<'a> {
13    /// A 32bpp RGBA data slice (4 bytes per pixel)
14    Rgba8(&'a [u8]),
15    /// A 64bpp float RGBA data slice (16 bytes per pixel)
16    Rgba32Float(&'a [f32]),
17}
18
19impl<'a> SourceImageData<'a> {
20    /// Cast the slice to [`SourceImageData::Rgba32Float`]. Return an error if the casting failed.
21    pub fn rgba32float<T: NoUninit>(data: &'a [T]) -> Result<SourceImageData<'a>, PodCastError> {
22        bytemuck::try_cast_slice(data).map(Self::Rgba32Float)
23    }
24
25    /// Cast the slice to [`SourceImageData::Rgba8`]. Return an error if the casting failed.
26    pub fn rgba8<T: NoUninit>(data: &'a [T]) -> Result<SourceImageData<'a>, PodCastError> {
27        bytemuck::try_cast_slice(data).map(Self::Rgba32Float)
28    }
29}
30
31#[derive(Debug, Clone)]
32pub struct SourceImage<'a> {
33    /// The input data of image pixels.
34    pub data: SourceImageData<'a>,
35    /// The size of image.
36    pub size: Extent3d,
37}
38
39static BASISU_ENCODER_INITIALIZED: OnceCell<()> = OnceCell::new();
40
41/// Init global data of encoder ([`enc_sys::bu_init`]), and basisu wasm if on web.
42pub async fn basisu_encoder_init() {
43    BASISU_ENCODER_INITIALIZED
44        .get_or_init(async || {
45            #[cfg(all(
46                target_arch = "wasm32",
47                target_vendor = "unknown",
48                target_os = "unknown",
49            ))]
50            crate::instantiate_basisu_wasm().await;
51            unsafe { enc_sys::bu_init() };
52        })
53        .await;
54}
55
56/// A wrapper of [`enc_sys::bu_enable_debug_printf`].
57pub fn basisu_encoder_enable_debug_printf(enable: bool) {
58    unsafe { enc_sys::bu_enable_debug_printf(enable as u32) };
59}
60
61/// Encoder that used to compress [`SourceImage`] to basis universal ktx2 file.
62pub struct BasisuEncoder {
63    params: u64,
64}
65
66impl Default for BasisuEncoder {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72#[derive(Debug, thiserror::Error, PartialEq)]
73#[non_exhaustive]
74pub enum BasisuEncodeError {
75    #[error("`BasisuEncoder::set_image_slice` only accepts image with 1 layer")]
76    SetImageSliceOnlyAcceptsOneLayer,
77    #[error("Image data is empty")]
78    EmptyImageData,
79    #[error("Image {image_size:?} Expects data length {expected_len}, got {data_len}")]
80    ImageUnmatchedDataAndSize {
81        image_size: Extent3d,
82        expected_len: usize,
83        data_len: usize,
84    },
85    #[error("bu_comp_params_set_image_* failed")]
86    BuSetImageFailed,
87    #[error("bu_compress_texture failed")]
88    BuCompressFailed,
89}
90
91#[derive(Debug, Clone, Copy)]
92#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
93pub struct BasisuEncoderParams {
94    /// Target file format — one of the BTF_* constants (e.g. BTF_ETC1S, BTF_UASTC_LDR_4X4).
95    pub basis_tex_format: BasisTextureFormat,
96    /// Unified Quality level [1, 100]. See [`common::BU_QUALITY_MIN`], [`common::BU_QUALITY_MAX`]. Note the recommended usable unified quality range is [1, 100], but the C API accepts [0, 100]. Use -1 to use older non-unified/direct codec-specific quality level or lambda (low 8-bits of flags_and_quality, or via low_level_uastc_rdo_or_dct_quality).
97    pub quality_level: i32,
98    /// Unified Encoder effort [0, 10]. See [`common::BU_EFFORT_MIN`], [`common::BU_EFFORT_MAX`]. See `BU_EFFORT_*` presets. Use -1 to use older non-unified/direct codec-specific effort level (low 8-bits of flags_and_quality for some codecs).
99    pub effort_level: i32,
100    /// Bitwise OR of `BU_COMP_FLAGS_*` constants. Controls output format, mipmaps, color space, etc. Low 8-bits are either the older non-unified quality level, or for some codecs the non-unified effort level.
101    pub flags_and_quality: u64,
102    /// Low-level (non-unified) quality or lambda parameter for UASTC RDO encoding. Typically 0.0 for defaults. Must be 0.0 if using unified (not -1) quality level.
103    pub low_level_uastc_rdo_or_dct_quality: f32,
104}
105
106impl BasisuEncoderParams {
107    pub const fn new_with_srgb_defaults(basis_tex_format: BasisTextureFormat) -> Self {
108        Self {
109            basis_tex_format,
110            quality_level: 75,
111            effort_level: 2,
112            flags_and_quality: common::BU_COMP_FLAGS_THREADED
113                | common::BU_COMP_FLAGS_SRGB
114                | common::BU_COMP_FLAGS_KTX2_OUTPUT
115                | common::BU_COMP_FLAGS_KTX2_UASTC_ZSTD,
116            low_level_uastc_rdo_or_dct_quality: 0.0,
117        }
118    }
119
120    pub const fn new_with_linear_defaults(basis_tex_format: BasisTextureFormat) -> Self {
121        Self {
122            basis_tex_format,
123            quality_level: 75,
124            effort_level: 2,
125            flags_and_quality: common::BU_COMP_FLAGS_THREADED
126                | common::BU_COMP_FLAGS_KTX2_OUTPUT
127                | common::BU_COMP_FLAGS_KTX2_UASTC_ZSTD,
128            low_level_uastc_rdo_or_dct_quality: 0.0,
129        }
130    }
131
132    /// Return [`Self`] with `common::BU_COMP_FLAGS_TEXTURE_TYPE_*` set according to the view dimension.
133    ///
134    /// Panic if the view dimension is D1 or D3.
135    pub const fn with_tex_type(mut self, tex_type: TextureViewDimension) -> Self {
136        self.flags_and_quality = self.flags_and_quality
137            & !(common::BU_COMP_FLAGS_TEXTURE_TYPE_MASK
138                << common::BU_COMP_FLAGS_TEXTURE_TYPE_SHIFT);
139
140        self.flags_and_quality = self.flags_and_quality
141            | match tex_type {
142                TextureViewDimension::D2 => common::BU_COMP_FLAGS_TEXTURE_TYPE_2D,
143                TextureViewDimension::D2Array => common::BU_COMP_FLAGS_TEXTURE_TYPE_2D_ARRAY,
144                TextureViewDimension::Cube | TextureViewDimension::CubeArray => {
145                    common::BU_COMP_FLAGS_TEXTURE_TYPE_CUBEMAP_ARRAY
146                }
147                TextureViewDimension::D1 | TextureViewDimension::D3 => {
148                    panic!("Compressing 1D or 3D texture is unsupported")
149                }
150            };
151        self
152    }
153
154    /// Bitwise OR the flags (See `BU_COMP_FLAGS_*`) to `self`.
155    pub const fn with_flags(mut self, flags: u64) -> Self {
156        self.flags_and_quality |= flags;
157        self
158    }
159}
160
161impl BasisuEncoder {
162    /// Create a encoder. Panic if [`basisu_encoder_init`] hasn't been called.
163    pub fn new() -> Self {
164        if !BASISU_ENCODER_INITIALIZED.is_initialized() {
165            panic!("`basisu_encoder_init` must be called before create encoder");
166        }
167        Self {
168            params: unsafe { enc_sys::bu_new_comp_params() },
169        }
170    }
171
172    /// Set the input image of the encoder and clear other image set before.
173    ///
174    /// All the layers of the image will be set. To compress it as a cubemap or texture array,
175    /// you will need to add flag to [`BasisuEncoderParams::flags_and_quality`] by calling [`BasisuEncoderParams::with_tex_type`].
176    ///
177    /// If you already have continuous data for the cubemap or texture array, this should be faster than [`Self::set_image_slice`] .
178    pub fn set_image(&mut self, image: SourceImage) -> Result<(), BasisuEncodeError> {
179        self.clear_image();
180
181        match image.data {
182            SourceImageData::Rgba8(data) => unsafe {
183                let pixel_bytes = 4;
184                let expected_len = (image.size.width
185                    * image.size.height
186                    * image.size.depth_or_array_layers
187                    * pixel_bytes) as usize;
188                if data.len() != expected_len {
189                    return Err(BasisuEncodeError::ImageUnmatchedDataAndSize {
190                        image_size: image.size,
191                        expected_len,
192                        data_len: data.len(),
193                    });
194                }
195
196                let Some(bu_heap) = BuHeap::new(data) else {
197                    return Err(BasisuEncodeError::EmptyImageData);
198                };
199                let ptr = u64::from(bu_heap.ptr());
200                for i in 0..image.size.depth_or_array_layers {
201                    if enc_sys::bu_comp_params_set_image_rgba32(
202                        self.params,
203                        i,
204                        ptr + (i * image.size.width * image.size.height * pixel_bytes) as u64,
205                        image.size.width,
206                        image.size.height,
207                        image.size.width * pixel_bytes,
208                    )
209                    .is_err()
210                    {
211                        return Err(BasisuEncodeError::BuSetImageFailed);
212                    }
213                }
214            },
215            SourceImageData::Rgba32Float(data) => unsafe {
216                let pixel_bytes = 16;
217                let expected_len = (image.size.width
218                    * image.size.height
219                    * image.size.depth_or_array_layers
220                    * pixel_bytes) as usize;
221                if data.len() != expected_len {
222                    return Err(BasisuEncodeError::ImageUnmatchedDataAndSize {
223                        image_size: image.size,
224                        expected_len,
225                        data_len: data.len(),
226                    });
227                }
228
229                let Some(bu_heap) = BuHeap::new(data) else {
230                    return Err(BasisuEncodeError::EmptyImageData);
231                };
232                let ptr = u64::from(bu_heap.ptr());
233                for i in 0..image.size.depth_or_array_layers {
234                    if enc_sys::bu_comp_params_set_image_float_rgba(
235                        self.params,
236                        i,
237                        ptr + (i * image.size.width * image.size.height * pixel_bytes) as u64,
238                        image.size.width,
239                        image.size.height,
240                        image.size.width * pixel_bytes,
241                    )
242                    .is_err()
243                    {
244                        return Err(BasisuEncodeError::BuSetImageFailed);
245                    }
246                }
247            },
248        }
249        Ok(())
250    }
251
252    /// Clear the input image of encoder that was set.
253    pub fn clear_image(&mut self) {
254        assert!(unsafe { enc_sys::bu_comp_params_clear(self.params) }.is_ok());
255    }
256
257    /// Set a image slice at index. Other image set before is not cleared.
258    ///
259    /// After set all the layers of the image, to compress it as a cubemap or texture array,
260    /// you will need to add flag to [`BasisuEncoderParams::flags_and_quality`] by calling [`BasisuEncoderParams::with_tex_type`].
261    ///
262    /// The input image array layer count must be 1, otherwise an error will be returned.
263    ///
264    /// If you already have continuous data for the cubemap or texture array, [`Self::set_image`] should faster.
265    pub fn set_image_slice(
266        &mut self,
267        index: u32,
268        image: SourceImage,
269    ) -> Result<(), BasisuEncodeError> {
270        if image.size.depth_or_array_layers != 1 {
271            return Err(BasisuEncodeError::SetImageSliceOnlyAcceptsOneLayer);
272        }
273
274        match image.data {
275            SourceImageData::Rgba8(data) => unsafe {
276                let pixel_bytes = 4;
277                let expected_len = (image.size.width
278                    * image.size.height
279                    * image.size.depth_or_array_layers
280                    * pixel_bytes) as usize;
281                if data.len() != expected_len {
282                    return Err(BasisuEncodeError::ImageUnmatchedDataAndSize {
283                        image_size: image.size,
284                        expected_len,
285                        data_len: data.len(),
286                    });
287                }
288
289                let Some(bu_heap) = BuHeap::new(data) else {
290                    return Err(BasisuEncodeError::EmptyImageData);
291                };
292                let ptr = u64::from(bu_heap.ptr());
293                if enc_sys::bu_comp_params_set_image_rgba32(
294                    self.params,
295                    index,
296                    ptr,
297                    image.size.width,
298                    image.size.height,
299                    image.size.width * pixel_bytes,
300                )
301                .is_err()
302                {
303                    return Err(BasisuEncodeError::BuSetImageFailed);
304                }
305            },
306            SourceImageData::Rgba32Float(data) => unsafe {
307                let pixel_bytes = 16;
308                let expected_len = (image.size.width
309                    * image.size.height
310                    * image.size.depth_or_array_layers
311                    * pixel_bytes) as usize;
312                if data.len() != expected_len {
313                    return Err(BasisuEncodeError::ImageUnmatchedDataAndSize {
314                        image_size: image.size,
315                        expected_len,
316                        data_len: data.len(),
317                    });
318                }
319
320                let Some(bu_heap) = BuHeap::new(data) else {
321                    return Err(BasisuEncodeError::EmptyImageData);
322                };
323                let ptr = u64::from(bu_heap.ptr());
324                if enc_sys::bu_comp_params_set_image_float_rgba(
325                    self.params,
326                    index,
327                    ptr,
328                    image.size.width,
329                    image.size.height,
330                    image.size.width * pixel_bytes,
331                )
332                .is_err()
333                {
334                    return Err(BasisuEncodeError::BuSetImageFailed);
335                }
336            },
337        }
338        Ok(())
339    }
340
341    /// Compress the inputted image and return the bytes of ktx2 file result.
342    pub fn compress(&mut self, params: BasisuEncoderParams) -> Result<Vec<u8>, BasisuEncodeError> {
343        unsafe {
344            if enc_sys::bu_compress_texture(
345                self.params,
346                params.basis_tex_format as u32,
347                params.quality_level,
348                params.effort_level,
349                params.flags_and_quality,
350                params.low_level_uastc_rdo_or_dct_quality,
351            )
352            .is_err()
353            {
354                return Err(BasisuEncodeError::BuCompressFailed);
355            }
356            let out_size = enc_sys::bu_comp_params_get_comp_data_size(self.params);
357            let out_ptr = enc_sys::bu_comp_params_get_comp_data_ofs(self.params);
358            let result = crate::copy_basisu_memory_to_host(out_ptr, out_size);
359            Ok(result)
360        }
361    }
362}
363
364impl Drop for BasisuEncoder {
365    fn drop(&mut self) {
366        assert!(unsafe { enc_sys::bu_delete_comp_params(self.params).is_ok() });
367    }
368}
369
370#[cfg(test)]
371mod tests {
372    use wgpu_types::Extent3d;
373
374    use crate::extra::{
375        BasisuEncodeError, BasisuEncoder, SourceImage, SourceImageData, basisu_encoder_init,
376        encoder::BASISU_ENCODER_INITIALIZED,
377    };
378
379    #[test]
380    #[should_panic]
381    fn encoder_create_before_init() {
382        if BASISU_ENCODER_INITIALIZED.is_initialized() {
383            panic!("Basisu is already initialized, panic to skip this test");
384        } else {
385            BasisuEncoder::new();
386        }
387    }
388
389    #[test]
390    fn invalid_image_data() {
391        block_on(basisu_encoder_init());
392        let mut encoder = BasisuEncoder::new();
393        assert_eq!(
394            encoder.set_image(SourceImage {
395                data: SourceImageData::Rgba8(&[]),
396                size: Extent3d {
397                    width: 1,
398                    height: 1,
399                    depth_or_array_layers: 1
400                },
401            }),
402            Err(BasisuEncodeError::ImageUnmatchedDataAndSize {
403                image_size: Extent3d {
404                    width: 1,
405                    height: 1,
406                    depth_or_array_layers: 1
407                },
408                expected_len: 4,
409                data_len: 0
410            })
411        );
412    }
413
414    /// Blocks on the supplied `future`.
415    /// This implementation will busy-wait until it is completed.
416    /// Consider enabling the `async-io` or `futures-lite` features.
417    pub fn block_on<T>(future: impl Future<Output = T>) -> T {
418        use core::task::{Context, Poll};
419
420        // Pin the future on the stack.
421        let mut future = core::pin::pin!(future);
422
423        // We don't care about the waker as we're just going to poll as fast as possible.
424        let cx = &mut Context::from_waker(core::task::Waker::noop());
425
426        // Keep polling until the future is ready.
427        loop {
428            match future.as_mut().poll(cx) {
429                Poll::Ready(output) => return output,
430                Poll::Pending => core::hint::spin_loop(),
431            }
432        }
433    }
434}