1use std::any::TypeId;
2
3use lerc_band_materialize::{
4 copy_band_values_into_slice, BandLayout as MaterializeLayout, BandMaterializer,
5};
6use lerc_core::{BandLayout, BandSetInfo, BlobInfo, Error, PixelData, Result};
7use ndarray::{ArrayD, IxDyn};
8
9use crate::allocation::{checked_mul, default_vec, vec_with_capacity};
10
11#[derive(Debug, Clone, PartialEq)]
12pub struct Decoded {
13 pub info: BlobInfo,
14 pub pixels: PixelData,
15 pub mask: Option<Vec<u8>>,
16}
17
18impl Decoded {
19 pub fn into_ndarray<T: NdArrayElement>(self) -> Result<ArrayD<T>> {
20 let shape = self.info.ndarray_shape();
21 self.pixels.into_ndarray(&shape)
22 }
23
24 pub fn into_mask_ndarray(self) -> Result<Option<ArrayD<u8>>> {
25 let shape = self.info.mask_ndarray_shape();
26 self.mask
27 .map(|mask| {
28 ArrayD::from_shape_vec(IxDyn(&shape), mask).map_err(|err| {
29 Error::InvalidBlob(format!("failed to build ndarray from decoded mask: {err}"))
30 })
31 })
32 .transpose()
33 }
34}
35
36#[derive(Debug, Clone, PartialEq)]
37pub struct DecodedF64 {
38 pub info: BlobInfo,
39 pub pixels: Vec<f64>,
40 pub mask: Option<Vec<u8>>,
41}
42
43impl DecodedF64 {
44 pub fn into_ndarray(self) -> Result<ArrayD<f64>> {
45 ArrayD::from_shape_vec(IxDyn(&self.info.ndarray_shape()), self.pixels).map_err(|err| {
46 Error::InvalidBlob(format!(
47 "failed to build ndarray from decoded pixels: {err}"
48 ))
49 })
50 }
51
52 pub fn into_mask_ndarray(self) -> Result<Option<ArrayD<u8>>> {
53 let shape = self.info.mask_ndarray_shape();
54 self.mask
55 .map(|mask| {
56 ArrayD::from_shape_vec(IxDyn(&shape), mask).map_err(|err| {
57 Error::InvalidBlob(format!("failed to build ndarray from decoded mask: {err}"))
58 })
59 })
60 .transpose()
61 }
62}
63
64#[derive(Debug, Clone, PartialEq)]
65pub struct DecodedBandSet {
66 pub info: BandSetInfo,
67 pub bands: Vec<PixelData>,
68 pub band_masks: Vec<Option<Vec<u8>>>,
69}
70
71impl DecodedBandSet {
72 pub fn into_ndarray<T: BandElement>(self) -> Result<ArrayD<T>> {
73 self.into_ndarray_with_layout(BandLayout::Interleaved)
74 }
75
76 pub fn into_ndarray_with_layout<T: BandElement>(self, layout: BandLayout) -> Result<ArrayD<T>> {
77 let shape = self.info.ndarray_shape_for_layout(layout);
78 let values = self.into_vec_with_layout(layout)?;
79 ArrayD::from_shape_vec(IxDyn(&shape), values).map_err(|err| {
80 Error::InvalidBlob(format!(
81 "failed to build ndarray from decoded band set: {err}"
82 ))
83 })
84 }
85
86 pub fn into_vec_with_layout<T: BandElement>(self, layout: BandLayout) -> Result<Vec<T>> {
87 if self.bands.len() == 1 {
88 return T::from_pixel_data(self.bands.into_iter().next().unwrap());
89 }
90
91 let mut materializer = BandMaterializer::new(
92 self.info.bands[0].pixel_count()?,
93 self.info.depth() as usize,
94 self.info.band_count(),
95 materialize_layout(layout),
96 )
97 .map_err(materialize_error)?;
98 for (band_index, band) in self.bands.into_iter().enumerate() {
99 copy_pixel_data_into_materializer(&mut materializer, band_index, band)?;
100 }
101 materializer.finish().map_err(materialize_error)
102 }
103
104 pub fn copy_into_slice<T: BandElement>(self, layout: BandLayout, out: &mut [T]) -> Result<()> {
105 let pixel_count = self.info.bands[0].pixel_count()?;
106 let depth = self.info.depth() as usize;
107 let band_count = self.info.band_count();
108 let expected_len = self.info.value_count()?;
109 if out.len() != expected_len {
110 return Err(Error::InvalidBlob(format!(
111 "output slice length {} does not match decoded band set length {}",
112 out.len(),
113 expected_len
114 )));
115 }
116
117 for (band_index, band) in self.bands.into_iter().enumerate() {
118 copy_pixel_data_into_layout_slice(
119 out,
120 band_index,
121 pixel_count,
122 depth,
123 band_count,
124 layout,
125 band,
126 )?;
127 }
128 Ok(())
129 }
130
131 pub fn into_band_mask_ndarray(self) -> Result<Option<ArrayD<u8>>> {
132 into_band_mask_ndarray(self.info, self.band_masks)
133 }
134}
135
136pub fn into_band_mask_ndarray(
137 info: BandSetInfo,
138 band_masks: Vec<Option<Vec<u8>>>,
139) -> Result<Option<ArrayD<u8>>> {
140 if band_masks.iter().all(Option::is_none) {
141 return Ok(None);
142 }
143
144 let pixel_count = info.bands[0].pixel_count()?;
145 let band_count = info.band_count();
146 let shape = info.band_mask_ndarray_shape();
147
148 if band_count == 1 {
149 let mask = band_masks
150 .into_iter()
151 .next()
152 .flatten()
153 .map(Ok)
154 .unwrap_or_else(|| default_vec(pixel_count, "decoded band mask"))?;
155 return ArrayD::from_shape_vec(IxDyn(&shape), mask)
156 .map(Some)
157 .map_err(|err| {
158 Error::InvalidBlob(format!("failed to build ndarray from decoded mask: {err}"))
159 });
160 }
161
162 let merged_len = checked_mul(pixel_count, band_count, "decoded band mask length")?;
163 let mut merged = vec_with_capacity(merged_len, "decoded band mask")?;
164 for pixel in 0..pixel_count {
165 for band_mask in &band_masks {
166 merged.push(band_mask.as_ref().map(|mask| mask[pixel]).unwrap_or(1));
167 }
168 }
169
170 ArrayD::from_shape_vec(IxDyn(&shape), merged)
171 .map(Some)
172 .map_err(|err| {
173 Error::InvalidBlob(format!(
174 "failed to build ndarray from decoded band mask: {err}"
175 ))
176 })
177}
178
179trait SupportedElementValue: Copy + 'static + IntoF64 {
180 const KIND: BandElementKind;
181}
182
183macro_rules! match_pixel_data_values {
184 ($band:expr, |$values:ident| $body:expr) => {
185 match $band {
186 PixelData::I8($values) => $body,
187 PixelData::U8($values) => $body,
188 PixelData::I16($values) => $body,
189 PixelData::U16($values) => $body,
190 PixelData::I32($values) => $body,
191 PixelData::U32($values) => $body,
192 PixelData::F32($values) => $body,
193 PixelData::F64($values) => $body,
194 }
195 };
196}
197
198fn copy_pixel_data_into_materializer<T: BandElement>(
199 materializer: &mut BandMaterializer<T>,
200 band_index: usize,
201 band: PixelData,
202) -> Result<()> {
203 match_pixel_data_values!(band, |values| {
204 copy_typed_values_into_materializer(materializer, band_index, &values)
205 })
206}
207
208fn copy_typed_values_into_materializer<T: BandElement, U: SupportedElementValue>(
209 materializer: &mut BandMaterializer<T>,
210 band_index: usize,
211 values: &[U],
212) -> Result<()> {
213 if T::KIND == U::KIND {
214 let typed = unsafe { cast_slice::<U, T>(values) };
215 return materializer
216 .copy_band(band_index, typed)
217 .map_err(materialize_error);
218 }
219 if T::KIND == BandElementKind::F64 {
220 return materializer
221 .copy_band_with(band_index, |index| {
222 unsafe_cast::<T, f64>(values[index].into_f64())
223 })
224 .map_err(materialize_error);
225 }
226 Err(Error::InvalidBlob(format!(
227 "cannot decode {} pixels into ndarray<{}>",
228 data_type_name::<U>(),
229 std::any::type_name::<T>()
230 .rsplit("::")
231 .next()
232 .unwrap_or("unknown"),
233 )))
234}
235
236fn copy_pixel_data_into_layout_slice<T: BandElement>(
237 out: &mut [T],
238 band_index: usize,
239 pixel_count: usize,
240 depth: usize,
241 band_count: usize,
242 layout: BandLayout,
243 band: PixelData,
244) -> Result<()> {
245 match_pixel_data_values!(band, |values| {
246 copy_typed_values_into_layout_slice(
247 out,
248 band_index,
249 pixel_count,
250 depth,
251 band_count,
252 layout,
253 &values,
254 )
255 })
256}
257
258fn copy_typed_values_into_layout_slice<T: BandElement, U: SupportedElementValue>(
259 out: &mut [T],
260 band_index: usize,
261 pixel_count: usize,
262 depth: usize,
263 band_count: usize,
264 layout: BandLayout,
265 values: &[U],
266) -> Result<()> {
267 if T::KIND == U::KIND {
268 let typed = unsafe { cast_slice::<U, T>(values) };
269 return copy_band_values_into_slice(
270 out,
271 typed,
272 pixel_count,
273 depth,
274 band_index,
275 band_count,
276 materialize_layout(layout),
277 )
278 .map_err(materialize_error);
279 }
280 if T::KIND == BandElementKind::F64 {
281 let band_len = pixel_count
282 .checked_mul(depth.max(1))
283 .ok_or_else(|| Error::InvalidBlob("decoded band length overflows usize".into()))?;
284 if values.len() != band_len {
285 return Err(Error::InvalidBlob(
286 "decoded band length does not match its metadata".into(),
287 ));
288 }
289 for (value_index, value) in values.iter().copied().enumerate() {
290 let out_index = match layout {
291 BandLayout::Interleaved => {
292 if depth <= 1 {
293 value_index * band_count + band_index
294 } else {
295 let pixel = value_index / depth;
296 let sample = value_index % depth;
297 (pixel * band_count + band_index) * depth + sample
298 }
299 }
300 BandLayout::Bsq => band_index * band_len + value_index,
301 };
302 out[out_index] = unsafe_cast::<T, f64>(value.into_f64());
303 }
304 return Ok(());
305 }
306 Err(Error::InvalidBlob(format!(
307 "cannot decode {} pixels into ndarray<{}>",
308 data_type_name::<U>(),
309 std::any::type_name::<T>()
310 .rsplit("::")
311 .next()
312 .unwrap_or("unknown"),
313 )))
314}
315
316unsafe fn cast_slice<U, T>(values: &[U]) -> &[T] {
317 &*(values as *const [U] as *const [T])
318}
319
320fn unsafe_cast<T, U: Copy>(value: U) -> T {
321 unsafe { std::mem::transmute_copy(&value) }
322}
323
324trait IntoF64 {
325 fn into_f64(self) -> f64;
326}
327
328impl IntoF64 for i8 {
329 fn into_f64(self) -> f64 {
330 self as f64
331 }
332}
333impl IntoF64 for u8 {
334 fn into_f64(self) -> f64 {
335 self as f64
336 }
337}
338impl IntoF64 for i16 {
339 fn into_f64(self) -> f64 {
340 self as f64
341 }
342}
343impl IntoF64 for u16 {
344 fn into_f64(self) -> f64 {
345 self as f64
346 }
347}
348impl IntoF64 for i32 {
349 fn into_f64(self) -> f64 {
350 self as f64
351 }
352}
353impl IntoF64 for u32 {
354 fn into_f64(self) -> f64 {
355 self as f64
356 }
357}
358impl IntoF64 for f32 {
359 fn into_f64(self) -> f64 {
360 self as f64
361 }
362}
363impl IntoF64 for f64 {
364 fn into_f64(self) -> f64 {
365 self
366 }
367}
368
369fn data_type_name<T: 'static>() -> &'static str {
370 if TypeId::of::<T>() == TypeId::of::<i8>() {
371 "i8"
372 } else if TypeId::of::<T>() == TypeId::of::<u8>() {
373 "u8"
374 } else if TypeId::of::<T>() == TypeId::of::<i16>() {
375 "i16"
376 } else if TypeId::of::<T>() == TypeId::of::<u16>() {
377 "u16"
378 } else if TypeId::of::<T>() == TypeId::of::<i32>() {
379 "i32"
380 } else if TypeId::of::<T>() == TypeId::of::<u32>() {
381 "u32"
382 } else if TypeId::of::<T>() == TypeId::of::<f32>() {
383 "f32"
384 } else if TypeId::of::<T>() == TypeId::of::<f64>() {
385 "f64"
386 } else {
387 "unknown"
388 }
389}
390
391pub trait NdArrayElement: Sized + Clone {
392 fn from_pixel_data(pixels: PixelData) -> Result<Vec<Self>>;
393}
394
395mod private {
396 pub trait Sealed {}
397
398 impl Sealed for i8 {}
399 impl Sealed for u8 {}
400 impl Sealed for i16 {}
401 impl Sealed for u16 {}
402 impl Sealed for i32 {}
403 impl Sealed for u32 {}
404 impl Sealed for f32 {}
405 impl Sealed for f64 {}
406}
407
408#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
409pub enum BandElementKind {
410 I8,
411 U8,
412 I16,
413 U16,
414 I32,
415 U32,
416 F32,
417 F64,
418}
419
420pub trait BandElement: NdArrayElement + private::Sealed + Copy + Default + 'static {
421 const KIND: BandElementKind;
422}
423
424macro_rules! impl_exact_ndarray_element {
425 ($ty:ty, $variant:ident, $name:literal) => {
426 impl NdArrayElement for $ty {
427 fn from_pixel_data(pixels: PixelData) -> Result<Vec<Self>> {
428 match pixels {
429 PixelData::$variant(values) => Ok(values),
430 other => Err(Error::InvalidBlob(format!(
431 "cannot decode {} pixels into ndarray<{}>",
432 other.data_type().name(),
433 $name
434 ))),
435 }
436 }
437 }
438 };
439}
440
441impl_exact_ndarray_element!(i8, I8, "i8");
442impl_exact_ndarray_element!(u8, U8, "u8");
443impl_exact_ndarray_element!(i16, I16, "i16");
444impl_exact_ndarray_element!(u16, U16, "u16");
445impl_exact_ndarray_element!(i32, I32, "i32");
446impl_exact_ndarray_element!(u32, U32, "u32");
447impl_exact_ndarray_element!(f32, F32, "f32");
448
449impl NdArrayElement for f64 {
450 fn from_pixel_data(pixels: PixelData) -> Result<Vec<Self>> {
451 Ok(pixels.to_f64())
452 }
453}
454
455macro_rules! impl_band_element {
456 ($ty:ty, $kind:ident) => {
457 impl BandElement for $ty {
458 const KIND: BandElementKind = BandElementKind::$kind;
459 }
460
461 impl SupportedElementValue for $ty {
462 const KIND: BandElementKind = BandElementKind::$kind;
463 }
464 };
465}
466
467impl_band_element!(i8, I8);
468impl_band_element!(u8, U8);
469impl_band_element!(i16, I16);
470impl_band_element!(u16, U16);
471impl_band_element!(i32, I32);
472impl_band_element!(u32, U32);
473impl_band_element!(f32, F32);
474impl_band_element!(f64, F64);
475
476fn materialize_layout(layout: BandLayout) -> MaterializeLayout {
477 match layout {
478 BandLayout::Interleaved => MaterializeLayout::Interleaved,
479 BandLayout::Bsq => MaterializeLayout::Bsq,
480 }
481}
482
483fn materialize_error(err: lerc_band_materialize::MaterializeError) -> Error {
484 Error::InvalidBlob(err.to_string())
485}
486
487trait PixelDataExt {
488 fn into_ndarray<T: NdArrayElement>(self, shape: &[usize]) -> Result<ArrayD<T>>;
489}
490
491impl PixelDataExt for PixelData {
492 fn into_ndarray<T: NdArrayElement>(self, shape: &[usize]) -> Result<ArrayD<T>> {
493 ArrayD::from_shape_vec(IxDyn(shape), T::from_pixel_data(self)?).map_err(|err| {
494 Error::InvalidBlob(format!(
495 "failed to build ndarray from decoded pixels: {err}"
496 ))
497 })
498 }
499}