1use serde::Serialize;
14
15use super::encoding::{Axis, CubeFace, NormalEncoding, NormalSpace, ScalarEncoding, narrow};
16use super::roles::{Occluded, OutputRole};
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
21pub struct Place {
22 pub output: usize,
24 pub lanes: u8,
26}
27
28impl Place {
29 pub fn first(&self) -> usize {
31 (self.lanes.trailing_zeros() as usize).min(3)
32 }
33
34 pub fn width(&self) -> u32 {
36 self.lanes.count_ones()
37 }
38
39 pub fn lane(&self, values: [f32; 4], k: usize) -> f32 {
41 values.get(self.first().saturating_add(k)).copied().unwrap_or(0.0)
42 }
43}
44
45#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
47pub struct Term {
48 pub place: Place,
50 pub encoding: ScalarEncoding,
52}
53
54impl Term {
55 pub fn decode(&self, values: [f32; 4], k: usize) -> f32 {
57 self.encoding.decode(self.place.lane(values, k))
58 }
59
60 pub fn colour(&self, values: [f32; 4]) -> [f32; 3] {
62 [0, 1, 2].map(|k| self.decode(values, k))
63 }
64}
65
66#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
69#[serde(rename_all = "snake_case")]
70pub enum Base {
71 Albedo(Place),
73 Colour(Place),
75}
76
77impl Base {
78 pub fn place(&self) -> Place {
80 match *self {
81 Base::Albedo(p) | Base::Colour(p) => p,
82 }
83 }
84
85 pub fn colour(&self, values: [f32; 4]) -> [f32; 3] {
87 let place = self.place();
88 [0, 1, 2].map(|k| place.lane(values, k))
89 }
90}
91
92#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
94pub struct Normal {
95 pub place: Place,
97 pub encoding: NormalEncoding,
99 pub space: NormalSpace,
101}
102
103impl Normal {
104 pub fn width(encoding: NormalEncoding) -> u32 {
106 match encoding {
107 NormalEncoding::Xy | NormalEncoding::XyUnorm | NormalEncoding::Octahedral { .. } => 2,
108 NormalEncoding::Xyz
109 | NormalEncoding::XyzUnorm
110 | NormalEncoding::XyzBiased { .. }
111 | NormalEncoding::OctahedralDiamond => 3,
112 NormalEncoding::CubeFace { .. } => 4,
113 }
114 }
115
116 pub fn vector(&self, values: [f32; 4]) -> [f32; 3] {
124 let lane = |k: usize| self.place.lane(values, k);
125 match self.encoding {
126 NormalEncoding::Xyz => [lane(0), lane(1), lane(2)],
127 NormalEncoding::XyzUnorm => [0, 1, 2].map(|k| (lane(k) - 0.5) * 2.0),
128 NormalEncoding::XyzBiased { bias_bits } => {
129 let bias = f32::from_bits(bias_bits);
130 [0, 1, 2].map(|k| (lane(k) - bias) * 2.0)
131 }
132 NormalEncoding::Xy => with_z(lane(0), lane(1)),
133 NormalEncoding::XyUnorm => with_z(lane(0) * 2.0 - 1.0, lane(1) * 2.0 - 1.0),
134 NormalEncoding::Octahedral { read, .. } => unfold(read.apply(lane(0)), read.apply(lane(1))),
135 NormalEncoding::OctahedralDiamond => {
136 let (n, _) = NormalEncoding::decode_octahedral_diamond([lane(0), lane(1), lane(2)]);
137 n.map(narrow)
138 }
139 NormalEncoding::CubeFace { minor, .. } => {
140 let thirds = (lane(3) * 3.0).round();
142 let axis = if thirds >= 2.0 {
143 Axis::Z
144 } else if thirds >= 1.0 {
145 Axis::Y
146 } else {
147 Axis::X
148 };
149 let negative = lane(2) > 0.5;
150 let face = CubeFace::new(axis, negative);
151 let (sc, tc) = (
152 (lane(0) - minor.bias()) / minor.scale(),
153 (lane(1) - minor.bias()) / minor.scale(),
154 );
155 let [(_, first_sign), (_, second_sign)] = face.minor;
156 let depth = (1.0 - sc * sc - tc * tc).max(0.0).sqrt() * if negative { -1.0 } else { 1.0 };
157 let n = face.place(depth, sc * narrow(first_sign), tc * narrow(second_sign), 0.0);
158 normalised(n).unwrap_or(n)
159 }
160 }
161 }
162
163 pub fn decode(&self, values: [f32; 4]) -> Option<[f32; 3]> {
165 normalised(self.vector(values))
166 }
167}
168
169fn normalised(n: [f32; 3]) -> Option<[f32; 3]> {
171 let length = n.iter().map(|x| x * x).sum::<f32>().sqrt();
172 (length > 1e-8).then(|| n.map(|x| x / length))
173}
174
175fn with_z(x: f32, y: f32) -> [f32; 3] {
177 [x, y, (1.0 - x * x - y * y).max(0.0).sqrt()]
178}
179
180fn unfold(x: f32, y: f32) -> [f32; 3] {
183 let z = 1.0 - x.abs() - y.abs();
184 let sign = |v: f32| if v < 0.0 { -1.0 } else { 1.0 };
185 let (x, y) = if z < 0.0 {
186 ((1.0 - y.abs()) * sign(x), (1.0 - x.abs()) * sign(y))
187 } else {
188 (x, y)
189 };
190 let n = [x, y, z];
191 normalised(n).unwrap_or(n)
192}
193
194#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
196#[serde(rename_all = "snake_case")]
197pub enum Roughness {
198 Roughness(Term),
200 Gloss(Term),
202 Exponent(Term),
207}
208
209impl Roughness {
210 pub fn term(&self) -> Term {
212 match *self {
213 Roughness::Roughness(t) | Roughness::Gloss(t) | Roughness::Exponent(t) => t,
214 }
215 }
216
217 pub fn decode(&self, values: [f32; 4]) -> f32 {
219 let value = self.term().decode(values, 0);
220 match self {
221 Roughness::Roughness(_) => value,
222 Roughness::Gloss(_) => 1.0 - value,
223 Roughness::Exponent(_) => (2.0 / (value.max(0.0) + 2.0)).powf(0.25),
224 }
225 }
226}
227
228#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize)]
230#[serde(rename_all = "snake_case")]
231pub enum Emission {
232 Colour(Term),
234 OfAlbedo(Term),
236}
237
238impl Emission {
239 pub fn term(&self) -> Term {
241 match *self {
242 Emission::Colour(t) | Emission::OfAlbedo(t) => t,
243 }
244 }
245
246 pub fn decode(&self, values: [f32; 4], albedo: [f32; 3]) -> [f32; 3] {
248 match self {
249 Emission::Colour(t) => t.colour(values),
250 Emission::OfAlbedo(t) => {
251 let e = t.decode(values, 0);
252 albedo.map(|a| a * e)
253 }
254 }
255 }
256}
257
258#[derive(Clone, Debug, Default, PartialEq, Serialize)]
260pub struct Surface {
261 pub base: Option<Base>,
263 pub normal: Option<Normal>,
265 pub roughness: Option<Roughness>,
267 pub metallic: Option<Term>,
270 pub reflectance: Option<Term>,
273 pub specular: Option<Term>,
275 pub emission: Option<Emission>,
277 pub occlusion: Option<Term>,
280 pub ambient_occlusion: Option<Term>,
283 pub refused: Vec<String>,
286}
287
288impl Surface {
289 pub fn read<'a>(outputs: impl IntoIterator<Item = (u8, &'a OutputRole)>) -> Surface {
292 let mut s = Surface::default();
293 let mut colour = None;
294 let mut refused = Vec::new();
295 for (output, (lanes, role)) in outputs.into_iter().enumerate() {
296 let place = Place { output, lanes };
297 let width = place.width();
298 let mut fits = |what: &str, wanted: &[u32]| -> bool {
299 if wanted.contains(&width) {
300 return true;
301 }
302 refused.push(format!(
303 "output {output}: {what} in {width} lanes ({lanes:#06b}), not {}",
304 wanted.iter().map(u32::to_string).collect::<Vec<_>>().join(" or ")
305 ));
306 false
307 };
308 let term = |encoding: ScalarEncoding| Term { place, encoding };
309 let linear = term(ScalarEncoding::LINEAR);
310 match *role {
311 OutputRole::Albedo if s.base.is_none() && fits("an albedo", &[3, 4]) => {
312 s.base = Some(Base::Albedo(place));
313 }
314 OutputRole::Colour if colour.is_none() && fits("a finished colour", &[3, 4]) => {
315 colour = Some(Base::Colour(place));
316 }
317 OutputRole::Normal { encoding, space }
318 if s.normal.is_none() && fits("a normal", &[Normal::width(encoding)]) =>
319 {
320 s.normal = Some(Normal {
321 place,
322 encoding,
323 space,
324 });
325 }
326 OutputRole::Roughness if s.roughness.is_none() && fits("a roughness", &[1]) => {
327 s.roughness = Some(Roughness::Roughness(linear));
328 }
329 OutputRole::Gloss { encoding } if s.roughness.is_none() && fits("a gloss", &[1]) => {
330 s.roughness = Some(Roughness::Gloss(term(encoding)));
331 }
332 OutputRole::SpecularExponent { encoding }
333 if s.roughness.is_none() && fits("a specular exponent", &[1]) =>
334 {
335 s.roughness = Some(Roughness::Exponent(term(encoding)));
336 }
337 OutputRole::Metallic if s.metallic.is_none() && fits("a metallic", &[1]) => {
338 s.metallic = Some(linear);
339 }
340 OutputRole::Reflectance { encoding }
341 if s.reflectance.is_none() && fits("a reflectance", &[3]) =>
342 {
343 s.reflectance = Some(term(encoding));
344 }
345 OutputRole::SpecularIntensity { encoding }
346 if s.specular.is_none() && fits("a specular intensity", &[1]) =>
347 {
348 s.specular = Some(term(encoding));
349 }
350 OutputRole::Emission { encoding } if s.emission.is_none() && fits("an emission", &[1, 3]) => {
351 s.emission = Some(if width == 1 {
352 Emission::OfAlbedo(term(encoding))
353 } else {
354 Emission::Colour(term(encoding))
355 });
356 }
357 OutputRole::Occlusion {
358 of: Occluded::AllLight,
359 encoding,
360 } if s.occlusion.is_none() && fits("an occlusion", &[1]) => {
361 s.occlusion = Some(term(encoding));
362 }
363 OutputRole::Occlusion {
364 of: Occluded::NaturalAmbient,
365 encoding,
366 } if s.ambient_occlusion.is_none() && fits("an occlusion", &[1]) => {
367 s.ambient_occlusion = Some(term(encoding));
368 }
369 _ => {}
370 }
371 }
372 s.base = s.base.or(colour);
373 if s.occlusion.is_some() {
374 s.ambient_occlusion = None;
375 }
376 s.refused = refused;
377 s
378 }
379}
380
381pub const WGSL: &str = include_str!("surface.wgsl");
392
393pub const ROWS: usize = 9;
395pub const BASE: usize = 0;
397pub const NORMAL: usize = 1;
399pub const ROUGHNESS: usize = 2;
401pub const METALLIC: usize = 3;
403pub const REFLECTANCE: usize = 4;
405pub const SPECULAR: usize = 5;
407pub const EMISSION: usize = 6;
409pub const OCCLUSION: usize = 7;
411pub const AMBIENT_OCCLUSION: usize = 8;
414
415pub fn normal_kind(encoding: NormalEncoding) -> (u32, [u32; 2]) {
417 match encoding {
418 NormalEncoding::Xyz => (1, [0, 0]),
419 NormalEncoding::XyzUnorm => (2, [0, 0]),
420 NormalEncoding::XyzBiased { bias_bits } => (3, [bias_bits, 0]),
421 NormalEncoding::Xy => (4, [0, 0]),
422 NormalEncoding::XyUnorm => (5, [0, 0]),
423 NormalEncoding::Octahedral { read, .. } => (6, [read.scale_bits, read.bias_bits]),
424 NormalEncoding::OctahedralDiamond => (7, [0, 0]),
425 NormalEncoding::CubeFace { minor, .. } => (8, [minor.scale_bits, minor.bias_bits]),
426 }
427}
428
429fn layout(place: Place, kind: u32, encoding: ScalarEncoding) -> u32 {
431 (u32::try_from(place.first()).unwrap_or(3) & 0x3)
432 | (place.width() & 0xf) << 4
433 | (kind & 0xff) << 8
434 | u32::from(encoding.power) << 16
435 | u32::from(encoding.root) << 24
436}
437
438pub fn rows(surface: &Surface, target: impl Fn(usize) -> Option<u32>) -> [[u32; 4]; ROWS] {
443 let mut rows = [[0u32; 4]; ROWS];
444 let at = |output: usize| target(output).map(|t| t.saturating_add(1));
445 let term_row = |t: &Term, kind: u32| -> [u32; 4] {
446 at(t.place.output).map_or([0; 4], |x| {
447 [x, layout(t.place, kind, t.encoding), t.encoding.scale_bits, 0]
448 })
449 };
450 if let Some(base) = &surface.base {
451 let (place, kind) = match *base {
452 Base::Albedo(p) => (p, 1),
453 Base::Colour(p) => (p, 2),
454 };
455 if let Some(x) = at(place.output) {
456 rows[BASE] = [x, layout(place, kind, ScalarEncoding::LINEAR), 0, 0];
457 }
458 }
459 if let Some(n) = &surface.normal {
460 let (kind, [a, b]) = normal_kind(n.encoding);
461 if let Some(x) = at(n.place.output) {
462 rows[NORMAL] = [x, layout(n.place, kind, ScalarEncoding::LINEAR), a, b];
463 }
464 }
465 if let Some(r) = &surface.roughness {
466 let kind = match r {
467 Roughness::Roughness(_) => 1,
468 Roughness::Gloss(_) => 2,
469 Roughness::Exponent(_) => 3,
470 };
471 rows[ROUGHNESS] = term_row(&r.term(), kind);
472 }
473 for (row, term) in [
474 (METALLIC, &surface.metallic),
475 (REFLECTANCE, &surface.reflectance),
476 (SPECULAR, &surface.specular),
477 (OCCLUSION, &surface.occlusion),
478 (AMBIENT_OCCLUSION, &surface.ambient_occlusion),
479 ] {
480 if let Some(t) = term
481 && let Some(slot) = rows.get_mut(row)
482 {
483 *slot = term_row(t, 1);
484 }
485 }
486 if let Some(e) = &surface.emission {
487 let kind = match e {
488 Emission::Colour(_) => 1,
489 Emission::OfAlbedo(_) => 2,
490 };
491 rows[EMISSION] = term_row(&e.term(), kind);
492 }
493 rows
494}
495
496#[cfg(test)]
497#[allow(clippy::unwrap_used, reason = "hand-worked values")]
498mod tests {
499 use super::*;
500 use crate::shader::LaneAffine;
501
502 fn close(a: [f32; 3], b: [f32; 3]) -> bool {
503 within(a, b, 1e-4)
504 }
505
506 fn within(a: [f32; 3], b: [f32; 3], tolerance: f32) -> bool {
507 a.iter().zip(&b).all(|(x, y)| (x - y).abs() < tolerance)
508 }
509
510 fn normal(encoding: NormalEncoding, lanes: u8) -> Normal {
511 Normal {
512 place: Place { output: 0, lanes },
513 encoding,
514 space: NormalSpace::World,
515 }
516 }
517
518 #[test]
521 fn each_normal_encoding_decodes_what_it_stores() {
522 let n = [0.6f32, 0.0, 0.8];
523 let unorm = n.map(|c| c * 0.5 + 0.5);
524 assert!(close(
525 normal(NormalEncoding::Xyz, 0b0111).vector([n[0], n[1], n[2], 9.0]),
526 n
527 ));
528 assert!(close(
529 normal(NormalEncoding::XyzUnorm, 0b0111).vector([unorm[0], unorm[1], unorm[2], 0.0]),
530 n
531 ));
532 let bias = 0.500_488_76f32;
533 let biased = n.map(|c| c * 0.5 + bias);
534 assert!(close(
535 normal(
536 NormalEncoding::XyzBiased {
537 bias_bits: bias.to_bits()
538 },
539 0b0111
540 )
541 .vector([biased[0], biased[1], biased[2], 0.0]),
542 n
543 ));
544 let long = normal(NormalEncoding::Xyz, 0b0111);
546 assert_eq!(long.vector([0.0, 2.0, 0.0, 0.0]), [0.0, 2.0, 0.0]);
547 assert_eq!(long.decode([0.0, 2.0, 0.0, 0.0]), Some([0.0, 1.0, 0.0]));
548 assert_eq!(long.decode([0.0; 4]), None);
549 assert!(close(
551 normal(NormalEncoding::Xy, 0b1100).vector([9.0, 9.0, 0.6, 0.0]),
552 n
553 ));
554 assert!(close(
555 normal(NormalEncoding::XyUnorm, 0b1100).vector([9.0, 9.0, 0.8, 0.5]),
556 n
557 ));
558 let e = 0.6f32 / 1.4;
560 let octahedral = normal(
561 NormalEncoding::Octahedral {
562 stored: LaneAffine::UNORM,
563 read: LaneAffine::SIGNED,
564 },
565 0b0011,
566 );
567 assert!(close(octahedral.vector([e * 0.5 + 0.5, 0.5, 0.0, 0.0]), n));
568 let below = octahedral.vector([1.0, (1.0 - e) * 0.5 + 0.5, 0.0, 0.0]);
570 assert!(close(below, [0.6, 0.0, -0.8]), "{below:?}");
571 }
572
573 #[test]
575 fn the_packed_normal_encodings_decode_what_their_encoders_store() {
576 let n = [0.36f64, -0.48, 0.8];
577 let bytes = NormalEncoding::encode_octahedral_diamond(n, 5);
578 let lanes = bytes.map(|b| f32::from(b) / 255.0);
579 let diamond =
580 normal(NormalEncoding::OctahedralDiamond, 0b0111).decode([lanes[0], lanes[1], lanes[2], 0.0]);
581 assert!(within(diamond.unwrap(), n.map(narrow), 0.01), "{diamond:?}");
582 let (minor, payload) = (LaneAffine::UNORM, LaneAffine::new(0.25, 0.1));
583 for n in [[0.36f32, -0.48, 0.8], [-0.9, 0.3, 0.3162], [0.2, -0.95, -0.24]] {
584 let lanes = NormalEncoding::encode_cube_face(minor, payload, n, 0.5);
585 let back = normal(NormalEncoding::CubeFace { minor, payload }, 0b1111).decode(lanes);
586 assert!(close(back.unwrap(), normalised(n).unwrap()), "{n:?}: {back:?}");
587 }
588 }
589
590 #[test]
593 fn the_roughness_is_read_from_a_gloss_a_roughness_or_an_exponent() {
594 let place = Place {
595 output: 0,
596 lanes: 0b0010,
597 };
598 let term = |encoding| Term { place, encoding };
599 let gloss = Roughness::Gloss(term(ScalarEncoding::square_of(0.95)));
602 assert!((gloss.decode([0.0, 0.64, 0.0, 0.0]) - 0.24).abs() < 1e-6);
603 let stored = Roughness::Roughness(term(ScalarEncoding::LINEAR));
604 assert_eq!(stored.decode([0.0, 0.3, 0.0, 0.0]), 0.3);
605 let exponent = Roughness::Exponent(term(ScalarEncoding::square_root_of(512.0)));
607 assert_eq!(exponent.decode([0.0; 4]), 1.0);
608 let stored_510 = (510.0f32 / 512.0).sqrt();
609 assert!((exponent.decode([0.0, stored_510, 0.0, 0.0]) - 0.25).abs() < 1e-4);
610 }
611
612 fn roles(outputs: &[(u8, OutputRole)]) -> Surface {
613 Surface::read(outputs.iter().map(|(lanes, role)| (*lanes, role)))
614 }
615
616 #[test]
620 fn the_reading_takes_the_first_output_of_each_role_and_refuses_lanes_that_do_not_fit() {
621 let linear = ScalarEncoding::LINEAR;
622 let s = roles(&[
623 (0b1111, OutputRole::Colour),
624 (0b0111, OutputRole::Albedo),
625 (0b0111, OutputRole::Albedo),
626 (
627 0b0001,
628 OutputRole::Occlusion {
629 of: Occluded::NaturalAmbient,
630 encoding: linear,
631 },
632 ),
633 (
634 0b0010,
635 OutputRole::Occlusion {
636 of: Occluded::AllLight,
637 encoding: linear,
638 },
639 ),
640 (0b0011, OutputRole::Metallic),
641 (0b0001, OutputRole::Metallic),
642 (0b1111, OutputRole::unknown("read by nothing")),
643 ]);
644 assert_eq!(
645 s.base,
646 Some(Base::Albedo(Place {
647 output: 1,
648 lanes: 0b0111
649 }))
650 );
651 assert_eq!(s.occlusion.map(|t| t.place.output), Some(4));
652 assert_eq!(s.ambient_occlusion, None);
653 assert_eq!(s.metallic.map(|t| t.place.output), Some(6));
654 let ambient = roles(&[(
655 0b0001,
656 OutputRole::Occlusion {
657 of: Occluded::NaturalAmbient,
658 encoding: linear,
659 },
660 )]);
661 assert_eq!(
662 (
663 ambient.occlusion,
664 ambient.ambient_occlusion.map(|t| t.place.output)
665 ),
666 (None, Some(0))
667 );
668 assert_eq!(s.refused, ["output 5: a metallic in 2 lanes (0b0011), not 1"]);
669 assert_eq!(
670 (s.normal, s.roughness, s.reflectance, s.emission),
671 (None, None, None, None)
672 );
673 let forward = roles(&[(0b1111, OutputRole::Colour)]);
675 assert_eq!(
676 forward.base,
677 Some(Base::Colour(Place {
678 output: 0,
679 lanes: 0b1111
680 }))
681 );
682 let scalar = roles(&[(0b0100, OutputRole::Emission { encoding: linear })]);
684 assert!(matches!(scalar.emission, Some(Emission::OfAlbedo(_))));
685 assert_eq!(
686 scalar
687 .emission
688 .unwrap()
689 .decode([0.0, 0.0, 2.0, 0.0], [0.5, 0.25, 1.0]),
690 [1.0, 0.5, 2.0]
691 );
692 let colour = roles(&[(0b0111, OutputRole::Emission { encoding: linear })]);
693 assert_eq!(
694 colour.emission.unwrap().decode([0.1, 0.2, 0.3, 0.0], [1.0; 3]),
695 [0.1, 0.2, 0.3]
696 );
697 }
698}