1const MAX_EDGE: u32 = 8192;
2
3#[derive(Clone, Debug, PartialEq)]
4pub struct Sky {
5 pub width: u32,
6 pub height: u32,
7 pub texels: Vec<[f32; 4]>,
8}
9
10#[derive(Clone, Debug, PartialEq)]
11pub struct Distribution {
12 pub cdf: Vec<f32>,
13 pub total: f32,
14}
15
16impl Sky {
17 pub fn parse(bytes: &[u8]) -> Result<Self, String> {
18 let (width, height, mut data) = header(bytes)?;
19 let mut texels = Vec::with_capacity(width as usize * height as usize);
20 let mut scan = vec![[0u8; 4]; width as usize];
21 for _ in 0..height {
22 read_scanline(&mut data, &mut scan)?;
23 for pixel in &scan {
24 texels.push(rgbe(*pixel));
25 }
26 }
27 Ok(Sky {
28 width,
29 height,
30 texels,
31 })
32 }
33
34 pub fn mean_luminance(&self) -> f32 {
35 self.texels.iter().map(luminance).sum::<f32>() / self.texels.len().max(1) as f32
36 }
37
38 pub fn mean_color(&self) -> [f32; 3] {
39 let mut sum = [0.0f64; 3];
40 for texel in &self.texels {
41 for (total, value) in sum.iter_mut().zip(texel) {
42 *total += f64::from(*value);
43 }
44 }
45 let count = self.texels.len().max(1) as f64;
46 sum.map(|total| (total / count) as f32)
47 }
48
49 pub fn capped_luminance(mut self, max: f32) -> Self {
50 for texel in &mut self.texels {
51 let lum = luminance(texel);
52 if lum > max {
53 let keep = max / lum;
54 texel[0] *= keep;
55 texel[1] *= keep;
56 texel[2] *= keep;
57 }
58 }
59 self
60 }
61
62 pub fn balanced_to(mut self, target: [f32; 3]) -> Self {
63 let mean = self.mean_color();
64 let want = luminance(&[target[0], target[1], target[2], 1.0]).max(1e-6);
65 let have = luminance(&[mean[0], mean[1], mean[2], 1.0]).max(1e-6);
66 let gain = [0, 1, 2].map(|k| (target[k] / want) / (mean[k] / have).max(1e-6));
67 for texel in &mut self.texels {
68 for (value, g) in texel.iter_mut().zip(gain) {
69 *value *= g;
70 }
71 }
72 self
73 }
74
75 pub fn scaled_to_mean_luminance(mut self, target: f32) -> Self {
76 let scale = target / self.mean_luminance().max(1e-4);
77 for texel in &mut self.texels {
78 texel[0] *= scale;
79 texel[1] *= scale;
80 texel[2] *= scale;
81 }
82 self
83 }
84
85 pub fn exposed(mut self, factor: f32) -> Self {
86 for texel in &mut self.texels {
87 texel[0] *= factor;
88 texel[1] *= factor;
89 texel[2] *= factor;
90 }
91 self
92 }
93
94 pub fn rotated(mut self, azimuth: f32) -> Self {
95 let w = self.width as usize;
96 if w == 0 || self.texels.len() != w * self.height as usize {
97 return self;
98 }
99 let shift = f64::from(azimuth) / std::f64::consts::TAU * w as f64;
100 let source = self.texels.clone();
101 for (row, out) in source.chunks(w).zip(self.texels.chunks_mut(w)) {
102 for (x, texel) in out.iter_mut().enumerate() {
103 let at = x as f64 - shift;
104 let base = at.floor();
105 let frac = (at - base) as f32;
106 let a = row[(base as i64).rem_euclid(w as i64) as usize];
107 let b = row[(base as i64 + 1).rem_euclid(w as i64) as usize];
108 for k in 0..4 {
109 texel[k] = a[k] + (b[k] - a[k]) * frac;
110 }
111 }
112 }
113 self
114 }
115
116 pub fn distribution(&self) -> Distribution {
117 let w = self.width as usize;
118 let h = self.height as usize;
119 if w == 0 || h == 0 || self.texels.len() != w * h {
120 return Distribution {
121 cdf: Vec::new(),
122 total: 0.0,
123 };
124 }
125 let mut cdf = vec![0.0f32; w * h + h];
126 let mut total = 0.0f64;
127 let mut rows = vec![0.0f64; h];
128 for y in 0..h {
129 let sin = ((y as f64 + 0.5) / h as f64 * std::f64::consts::PI).sin();
130 let mut acc = 0.0f64;
131 for x in 0..w {
132 let t = self.texels[y * w + x];
133 acc += (0.2126 * t[0] + 0.7152 * t[1] + 0.0722 * t[2]).max(0.0) as f64 * sin;
134 cdf[y * w + x] = acc as f32;
135 }
136 if acc > 0.0 {
137 for x in 0..w {
138 cdf[y * w + x] = (cdf[y * w + x] as f64 / acc) as f32;
139 }
140 } else {
141 for x in 0..w {
142 cdf[y * w + x] = (x + 1) as f32 / w as f32;
143 }
144 }
145 rows[y] = acc;
146 total += acc;
147 }
148 let mut acc = 0.0f64;
149 for y in 0..h {
150 acc += rows[y];
151 cdf[w * h + y] = if total > 0.0 {
152 (acc / total) as f32
153 } else {
154 (y + 1) as f32 / h as f32
155 };
156 }
157 Distribution {
158 cdf,
159 total: total as f32,
160 }
161 }
162}
163
164fn luminance(texel: &[f32; 4]) -> f32 {
165 0.2126 * texel[0] + 0.7152 * texel[1] + 0.0722 * texel[2]
166}
167
168fn header(bytes: &[u8]) -> Result<(u32, u32, &[u8]), String> {
169 let mut at = 0usize;
170 let line = |at: &mut usize| -> Result<String, String> {
171 if *at >= bytes.len() {
172 return Err("truncated hdr".into());
173 }
174 let start = *at;
175 while *at < bytes.len() && bytes[*at] != b'\n' {
176 *at += 1;
177 }
178 if *at >= bytes.len() {
179 return Err("truncated hdr".into());
180 }
181 let mut end = *at;
182 if end > start && bytes[end - 1] == b'\r' {
183 end -= 1;
184 }
185 let text = String::from_utf8_lossy(&bytes[start..end]).into_owned();
186 *at += 1;
187 Ok(text)
188 };
189 let magic = line(&mut at)?;
190 if !magic.starts_with("#?") {
191 return Err("not a Radiance HDR".into());
192 }
193 loop {
194 let header = line(&mut at)?;
195 if header.is_empty() {
196 break;
197 }
198 if header.starts_with("FORMAT")
199 && header.contains('=')
200 && !header.contains("32-bit_rle_rgbe")
201 {
202 return Err(format!("unsupported {header}"));
203 }
204 }
205 let size = line(&mut at)?;
206 let parts: Vec<&str> = size.split_whitespace().collect();
207 if parts.len() != 4 || parts[0] != "-Y" || parts[2] != "+X" {
208 return Err(format!("unsupported orientation {size}"));
209 }
210 let height: u32 = parts[1].parse().map_err(|_| "bad hdr height")?;
211 let width: u32 = parts[3].parse().map_err(|_| "bad hdr width")?;
212 if width == 0 || height == 0 || width > MAX_EDGE || height > MAX_EDGE {
213 return Err(format!(
214 "hdr is {width} by {height}, the limit is {MAX_EDGE}"
215 ));
216 }
217 Ok((width, height, &bytes[at..]))
218}
219
220fn read_scanline(data: &mut &[u8], scan: &mut [[u8; 4]]) -> Result<(), String> {
221 let width = scan.len() as u32;
222 let rle = (8..32768).contains(&width)
223 && data.len() >= 4
224 && data[0] == 2
225 && data[1] == 2
226 && data[2] & 0x80 == 0
227 && (u32::from(data[2]) << 8 | u32::from(data[3])) == width;
228 if rle {
229 *data = &data[4..];
230 for channel in 0..4 {
231 read_channel(data, channel, scan)?;
232 }
233 return Ok(());
234 }
235 for pixel in scan.iter_mut() {
236 let bytes = take(data, 4)?;
237 pixel.copy_from_slice(bytes);
238 }
239 Ok(())
240}
241
242fn read_channel(data: &mut &[u8], channel: usize, scan: &mut [[u8; 4]]) -> Result<(), String> {
243 let width = scan.len();
244 let mut x = 0;
245 while x < width {
246 let count = take(data, 1)?[0] as usize;
247 if count == 0 {
248 return Err("bad hdr run".into());
249 }
250 if count > 128 {
251 let run = count - 128;
252 if x + run > width {
253 return Err("hdr run overruns the row".into());
254 }
255 let value = take(data, 1)?[0];
256 for pixel in scan.iter_mut().skip(x).take(run) {
257 pixel[channel] = value;
258 }
259 x += run;
260 } else if x + count > width {
261 return Err("hdr literal overruns the row".into());
262 } else {
263 let literal = take(data, count)?;
264 for (pixel, value) in scan.iter_mut().skip(x).take(count).zip(literal) {
265 pixel[channel] = *value;
266 }
267 x += count;
268 }
269 }
270 Ok(())
271}
272
273fn take<'a>(data: &mut &'a [u8], len: usize) -> Result<&'a [u8], String> {
274 if data.len() < len {
275 return Err("truncated hdr".into());
276 }
277 let (head, tail) = data.split_at(len);
278 *data = tail;
279 Ok(head)
280}
281
282fn rgbe(pixel: [u8; 4]) -> [f32; 4] {
283 if pixel[3] == 0 {
284 [0.0, 0.0, 0.0, 1.0]
285 } else {
286 let scale = 2f32.powi(i32::from(pixel[3]) - 136);
287 [
288 f32::from(pixel[0]) * scale,
289 f32::from(pixel[1]) * scale,
290 f32::from(pixel[2]) * scale,
291 1.0,
292 ]
293 }
294}
295
296#[cfg(test)]
297mod tests {
298 use super::*;
299
300 fn sums_to_one(sky: &Sky) {
301 let distribution = sky.distribution();
302 let w = sky.width as usize;
303 let h = sky.height as usize;
304 assert_eq!(distribution.cdf.len(), w * h + h);
305 let marginal = &distribution.cdf[w * h..];
306 let mut sum = 0.0f64;
307 for y in 0..h {
308 let row_mass = f64::from(marginal[y])
309 - if y == 0 {
310 0.0
311 } else {
312 f64::from(marginal[y - 1])
313 };
314 assert!(row_mass >= -1e-6);
315 for x in 0..w {
316 let here = distribution.cdf[y * w + x];
317 let prev = if x == 0 {
318 0.0
319 } else {
320 distribution.cdf[y * w + x - 1]
321 };
322 assert!(here + 1e-5 >= prev);
323 sum += (f64::from(here) - f64::from(prev)) * row_mass;
324 }
325 assert!((distribution.cdf[y * w + w - 1] - 1.0).abs() < 1e-5);
326 }
327 assert!((marginal[h - 1] - 1.0).abs() < 1e-5);
328 assert!((sum - 1.0).abs() < 1e-4, "{sum}");
329 }
330
331 #[test]
332 fn a_small_hdr_decodes_and_its_cdf_sums_to_one() {
333 let sky = Sky::parse(include_bytes!("../tests/data/sky.hdr")).unwrap();
334 assert_eq!((sky.width, sky.height), (8, 2));
335 assert_eq!(sky.texels[0], [1.0, 0.0, 0.0, 1.0]);
336 assert_eq!(sky.texels[8], [0.0, 0.5, 0.0, 1.0]);
337 let distribution = sky.distribution();
338 let row0 = distribution.cdf[8 * 2] - 0.0;
339 let row1 = distribution.cdf[8 * 2 + 1] - distribution.cdf[8 * 2];
340 assert!(row1 > row0);
341 sums_to_one(&sky);
342
343 let flat = Sky::parse(include_bytes!("../tests/data/flat.hdr")).unwrap();
344 assert_eq!(
345 flat.texels,
346 vec![[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]
347 );
348 sums_to_one(&flat);
349 assert!(Sky::parse(b"not a sky").is_err());
350 }
351
352 #[test]
353 fn a_run_header_of_128_is_a_full_literal() {
354 let mut bytes = vec![128u8];
355 bytes.extend((0..128).map(|value| value as u8));
356 bytes.extend([130u8, 7]);
357 let mut data = bytes.as_slice();
358 let mut scan = vec![[0u8; 4]; 130];
359 read_channel(&mut data, 1, &mut scan).unwrap();
360 assert!(data.is_empty());
361 assert!((0..128).all(|x| scan[x][1] == x as u8));
362 assert_eq!(scan[128][1], 7);
363 assert_eq!(scan[129][1], 7);
364 let mut zero: &[u8] = &[0u8];
365 assert!(read_channel(&mut zero, 0, &mut [[0u8; 4]; 4]).is_err());
366 }
367
368 #[test]
369 fn a_black_row_still_has_a_cdf_that_ends_at_one() {
370 let sky = Sky {
371 width: 1,
372 height: 2,
373 texels: vec![[0.0, 0.0, 0.0, 1.0], [1.0, 1.0, 1.0, 1.0]],
374 };
375 sums_to_one(&sky);
376 assert_eq!(sky.distribution().cdf[0], 1.0);
377 }
378
379 fn grid(width: u32, height: u32, rgb: &[[f32; 3]]) -> Sky {
380 Sky {
381 width,
382 height,
383 texels: rgb.iter().map(|c| [c[0], c[1], c[2], 1.0]).collect(),
384 }
385 }
386
387 fn close(a: f32, b: f32, tolerance: f32) {
388 assert!(
389 (a - b).abs() <= tolerance * b.abs().max(1e-6),
390 "{a} against {b}"
391 );
392 }
393
394 fn bright_sky() -> Sky {
395 let mut texels = Vec::new();
396 for y in 0..8u32 {
397 for x in 0..16u32 {
398 let n = (x * 7 + y * 13) % 17;
399 let sun = if (x, y) == (5, 2) || (x, y) == (6, 2) {
400 400.0 + x as f32
401 } else {
402 0.0
403 };
404 texels.push([
405 0.9 + n as f32 * 0.31 + sun,
406 0.4 + n as f32 * 0.17 + sun * 0.8,
407 0.1 + (16 - n) as f32 * 0.05 + sun * 0.5,
408 1.0,
409 ]);
410 }
411 }
412 Sky {
413 width: 16,
414 height: 8,
415 texels,
416 }
417 }
418
419 fn prepared_port(sky: &Sky, cap: f32, mean: f32, balance: Option<[f32; 3]>) -> Vec<[f32; 4]> {
420 let mut texels = sky.texels.clone();
421 let lum = |t: &[f32; 4]| 0.2126 * t[0] + 0.7152 * t[1] + 0.0722 * t[2];
422 for t in &mut texels {
423 let l = lum(t);
424 if l > cap {
425 let k = cap / l;
426 t[0] *= k;
427 t[1] *= k;
428 t[2] *= k;
429 }
430 }
431 if let Some(target) = balance {
432 let mut sum = [0.0f64; 3];
433 for t in &texels {
434 for (m, v) in sum.iter_mut().zip(t) {
435 *m += f64::from(*v);
436 }
437 }
438 let n = texels.len().max(1) as f64;
439 let mean_colour = sum.map(|v| (v / n) as f32);
440 let want = lum(&[target[0], target[1], target[2], 1.0]).max(1e-6);
441 let have = lum(&[mean_colour[0], mean_colour[1], mean_colour[2], 1.0]).max(1e-6);
442 let gain = [0, 1, 2].map(|k| (target[k] / want) / (mean_colour[k] / have).max(1e-6));
443 for t in &mut texels {
444 for (v, g) in t.iter_mut().zip(gain) {
445 *v *= g;
446 }
447 }
448 }
449 let have = texels.iter().map(lum).sum::<f32>() / texels.len().max(1) as f32;
450 let scale = mean / have.max(1e-4);
451 for t in &mut texels {
452 t[0] *= scale;
453 t[1] *= scale;
454 t[2] *= scale;
455 }
456 texels
457 }
458
459 fn prepare(sky: Sky, cap: f32, mean: f32, balance: Option<[f32; 3]>) -> Sky {
460 let sky = sky.capped_luminance(cap);
461 let sky = match balance {
462 Some(target) => sky.balanced_to(target),
463 None => sky,
464 };
465 sky.scaled_to_mean_luminance(mean)
466 }
467
468 #[test]
469 fn capping_scales_a_bright_texel_to_the_cap_and_keeps_its_hue() {
470 let sky = grid(2, 1, &[[120.0, 60.0, 30.0], [3.0, 6.0, 9.0]]);
471 let lum = luminance(&sky.texels[0]);
472 let capped = sky.clone().capped_luminance(40.0);
473 let keep = 40.0 / lum;
474 close(capped.texels[0][0], 120.0 * keep, 1e-6);
475 close(capped.texels[0][1], 60.0 * keep, 1e-6);
476 close(capped.texels[0][2], 30.0 * keep, 1e-6);
477 close(luminance(&capped.texels[0]), 40.0, 1e-5);
478 close(capped.texels[0][0] / capped.texels[0][1], 2.0, 1e-5);
479 assert_eq!(capped.texels[0][3], 1.0);
480 assert_eq!(capped.texels[1], sky.texels[1]);
481 }
482
483 #[test]
484 fn mean_luminance_and_colour_average_every_texel() {
485 let sky = grid(2, 1, &[[1.0, 2.0, 4.0], [3.0, 2.0, 0.0]]);
486 assert_eq!(sky.mean_color(), [2.0, 2.0, 2.0]);
487 close(sky.mean_luminance(), 2.0, 1e-6);
488 }
489
490 #[test]
491 fn balancing_moves_the_mean_to_the_target_hue_at_the_same_brightness() {
492 let sky = grid(2, 1, &[[1.0, 2.0, 4.0], [3.0, 2.0, 0.0]]);
493 let target = [4.0, 2.0, 1.0];
494 let want = 0.2126 * 4.0 + 0.7152 * 2.0 + 0.0722;
495 let balanced = sky.clone().balanced_to(target);
496 close(balanced.texels[0][0], 4.0 / want, 1e-5);
497 close(balanced.texels[0][1], 4.0 / want, 1e-5);
498 close(balanced.texels[0][2], 4.0 / want, 1e-5);
499 close(balanced.texels[1][0], 12.0 / want, 1e-5);
500 assert_eq!(balanced.texels[1][2], 0.0);
501 assert_eq!(balanced.texels[0][3], 1.0);
502 let mean = balanced.mean_color();
503 close(mean[0] / mean[2], 4.0, 1e-5);
504 close(mean[1] / mean[2], 2.0, 1e-5);
505 close(balanced.mean_luminance(), sky.mean_luminance(), 1e-5);
506 }
507
508 #[test]
509 fn scaling_to_a_mean_luminance_multiplies_every_texel_alike() {
510 let sky = grid(2, 1, &[[2.0, 2.0, 2.0], [0.0, 0.0, 0.0]]);
511 let scaled = sky.scaled_to_mean_luminance(0.4);
512 close(scaled.texels[0][0], 0.8, 1e-5);
513 close(scaled.texels[0][2], 0.8, 1e-5);
514 assert_eq!(scaled.texels[1], [0.0, 0.0, 0.0, 1.0]);
515 close(scaled.mean_luminance(), 0.4, 1e-5);
516 }
517
518 #[test]
519 fn exposure_multiplies_colour_and_not_alpha() {
520 let sky = grid(1, 1, &[[1.0, 2.0, 3.0]]).exposed(0.5);
521 assert_eq!(sky.texels[0], [0.5, 1.0, 1.5, 1.0]);
522 }
523
524 #[test]
525 fn rotation_by_whole_and_half_texels_shifts_the_columns() {
526 let sky = grid(
527 4,
528 1,
529 &[
530 [1.0, 0.0, 0.0],
531 [2.0, 0.0, 0.0],
532 [3.0, 0.0, 0.0],
533 [4.0, 0.0, 0.0],
534 ],
535 );
536 let whole = sky.clone().rotated(std::f32::consts::FRAC_PI_2);
537 let reds: Vec<f32> = whole.texels.iter().map(|t| t[0]).collect();
538 for (got, want) in reds.iter().zip([4.0, 1.0, 2.0, 3.0]) {
539 close(*got, want, 1e-5);
540 }
541 let half = sky.clone().rotated(std::f32::consts::FRAC_PI_4);
542 let reds: Vec<f32> = half.texels.iter().map(|t| t[0]).collect();
543 for (got, want) in reds.iter().zip([2.5, 1.5, 2.5, 3.5]) {
544 close(*got, want, 1e-5);
545 }
546 let back = sky.clone().rotated(std::f32::consts::TAU);
547 for (got, want) in back.texels.iter().zip(&sky.texels) {
548 close(got[0], want[0], 1e-5);
549 }
550 let backwards = sky.rotated(-std::f32::consts::FRAC_PI_2);
551 let reds: Vec<f32> = backwards.texels.iter().map(|t| t[0]).collect();
552 for (got, want) in reds.iter().zip([2.0, 3.0, 4.0, 1.0]) {
553 close(*got, want, 1e-5);
554 }
555 }
556
557 #[test]
558 fn a_cap_balance_and_mean_match_a_straight_port() {
559 let sky = bright_sky();
560 let cap = 40.0;
561 let mean = 0.25;
562 assert!(sky.texels.iter().any(|t| luminance(t) > cap));
563 for balance in [None, Some([0.7, 0.8, 0.6])] {
564 let want = prepared_port(&sky, cap, mean, balance);
565 let got = prepare(sky.clone(), cap, mean, balance);
566 assert_eq!((got.width, got.height), (16, 8));
567 assert_eq!(got.texels.len(), want.len());
568 for (g, w) in got.texels.iter().zip(&want) {
569 for k in 0..4 {
570 close(g[k], w[k], 1e-6);
571 }
572 }
573 close(got.mean_luminance(), mean, 1e-5);
574 }
575 }
576
577 #[test]
578 fn the_distribution_reflects_the_prepared_texels() {
579 let raw = grid(1, 2, &[[1.0, 1.0, 1.0], [3.0, 3.0, 3.0]]);
580 let sin = (std::f64::consts::PI / 4.0).sin();
581 close(raw.distribution().total, (4.0 * sin) as f32, 1e-5);
582 let prepared = raw
583 .clone()
584 .capped_luminance(10.0)
585 .scaled_to_mean_luminance(0.5);
586 let scale = 0.5 / raw.mean_luminance();
587 close(
588 prepared.distribution().total,
589 (4.0 * sin * f64::from(scale)) as f32,
590 1e-5,
591 );
592 close(
593 prepared.distribution().total,
594 raw.distribution().total * scale,
595 1e-5,
596 );
597 sums_to_one(&prepared);
598 let bright = prepare(bright_sky(), 40.0, 0.25, None);
599 sums_to_one(&bright);
600 let total = bright.distribution().total;
601 let mut expected = 0.0f64;
602 for y in 0..8usize {
603 let sin = ((y as f64 + 0.5) / 8.0 * std::f64::consts::PI).sin();
604 for x in 0..16usize {
605 expected += f64::from(luminance(&bright.texels[y * 16 + x])) * sin;
606 }
607 }
608 close(total, expected as f32, 1e-5);
609 }
610}