Skip to content

Commit 56a25ca

Browse files
committed
Simplify [u8]->sample cast
1 parent 7eee238 commit 56a25ca

File tree

1 file changed

+30
-66
lines changed

1 file changed

+30
-66
lines changed

src/codecs/tiff.rs

Lines changed: 30 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -433,72 +433,32 @@ fn cmyk_to_rgb(cmyk: &[u8]) -> [u8; 3] {
433433
]
434434
}
435435

436-
enum DtypeContainer<'a, T> {
437-
Slice(&'a [T]),
438-
Vec(Vec<T>),
439-
}
440-
441-
impl<T> DtypeContainer<'_, T> {
442-
fn as_slice(&self) -> &[T] {
443-
match self {
444-
DtypeContainer::Slice(slice) => slice,
445-
DtypeContainer::Vec(vec) => vec,
446-
}
447-
}
448-
}
449-
450-
fn u8_slice_as_f32(buf: &[u8]) -> ImageResult<DtypeContainer<f32>> {
451-
match bytemuck::try_cast_slice(buf) {
452-
Ok(slice) => Ok(DtypeContainer::<f32>::Slice(slice)),
453-
Err(err) => {
454-
match err {
455-
bytemuck::PodCastError::TargetAlignmentGreaterAndInputNotAligned => {
456-
// If the buffer is not aligned for a f32 slice, copy the buffer into a new Vec<f32>
457-
let mut vec = vec![0.0; buf.len() / 4];
458-
for (i, chunk) in buf.chunks_exact(4).enumerate() {
459-
let f32_val = f32::from_ne_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
460-
vec[i] = f32_val;
461-
}
462-
Ok(DtypeContainer::Vec(vec))
463-
}
464-
_ => {
465-
// If the buffer is not the correct length for a f32 slice, err.
466-
Err(ImageError::Parameter(ParameterError::from_kind(
467-
ParameterErrorKind::Generic(format!("{:?}", err)),
468-
)))
469-
}
470-
}
471-
}
472-
}
473-
}
474-
475-
fn u8_slice_as_u16(buf: &[u8]) -> ImageResult<DtypeContainer<u16>> {
476-
match bytemuck::try_cast_slice(buf) {
477-
Ok(slice) => Ok(DtypeContainer::<u16>::Slice(slice)),
478-
Err(err) => {
436+
/// Convert a slice of sample bytes to its semantic type, being a `Pod`.
437+
fn u8_slice_as_pod<P: bytemuck::Pod>(buf: &[u8]) -> ImageResult<std::borrow::Cow<'_, [P]>> {
438+
bytemuck::try_cast_slice(buf)
439+
.map(std::borrow::Cow::Borrowed)
440+
.or_else(|err| {
479441
match err {
480442
bytemuck::PodCastError::TargetAlignmentGreaterAndInputNotAligned => {
481-
// If the buffer is not aligned for a f32 slice, copy the buffer into a new Vec<f32>
482-
let mut vec = vec![0; buf.len() / 2];
483-
for (i, chunk) in buf.chunks_exact(2).enumerate() {
484-
let u16_val = u16::from_ne_bytes([chunk[0], chunk[1]]);
485-
vec[i] = u16_val;
486-
}
487-
Ok(DtypeContainer::Vec(vec))
443+
// If the buffer is not aligned for a native slice, copy the buffer into a Vec,
444+
// aligning it in the process. This is only done if the element count can be
445+
// represented exactly.
446+
let vec = bytemuck::allocation::pod_collect_to_vec(buf);
447+
Ok(std::borrow::Cow::Owned(vec))
488448
}
449+
/* only expecting: bytemuck::PodCastError::OutputSliceWouldHaveSlop */
489450
_ => {
490-
// If the buffer is not aligned or the correct length for a u16 slice, err.
491-
//
492-
// `bytemuck::PodCastError` of bytemuck-1.2.0 does not implement
493-
// `Error` and `Display` trait.
451+
// `bytemuck::PodCastError` of bytemuck-1.2.0 does not implement `Error` and
452+
// `Display` trait.
494453
// See <https://github.com/Lokathor/bytemuck/issues/22>.
495454
Err(ImageError::Parameter(ParameterError::from_kind(
496-
ParameterErrorKind::Generic(format!("{err:?}")),
455+
ParameterErrorKind::Generic(format!(
456+
"Casting samples to their representation failed: {err:?}",
457+
)),
497458
)))
498459
}
499460
}
500-
}
501-
}
461+
})
502462
}
503463

504464
impl<W: Write + Seek> TiffEncoder<W> {
@@ -562,20 +522,24 @@ impl<W: Write + Seek> TiffEncoder<W> {
562522
ExtendedColorType::Rgb8 => encoder.write_image::<RGB8>(width, height, buf),
563523
ExtendedColorType::Rgba8 => encoder.write_image::<RGBA8>(width, height, buf),
564524
ExtendedColorType::L16 => {
565-
encoder.write_image::<Gray16>(width, height, u8_slice_as_u16(buf)?.as_slice())
525+
encoder.write_image::<Gray16>(width, height, u8_slice_as_pod::<u16>(buf)?.as_ref())
566526
}
567527
ExtendedColorType::Rgb16 => {
568-
encoder.write_image::<RGB16>(width, height, u8_slice_as_u16(buf)?.as_slice())
528+
encoder.write_image::<RGB16>(width, height, u8_slice_as_pod::<u16>(buf)?.as_ref())
569529
}
570530
ExtendedColorType::Rgba16 => {
571-
encoder.write_image::<RGBA16>(width, height, u8_slice_as_u16(buf)?.as_slice())
572-
}
573-
ExtendedColorType::Rgb32F => {
574-
encoder.write_image::<RGB32Float>(width, height, u8_slice_as_f32(buf)?.as_slice())
575-
}
576-
ExtendedColorType::Rgba32F => {
577-
encoder.write_image::<RGBA32Float>(width, height, u8_slice_as_f32(buf)?.as_slice())
531+
encoder.write_image::<RGBA16>(width, height, u8_slice_as_pod::<u16>(buf)?.as_ref())
578532
}
533+
ExtendedColorType::Rgb32F => encoder.write_image::<RGB32Float>(
534+
width,
535+
height,
536+
u8_slice_as_pod::<f32>(buf)?.as_ref(),
537+
),
538+
ExtendedColorType::Rgba32F => encoder.write_image::<RGBA32Float>(
539+
width,
540+
height,
541+
u8_slice_as_pod::<f32>(buf)?.as_ref(),
542+
),
579543
_ => {
580544
return Err(ImageError::Unsupported(
581545
UnsupportedError::from_format_and_kind(

0 commit comments

Comments
 (0)