terrain_codec/heightmap/
container.rs

1//! PNG / WebP / AVIF container helpers for heightmap RGB bytes.
2//!
3//! Each container is gated on its own cargo feature so callers only pay
4//! the compile-time cost of what they actually use:
5//!
6//! | Feature | Provides                           | Backend                            |
7//! |---------|------------------------------------|------------------------------------|
8//! | `png`   | [`rgb_to_png`] / [`rgba_to_png`]   | `image/png`                        |
9//! | `webp`  | [`rgb_to_webp`] / [`rgba_to_webp`] | `image/webp` (lossless)            |
10//! | `avif`  | [`rgb_to_avif`] / [`rgba_to_avif`] | `image/avif` (ravif, encode-only)  |
11//!
12//! Each format has an `rgb_*` (3-channel, the heightmap path) and an
13//! `rgba_*` (4-channel, e.g. a packed watermask) variant.
14//!
15//! [`decode_image`] auto-detects whichever formats are compiled in. WebP
16//! encoding is **lossless** — lossy WebP would need `libwebp` which is
17//! out of scope here.
18//!
19//! For runtime-chosen container format use the [`ContainerFormat`] enum
20//! and the dispatching [`rgb_to_container`]; calling with a format whose
21//! feature wasn't enabled returns [`ContainerError::Unsupported`] rather
22//! than failing to compile.
23
24use std::fmt;
25use std::io::Cursor;
26use std::str::FromStr;
27
28#[cfg(feature = "avif")]
29use image::codecs::avif::AvifEncoder;
30#[cfg(feature = "png")]
31use image::codecs::png::PngEncoder;
32#[cfg(feature = "webp")]
33use image::codecs::webp::WebPEncoder;
34use image::{ExtendedColorType, ImageEncoder, ImageReader};
35
36/// Re-export of [`image::ImageError`] for callers that don't want to
37/// pull in the `image` crate directly.
38pub type ImageError = image::ImageError;
39
40/// A decoded image returned by [`decode_image`].
41#[derive(Debug, Clone)]
42pub struct DecodedImage {
43    /// Flat row-major RGB bytes (3 bytes per pixel).
44    pub rgb: Vec<u8>,
45    /// Image width in pixels.
46    pub width: u32,
47    /// Image height in pixels.
48    pub height: u32,
49}
50
51/// Identifies one of the supported image container formats for the
52/// runtime-dispatched [`rgb_to_container`] entry point.
53///
54/// All three variants are always present in the enum so callers can
55/// parse user-supplied format names regardless of which cargo features
56/// were enabled at compile time. Encoding into a format whose feature is
57/// not enabled returns [`ContainerError::Unsupported`].
58#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
59pub enum ContainerFormat {
60    /// PNG.
61    Png,
62    /// Lossless WebP.
63    Webp,
64    /// AVIF (encode-only).
65    Avif,
66}
67
68impl ContainerFormat {
69    /// All variants, in declaration order.
70    pub const ALL: [Self; 3] = [Self::Png, Self::Webp, Self::Avif];
71
72    /// Canonical lowercase name (`"png"` / `"webp"` / `"avif"`).
73    pub const fn name(self) -> &'static str {
74        match self {
75            Self::Png => "png",
76            Self::Webp => "webp",
77            Self::Avif => "avif",
78        }
79    }
80
81    /// IANA MIME type for the format.
82    pub const fn mime_type(self) -> &'static str {
83        match self {
84            Self::Png => "image/png",
85            Self::Webp => "image/webp",
86            Self::Avif => "image/avif",
87        }
88    }
89
90    /// Whether the encoder for this format was compiled in at build time.
91    pub const fn is_enabled(self) -> bool {
92        match self {
93            Self::Png => cfg!(feature = "png"),
94            Self::Webp => cfg!(feature = "webp"),
95            Self::Avif => cfg!(feature = "avif"),
96        }
97    }
98}
99
100impl fmt::Display for ContainerFormat {
101    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
102        f.write_str(self.name())
103    }
104}
105
106/// Error returned by [`ContainerFormat::from_str`] for an unrecognised name.
107#[derive(Debug, Clone, PartialEq, Eq)]
108pub struct ParseContainerFormatError {
109    /// The input string that failed to parse.
110    pub input: String,
111}
112
113impl fmt::Display for ParseContainerFormatError {
114    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
115        write!(
116            f,
117            "unknown container format `{}` (expected one of: png, webp, avif)",
118            self.input
119        )
120    }
121}
122
123impl std::error::Error for ParseContainerFormatError {}
124
125impl FromStr for ContainerFormat {
126    type Err = ParseContainerFormatError;
127
128    /// Parses case-insensitively. Accepts the canonical lowercase names
129    /// as well as the `image/<name>` MIME shorthand.
130    fn from_str(s: &str) -> Result<Self, Self::Err> {
131        match s.to_ascii_lowercase().as_str() {
132            "png" | "image/png" => Ok(Self::Png),
133            "webp" | "image/webp" => Ok(Self::Webp),
134            "avif" | "image/avif" => Ok(Self::Avif),
135            _ => Err(ParseContainerFormatError {
136                input: s.to_string(),
137            }),
138        }
139    }
140}
141
142/// Error returned by [`rgb_to_container`].
143#[derive(Debug)]
144pub enum ContainerError {
145    /// The underlying `image` encoder failed.
146    Image(ImageError),
147    /// The requested format's cargo feature was not enabled at build time.
148    Unsupported(ContainerFormat),
149}
150
151impl fmt::Display for ContainerError {
152    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153        match self {
154            Self::Image(e) => write!(f, "container encoding failed: {e}"),
155            Self::Unsupported(fmt_) => write!(
156                f,
157                "container format `{fmt_}` is not supported in this build — enable the `{fmt_}` cargo feature"
158            ),
159        }
160    }
161}
162
163impl std::error::Error for ContainerError {
164    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
165        match self {
166            Self::Image(e) => Some(e),
167            Self::Unsupported(_) => None,
168        }
169    }
170}
171
172impl From<ImageError> for ContainerError {
173    fn from(value: ImageError) -> Self {
174        Self::Image(value)
175    }
176}
177
178/// Wrap raw `width × height × 3` RGB bytes in the chosen container format.
179///
180/// This is the runtime-dispatched counterpart of the per-format
181/// [`rgb_to_png`] / [`rgb_to_webp`] / [`rgb_to_avif`] functions. Useful
182/// when the format is determined at runtime (CLI flag, query param,
183/// `Accept` header).
184///
185/// Returns [`ContainerError::Unsupported`] when the requested format's
186/// cargo feature was not enabled.
187///
188/// # Panics
189///
190/// Panics if `rgb.len() != (width * height * 3) as usize`.
191pub fn rgb_to_container(
192    format: ContainerFormat,
193    rgb: &[u8],
194    width: u32,
195    height: u32,
196) -> Result<Vec<u8>, ContainerError> {
197    let mut out = Vec::new();
198    rgb_to_container_to_writer(format, rgb, width, height, &mut out)?;
199    Ok(out)
200}
201
202/// Stream-encode raw `width × height × 3` RGB bytes into the chosen
203/// container format, writing directly to `writer` without an intermediate
204/// `Vec<u8>`. Pair with [`crate::heightmap::encode_to`] for an
205/// allocation-free DEM → image pipeline.
206pub fn rgb_to_container_to_writer<W: std::io::Write>(
207    format: ContainerFormat,
208    rgb: &[u8],
209    width: u32,
210    height: u32,
211    writer: W,
212) -> Result<(), ContainerError> {
213    match format {
214        ContainerFormat::Png => {
215            #[cfg(feature = "png")]
216            {
217                rgb_to_png_to_writer(rgb, width, height, writer)?;
218                Ok(())
219            }
220            #[cfg(not(feature = "png"))]
221            {
222                let _ = (rgb, width, height, writer);
223                Err(ContainerError::Unsupported(ContainerFormat::Png))
224            }
225        }
226        ContainerFormat::Webp => {
227            #[cfg(feature = "webp")]
228            {
229                rgb_to_webp_to_writer(rgb, width, height, writer)?;
230                Ok(())
231            }
232            #[cfg(not(feature = "webp"))]
233            {
234                let _ = (rgb, width, height, writer);
235                Err(ContainerError::Unsupported(ContainerFormat::Webp))
236            }
237        }
238        ContainerFormat::Avif => {
239            #[cfg(feature = "avif")]
240            {
241                rgb_to_avif_to_writer(rgb, width, height, writer)?;
242                Ok(())
243            }
244            #[cfg(not(feature = "avif"))]
245            {
246                let _ = (rgb, width, height, writer);
247                Err(ContainerError::Unsupported(ContainerFormat::Avif))
248            }
249        }
250    }
251}
252
253/// Encode raw `width × height × 3` RGB bytes as PNG directly to a writer.
254///
255/// Available behind the `png` cargo feature. Combine with
256/// [`crate::heightmap::encode_to`] to wrap a DEM tile in PNG with no
257/// intermediate `Vec<u8>`.
258///
259/// # Panics
260///
261/// Panics if `rgb.len() != (width * height * 3) as usize`.
262#[cfg(feature = "png")]
263pub fn rgb_to_png_to_writer<W: std::io::Write>(
264    rgb: &[u8],
265    width: u32,
266    height: u32,
267    writer: W,
268) -> Result<(), ImageError> {
269    assert_rgb_len(rgb, width, height);
270    PngEncoder::new(writer).write_image(rgb, width, height, ExtendedColorType::Rgb8)
271}
272
273/// Wrap raw `width × height × 3` RGB bytes in a PNG container `Vec<u8>`.
274///
275/// Available behind the `png` cargo feature.
276///
277/// # Errors
278///
279/// Returns [`ImageError`] if the underlying encoder fails (very rare for
280/// valid RGB inputs — typically only OOM).
281#[cfg(feature = "png")]
282pub fn rgb_to_png(rgb: &[u8], width: u32, height: u32) -> Result<Vec<u8>, ImageError> {
283    let mut out = Vec::with_capacity(rgb.len());
284    rgb_to_png_to_writer(rgb, width, height, &mut out)?;
285    Ok(out)
286}
287
288/// Encode raw `width × height × 3` RGB bytes as lossless WebP directly to a writer.
289///
290/// Available behind the `webp` cargo feature.
291///
292/// # Panics
293///
294/// Panics if `rgb.len() != (width * height * 3) as usize`.
295#[cfg(feature = "webp")]
296pub fn rgb_to_webp_to_writer<W: std::io::Write>(
297    rgb: &[u8],
298    width: u32,
299    height: u32,
300    writer: W,
301) -> Result<(), ImageError> {
302    assert_rgb_len(rgb, width, height);
303    WebPEncoder::new_lossless(writer).write_image(rgb, width, height, ExtendedColorType::Rgb8)
304}
305
306/// Wrap raw `width × height × 3` RGB bytes in a lossless WebP container `Vec<u8>`.
307///
308/// Available behind the `webp` cargo feature.
309#[cfg(feature = "webp")]
310pub fn rgb_to_webp(rgb: &[u8], width: u32, height: u32) -> Result<Vec<u8>, ImageError> {
311    let mut out = Vec::with_capacity(rgb.len() / 2);
312    rgb_to_webp_to_writer(rgb, width, height, &mut out)?;
313    Ok(out)
314}
315
316/// Encode raw `width × height × 3` RGB bytes as AVIF directly to a writer.
317///
318/// Available behind the `avif` cargo feature.
319///
320/// # Panics
321///
322/// Panics if `rgb.len() != (width * height * 3) as usize`.
323#[cfg(feature = "avif")]
324pub fn rgb_to_avif_to_writer<W: std::io::Write>(
325    rgb: &[u8],
326    width: u32,
327    height: u32,
328    writer: W,
329) -> Result<(), ImageError> {
330    assert_rgb_len(rgb, width, height);
331    AvifEncoder::new(writer).write_image(rgb, width, height, ExtendedColorType::Rgb8)
332}
333
334/// Wrap raw `width × height × 3` RGB bytes in an AVIF container `Vec<u8>`.
335///
336/// Available behind the `avif` cargo feature, which pulls in the pure-Rust
337/// [`ravif`](https://docs.rs/ravif) encoder.
338///
339/// **Encode-only:** [`decode_image`] cannot decode AVIF without the
340/// system `libdav1d` library. If you need to decode AVIF, enable
341/// `image/avif-native` in your own dependency declaration and provide
342/// libdav1d at link time.
343#[cfg(feature = "avif")]
344pub fn rgb_to_avif(rgb: &[u8], width: u32, height: u32) -> Result<Vec<u8>, ImageError> {
345    let mut out = Vec::with_capacity(rgb.len() / 4);
346    rgb_to_avif_to_writer(rgb, width, height, &mut out)?;
347    Ok(out)
348}
349
350// --- RGBA variants ---
351//
352// The DEM/heightmap path is RGB, but the same containers are handy for
353// 4-channel payloads (e.g. a watermask packed as RGBA). These mirror the
354// RGB functions exactly, differing only in `ExtendedColorType::Rgba8` and
355// the `× 4` length check.
356
357/// Encode raw `width × height × 4` RGBA bytes as PNG directly to a writer.
358///
359/// Available behind the `png` cargo feature.
360///
361/// # Panics
362///
363/// Panics if `rgba.len() != (width * height * 4) as usize`.
364#[cfg(feature = "png")]
365pub fn rgba_to_png_to_writer<W: std::io::Write>(
366    rgba: &[u8],
367    width: u32,
368    height: u32,
369    writer: W,
370) -> Result<(), ImageError> {
371    assert_rgba_len(rgba, width, height);
372    PngEncoder::new(writer).write_image(rgba, width, height, ExtendedColorType::Rgba8)
373}
374
375/// Wrap raw `width × height × 4` RGBA bytes in a PNG container `Vec<u8>`.
376///
377/// Available behind the `png` cargo feature.
378///
379/// # Errors
380///
381/// Returns [`ImageError`] if the underlying encoder fails (very rare for
382/// valid RGBA inputs — typically only OOM).
383#[cfg(feature = "png")]
384pub fn rgba_to_png(rgba: &[u8], width: u32, height: u32) -> Result<Vec<u8>, ImageError> {
385    let mut out = Vec::with_capacity(rgba.len());
386    rgba_to_png_to_writer(rgba, width, height, &mut out)?;
387    Ok(out)
388}
389
390/// Encode raw `width × height × 4` RGBA bytes as lossless WebP directly to a writer.
391///
392/// Available behind the `webp` cargo feature.
393///
394/// # Panics
395///
396/// Panics if `rgba.len() != (width * height * 4) as usize`.
397#[cfg(feature = "webp")]
398pub fn rgba_to_webp_to_writer<W: std::io::Write>(
399    rgba: &[u8],
400    width: u32,
401    height: u32,
402    writer: W,
403) -> Result<(), ImageError> {
404    assert_rgba_len(rgba, width, height);
405    WebPEncoder::new_lossless(writer).write_image(rgba, width, height, ExtendedColorType::Rgba8)
406}
407
408/// Wrap raw `width × height × 4` RGBA bytes in a lossless WebP container `Vec<u8>`.
409///
410/// Available behind the `webp` cargo feature.
411#[cfg(feature = "webp")]
412pub fn rgba_to_webp(rgba: &[u8], width: u32, height: u32) -> Result<Vec<u8>, ImageError> {
413    let mut out = Vec::with_capacity(rgba.len() / 2);
414    rgba_to_webp_to_writer(rgba, width, height, &mut out)?;
415    Ok(out)
416}
417
418/// Encode raw `width × height × 4` RGBA bytes as AVIF directly to a writer.
419///
420/// Available behind the `avif` cargo feature.
421///
422/// # Panics
423///
424/// Panics if `rgba.len() != (width * height * 4) as usize`.
425#[cfg(feature = "avif")]
426pub fn rgba_to_avif_to_writer<W: std::io::Write>(
427    rgba: &[u8],
428    width: u32,
429    height: u32,
430    writer: W,
431) -> Result<(), ImageError> {
432    assert_rgba_len(rgba, width, height);
433    AvifEncoder::new(writer).write_image(rgba, width, height, ExtendedColorType::Rgba8)
434}
435
436/// Wrap raw `width × height × 4` RGBA bytes in an AVIF container `Vec<u8>`.
437///
438/// Available behind the `avif` cargo feature (pure-Rust [`ravif`](https://docs.rs/ravif)
439/// encoder). Encode-only — see [`rgb_to_avif`] for the decode caveat.
440#[cfg(feature = "avif")]
441pub fn rgba_to_avif(rgba: &[u8], width: u32, height: u32) -> Result<Vec<u8>, ImageError> {
442    let mut out = Vec::with_capacity(rgba.len() / 4);
443    rgba_to_avif_to_writer(rgba, width, height, &mut out)?;
444    Ok(out)
445}
446
447/// Decode container bytes to raw RGB. The format is auto-detected from
448/// the bytes' header.
449///
450/// Only formats whose cargo features are enabled will be recognised —
451/// e.g. with just `png` on, this can decode PNG but not WebP. AVIF
452/// decoding additionally requires `image/avif-native` (libdav1d) which
453/// is not enabled by our `avif` feature.
454///
455/// Pixels with alpha are dropped (the `image` crate decodes to RGBA
456/// internally and we keep only the RGB channels).
457///
458/// # Errors
459///
460/// Returns [`ImageError`] if the bytes are not in a recognised format
461/// or the decoder fails.
462pub fn decode_image(bytes: &[u8]) -> Result<DecodedImage, ImageError> {
463    let reader = ImageReader::new(Cursor::new(bytes)).with_guessed_format()?;
464    let img = reader.decode()?;
465    let width = img.width();
466    let height = img.height();
467    let rgb = img.into_rgb8().into_raw();
468    Ok(DecodedImage { rgb, width, height })
469}
470
471#[track_caller]
472fn assert_rgb_len(rgb: &[u8], width: u32, height: u32) {
473    let expected = (width as usize) * (height as usize) * 3;
474    assert_eq!(
475        rgb.len(),
476        expected,
477        "rgb length mismatch: expected {expected}, got {}",
478        rgb.len()
479    );
480}
481
482#[cfg(any(feature = "png", feature = "webp", feature = "avif"))]
483#[track_caller]
484fn assert_rgba_len(rgba: &[u8], width: u32, height: u32) {
485    let expected = (width as usize) * (height as usize) * 4;
486    assert_eq!(
487        rgba.len(),
488        expected,
489        "rgba length mismatch: expected {expected}, got {}",
490        rgba.len()
491    );
492}
493
494#[cfg(test)]
495mod tests {
496    use super::*;
497    use crate::heightmap::{HeightmapFormat, decode, encode};
498
499    fn sample_rgb(width: u32, height: u32) -> Vec<u8> {
500        let elevations: Vec<f32> = (0..(width * height) as usize)
501            .map(|i| i as f32 * 10.0)
502            .collect();
503        encode(HeightmapFormat::Terrarium, &elevations, width, height)
504    }
505
506    #[test]
507    fn container_format_round_trips_through_from_str() {
508        for fmt in ContainerFormat::ALL {
509            let parsed: ContainerFormat = fmt.to_string().parse().unwrap();
510            assert_eq!(parsed, fmt);
511            // MIME alias also works.
512            let mime: ContainerFormat = fmt.mime_type().parse().unwrap();
513            assert_eq!(mime, fmt);
514        }
515        assert!("bogus".parse::<ContainerFormat>().is_err());
516    }
517
518    #[test]
519    fn is_enabled_reflects_features() {
520        assert_eq!(ContainerFormat::Png.is_enabled(), cfg!(feature = "png"));
521        assert_eq!(ContainerFormat::Webp.is_enabled(), cfg!(feature = "webp"));
522        assert_eq!(ContainerFormat::Avif.is_enabled(), cfg!(feature = "avif"));
523    }
524
525    #[test]
526    fn dispatch_returns_unsupported_for_disabled_features() {
527        let rgb = sample_rgb(4, 4);
528        for fmt in ContainerFormat::ALL {
529            let result = rgb_to_container(fmt, &rgb, 4, 4);
530            match (fmt.is_enabled(), &result) {
531                (true, Ok(_)) => {}
532                (false, Err(ContainerError::Unsupported(f))) => assert_eq!(*f, fmt),
533                other => panic!(
534                    "unexpected combination: enabled={:?} {other:?}",
535                    fmt.is_enabled()
536                ),
537            }
538        }
539    }
540
541    #[cfg(feature = "png")]
542    #[test]
543    fn png_roundtrip_through_codec() {
544        let width = 8u32;
545        let height = 8u32;
546        let elevations: Vec<f32> = (0..(width * height) as usize)
547            .map(|i| i as f32 * 10.0)
548            .collect();
549
550        for fmt in [
551            HeightmapFormat::Terrarium,
552            HeightmapFormat::Mapbox,
553            HeightmapFormat::Gsi,
554        ] {
555            let rgb = encode(fmt, &elevations, width, height);
556            let png = rgb_to_png(&rgb, width, height).unwrap();
557            assert_eq!(
558                &png[..8],
559                b"\x89PNG\r\n\x1a\n",
560                "{fmt} should produce PNG magic"
561            );
562            let DecodedImage {
563                rgb: rgb_back,
564                width: w2,
565                height: h2,
566            } = decode_image(&png).unwrap();
567            assert_eq!((w2, h2), (width, height));
568            assert_eq!(rgb_back, rgb);
569            let elev_back = decode(fmt, &rgb_back, width, height);
570            for (a, b) in elevations.iter().zip(&elev_back) {
571                assert!((a - b).abs() < 0.5, "{fmt}: {a} → {b}");
572            }
573        }
574    }
575
576    #[cfg(feature = "avif")]
577    #[test]
578    fn avif_encodes_to_valid_container() {
579        let rgb = sample_rgb(8, 8);
580        let avif = rgb_to_avif(&rgb, 8, 8).unwrap();
581        // AVIF files have an `ftypavif` brand in the first ISO BMFF box.
582        assert!(
583            avif.windows(8).any(|w| w == b"ftypavif"),
584            "expected AVIF brand in output"
585        );
586    }
587
588    #[cfg(all(feature = "webp", feature = "png"))]
589    #[test]
590    fn webp_roundtrip_through_codec() {
591        let rgb = sample_rgb(8, 8);
592        let webp = rgb_to_webp(&rgb, 8, 8).unwrap();
593        // WebP files start with "RIFF" .... "WEBP".
594        assert_eq!(&webp[..4], b"RIFF");
595        assert_eq!(&webp[8..12], b"WEBP");
596        let decoded = decode_image(&webp).unwrap();
597        assert_eq!((decoded.width, decoded.height), (8, 8));
598        assert_eq!(decoded.rgb, rgb);
599    }
600
601    #[cfg(any(feature = "png", feature = "webp", feature = "avif"))]
602    fn sample_rgba(width: u32, height: u32) -> Vec<u8> {
603        (0..(width * height) as usize)
604            .flat_map(|i| [(i % 256) as u8, (i % 200) as u8, (i % 100) as u8, 255])
605            .collect()
606    }
607
608    #[cfg(feature = "png")]
609    #[test]
610    fn rgba_png_encodes_and_decodes_dimensions() {
611        let rgba = sample_rgba(8, 8);
612        let png = rgba_to_png(&rgba, 8, 8).unwrap();
613        assert_eq!(&png[..8], b"\x89PNG\r\n\x1a\n");
614        // decode_image drops alpha → 3-channel RGB at the same dimensions.
615        let decoded = decode_image(&png).unwrap();
616        assert_eq!((decoded.width, decoded.height), (8, 8));
617        assert_eq!(decoded.rgb.len(), 8 * 8 * 3);
618    }
619
620    #[cfg(feature = "webp")]
621    #[test]
622    fn rgba_webp_has_riff_webp_magic() {
623        let rgba = sample_rgba(8, 8);
624        let webp = rgba_to_webp(&rgba, 8, 8).unwrap();
625        assert_eq!(&webp[..4], b"RIFF");
626        assert_eq!(&webp[8..12], b"WEBP");
627    }
628
629    #[cfg(feature = "png")]
630    #[test]
631    #[should_panic(expected = "rgba length mismatch")]
632    fn rgba_length_mismatch_panics() {
633        rgba_to_png(&[0u8; 10], 8, 8).unwrap();
634    }
635}
terrain_codec/heightmap/container.rs

terrain_codec/heightmap/
container.rs
