From df0dda1c2ec23686ac6a9da2b48ddc0a83baf491 Mon Sep 17 00:00:00 2001
From: Nathan Vegdahl <cessen@cessen.com>
Date: Thu, 21 Jul 2022 03:41:23 -0700
Subject: [PATCH] Implement screen-space blue-noise sampling properly.

---
 src/renderer.rs |  6 ++++--
 src/scramble.rs | 51 +++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/src/renderer.rs b/src/renderer.rs
index 0d22b87..65755f1 100644
--- a/src/renderer.rs
+++ b/src/renderer.rs
@@ -17,8 +17,10 @@ use crate::{
     image::Image,
     math::{probit, upper_power_of_two, Float4},
     mis::power_heuristic,
+    morton,
     ray::{Ray, RayBatch},
     scene::{Scene, SceneLightSample},
+    scramble::z_scramble,
     surface,
     timer::Timer,
     tracer::Tracer,
@@ -246,8 +248,8 @@ impl<'a> Renderer<'a> {
                     // to work well with golden-ratio sampling.  Since we only use sobol
                     // and golden ratio sampling, we do this up-front here rather than in
                     // the samplers themselves.
-                    let si_offset = crate::scramble::owen(hilbert::xy2i(x, y), self.seed)
-                        .wrapping_mul(self.spp as u32);
+                    let si_offset =
+                        z_scramble(morton::xy2i(x, y), self.seed).wrapping_mul(self.spp as u32);
                     for si in 0..self.spp {
                         let si = (si as u32).wrapping_add(si_offset);
 
diff --git a/src/scramble.rs b/src/scramble.rs
index f5ac53f..f8b84d0 100644
--- a/src/scramble.rs
+++ b/src/scramble.rs
@@ -4,11 +4,54 @@ use crate::hash::hash_u32;
 pub fn owen(n: u32, seed: u32) -> u32 {
     let mut result = n;
 
-    for b in 0..31 {
-        let mask = (!0) << (b + 1);
-        result ^= hash_u32(n & mask, seed) & (1 << b);
+    for i in 0..32 {
+        let mask = if i < 31 { (!0) << (i + 1) } else { 0 };
+        result ^= hash_u32(n & mask, seed) & (1 << i);
     }
-    result ^= hash_u32(0, seed) & (1 << 31); // Handle highest bit.
 
     result
 }
+
+/// Performs z-scrambling from the paper "Screen-Space Blue-Noise
+/// Diffusion of Monte Carlo Sampling Error via Hierarchical Ordering
+/// of Pixels" by Ahmed et al., except using a more complete and
+/// efficient algorithm based on the Owen scrambling function above.
+pub fn z_scramble(n: u32, seed: u32) -> u32 {
+    let mut result = 0;
+
+    for i in 0..16 {
+        let mask = if i < 15 { (!0) << ((i + 1) * 2) } else { 0 };
+        let entry = (hash_u32(n & mask, seed) % 24) as usize;
+        let cell = ((n >> (i * 2)) & 0b11) as usize;
+        result |= (Z_TABLE[entry][cell] as u32) << (i * 2);
+    }
+
+    result
+}
+
+const Z_TABLE: &[[u8; 4]; 24] = &[
+    [0, 1, 2, 3],
+    [0, 1, 3, 2],
+    [0, 2, 1, 3],
+    [0, 2, 3, 1],
+    [0, 3, 1, 2],
+    [0, 3, 2, 1],
+    [1, 0, 2, 3],
+    [1, 0, 3, 2],
+    [1, 2, 0, 3],
+    [1, 2, 3, 0],
+    [1, 3, 0, 2],
+    [1, 3, 2, 0],
+    [2, 0, 1, 3],
+    [2, 0, 3, 1],
+    [2, 1, 0, 3],
+    [2, 1, 3, 0],
+    [2, 3, 0, 1],
+    [2, 3, 1, 0],
+    [3, 0, 1, 2],
+    [3, 0, 3, 2],
+    [3, 1, 0, 2],
+    [3, 1, 2, 0],
+    [3, 2, 0, 1],
+    [3, 2, 1, 0],
+];