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 Rgba8(&'a [u8]),
15 Rgba32Float(&'a [f32]),
17}
18
19impl<'a> SourceImageData<'a> {
20 pub fn rgba32float<T: NoUninit>(data: &'a [T]) -> Result<SourceImageData<'a>, PodCastError> {
22 bytemuck::try_cast_slice(data).map(Self::Rgba32Float)
23 }
24
25 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 pub data: SourceImageData<'a>,
35 pub size: Extent3d,
37}
38
39static BASISU_ENCODER_INITIALIZED: OnceCell<()> = OnceCell::new();
40
41pub 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
56pub fn basisu_encoder_enable_debug_printf(enable: bool) {
58 unsafe { enc_sys::bu_enable_debug_printf(enable as u32) };
59}
60
61pub 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 pub basis_tex_format: BasisTextureFormat,
96 pub quality_level: i32,
98 pub effort_level: i32,
100 pub flags_and_quality: u64,
102 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 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 pub const fn with_flags(mut self, flags: u64) -> Self {
156 self.flags_and_quality |= flags;
157 self
158 }
159}
160
161impl BasisuEncoder {
162 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 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 pub fn clear_image(&mut self) {
254 assert!(unsafe { enc_sys::bu_comp_params_clear(self.params) }.is_ok());
255 }
256
257 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 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 pub fn block_on<T>(future: impl Future<Output = T>) -> T {
418 use core::task::{Context, Poll};
419
420 let mut future = core::pin::pin!(future);
422
423 let cx = &mut Context::from_waker(core::task::Waker::noop());
425
426 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}