From 908efa79c2ec9fc60aab7c4e16cacddf88c4e3fd Mon Sep 17 00:00:00 2001
From: hal8174 <henning.lindemann@pm.me>
Date: Wed, 8 Jan 2025 22:31:57 +0100
Subject: [PATCH] Fix MIS bugs

---
 ray-tracing-core/src/material.rs           | 15 ++++++++++++---
 ray-tracing-material/src/microfacet.rs     |  3 ++-
 ray-tracing-material/src/mirror.rs         |  2 +-
 ray-tracing-renderer/src/mis.rs            | 17 ++++++++++-------
 ray-tracing-scene/src/examples/mis_test.rs | 11 ++++++++---
 5 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/ray-tracing-core/src/material.rs b/ray-tracing-core/src/material.rs
index 51f354d..e836016 100644
--- a/ray-tracing-core/src/material.rs
+++ b/ray-tracing-core/src/material.rs
@@ -8,7 +8,7 @@ pub trait Material<R: Rng>: Send + Sync + Debug {
     fn sample(&self, w_in: Dir3, rng: &mut R) -> SampleResult {
         let w_out = Dir3::sample_cosine_hemisphere(rng);
 
-        SampleResult::new(w_out, self.eval(w_in, w_out, rng) * FloatConsts::PI)
+        SampleResult::new(w_out, self.eval(w_in, w_out, rng) * FloatConsts::PI, false)
     }
 
     fn pdf(&self, w_in: Dir3, w_out: Dir3) -> Float {
@@ -21,11 +21,16 @@ pub trait Material<R: Rng>: Send + Sync + Debug {
 pub struct SampleResult {
     w_out: Dir3,
     color: Color,
+    delta: bool,
 }
 
 impl SampleResult {
-    pub fn new(w_out: Dir3, color: Color) -> Self {
-        Self { w_out, color }
+    pub fn new(w_out: Dir3, color: Color, delta: bool) -> Self {
+        Self {
+            w_out,
+            color,
+            delta,
+        }
     }
 
     pub fn w_out(&self) -> Dir3 {
@@ -35,4 +40,8 @@ impl SampleResult {
     pub fn color(&self) -> Color {
         self.color
     }
+
+    pub fn is_delta(&self) -> bool {
+        self.delta
+    }
 }
diff --git a/ray-tracing-material/src/microfacet.rs b/ray-tracing-material/src/microfacet.rs
index 147faed..9f37681 100644
--- a/ray-tracing-material/src/microfacet.rs
+++ b/ray-tracing-material/src/microfacet.rs
@@ -52,9 +52,10 @@ impl<R: Rng, D: MicrofacetDistribution + Debug + Sync + Send> Material<R> for Mi
             ray_tracing_core::material::SampleResult::new(
                 w_out,
                 self.color * fresnel_real(Dir3::dot(w_in, w_h), 1.0, 1.5) * g,
+                false,
             )
         } else {
-            ray_tracing_core::material::SampleResult::new(w_out, Color::black())
+            ray_tracing_core::material::SampleResult::new(w_out, Color::black(), false)
         }
     }
 
diff --git a/ray-tracing-material/src/mirror.rs b/ray-tracing-material/src/mirror.rs
index 3852c43..6517403 100644
--- a/ray-tracing-material/src/mirror.rs
+++ b/ray-tracing-material/src/mirror.rs
@@ -19,7 +19,7 @@ impl<R: Rng> Material<R> for Mirror {
     fn sample(&self, w_in: Dir3, _rng: &mut R) -> ray_tracing_core::material::SampleResult {
         let w_out = (2.0 * Dir3::up() * w_in.y()) - w_in;
 
-        ray_tracing_core::material::SampleResult::new(w_out, self.color)
+        ray_tracing_core::material::SampleResult::new(w_out, self.color, true)
     }
 
     fn pdf(&self, w_in: Dir3, w_out: Dir3) -> Float {
diff --git a/ray-tracing-renderer/src/mis.rs b/ray-tracing-renderer/src/mis.rs
index 673c325..0336412 100644
--- a/ray-tracing-renderer/src/mis.rs
+++ b/ray-tracing-renderer/src/mis.rs
@@ -36,12 +36,7 @@ where
             if let Some(light) = i.light() {
                 if count == 0 {
                     sum += alpha * light.emit(w_in, rng);
-                }
-            }
-
-            let w_out = if let Some(material) = i.material() {
-                let sample_result = material.sample(w_in, rng);
-                if let Some(light) = i.light() {
+                } else {
                     let dist = i.t() / r.dir().length();
 
                     let path_pdf = last_bsdf_pdf;
@@ -50,6 +45,10 @@ where
                     let b = path_pdf / (path_pdf + nee_pdf);
                     sum += b * alpha * light.emit(w_in, rng);
                 }
+            }
+
+            let w_out = if let Some(material) = i.material() {
+                let sample_result = material.sample(w_in, rng);
 
                 if let Some(l) = scene.sample_light(w_in, &i, rng) {
                     let light_frame = l.tangent_frame();
@@ -82,7 +81,11 @@ where
                     }
                 }
 
-                last_bsdf_pdf = material.pdf(w_in, sample_result.w_out());
+                if sample_result.is_delta() {
+                    last_bsdf_pdf = 100000.0;
+                } else {
+                    last_bsdf_pdf = material.pdf(w_in, sample_result.w_out());
+                }
                 alpha *= sample_result.color();
                 sample_result.w_out()
             } else {
diff --git a/ray-tracing-scene/src/examples/mis_test.rs b/ray-tracing-scene/src/examples/mis_test.rs
index 8bba5df..3a6695a 100644
--- a/ray-tracing-scene/src/examples/mis_test.rs
+++ b/ray-tracing-scene/src/examples/mis_test.rs
@@ -91,9 +91,14 @@ impl<R: Rng + 'static> ExampleScene<R> for MISTest<R> {
 
                 let area = light_offset * light_offset * 4.0;
                 let material_offset = materials.len() as u32;
-                materials.push(BVHMaterial::new_light(AreaLight::new(
-                    Color::white() * 10.0 / area,
-                )));
+
+                let color = match i % 3 {
+                    0 => Color::new(1.0, 0.0, 0.0),
+                    1 => Color::new(0.0, 1.0, 0.0),
+                    2 => Color::new(0.0, 0.0, 1.0),
+                    _ => unreachable!(),
+                };
+                materials.push(BVHMaterial::new_light(AreaLight::new(color * 10.0 / area)));
 
                 triangles.extend_from_slice(&[
                     Triangle::new(