2 files changed, 33 insertions, 5 deletions

diff --git a/hitlist.h b/hitlist.h
index d498b26..02c942a 100644
--- a/hitlist.h
+++ b/hitlist.h
@@ -44,7 +44,7 @@ public:
 
     // Given a ray, compute the color.
     pixel<T> color(ray3d r, random_uv_gen_3d &ruvg, uint_fast32_t max_recursion_depth = 64) const {
-        double decay = 1;
+        assert(r.decay() == 1.0);
         while (max_recursion_depth-- > 0) {
             // Detect hits
             bool hit = false;
@@ -90,8 +90,8 @@ public:
 #ifdef DIFFUSE_HEMI
                 vec3d diffuse_target = hit_point + ruvg.hemisphere(nv);
 #endif
-                decay *= 0.5; // lose 50% light when diffused
-                r = ray3d{hit_point, diffuse_target - hit_point}; // the new diffused ray we trace on
+                r.decay(0.5); // lose 50% light when diffused
+                r.source(hit_point); r.direction((diffuse_target - hit_point).unit_vec()); // the new diffused ray we trace on
                 continue;
 #endif
             }
@@ -105,7 +105,7 @@ public:
                     u
             );
 #ifdef T_DIFFUSE
-            return decay * c;
+            return r.hit(c);
 #else
             return c;
 #endif
diff --git a/ray.h b/ray.h
index db28d0f..4d5401c 100644
--- a/ray.h
+++ b/ray.h
@@ -12,6 +12,7 @@ template<typename T>
 class ray3 {
     vec3<T> source_;
     vec3<T> direction_; // unit vector
+    double decay_; // How much power remaining
 
 public:
     ~ray3() = default;
@@ -23,27 +24,54 @@ public:
     ray3<T> &operator=(const ray3<T> &other) {
         source_ = other.source_;
         direction_ = other.direction_;
+        decay_ = other.decay_;
         return *this;
     }
 
-    ray3(const vec3<T> &source, const vec3<T> &direction) : source_(source), direction_(direction.unit_vec()) {}
+    ray3(const vec3<T> &source, const vec3<T> &direction) :
+            source_(source), direction_(direction.unit_vec()), decay_(1.0) {}
 
     // Get the source point from where the ray emits.
     vec3<T> source() const {
         return source_;
     }
 
+    void source(const vec3<T> &s) {
+        source_ = s;
+    }
+
     // Get the unit vector along the ray's direction.
     vec3<T> direction() const {
         return direction_;
     }
 
+    // d should be a unit vector
+    void direction(const vec3<T> &d) {
+        assert(fabs(d.mod2() - 1.0) < 1e-6);
+        direction_ = d;
+    }
+
     // Compute the point this ray reaches at the time `t`.
     template<typename U>
     vec3<T> at(U t) const {
         return source_ + direction_ * t;
     }
 
+    // Get the final color we got at the ray's origin, by tracing back along with the ray's path.
+    // Currently, here is only a simple power decay accumulation.
+    template<typename U>
+    pixel<U> hit(const pixel<U> &color) const {
+        return decay_ * color;
+    }
+
+    void decay(double a) {
+        decay_ *= a;
+    }
+
+    double decay() const {
+        return decay_;
+    }
+
     // Get a ray starts from zero and directs to undefined direction.
     static ray3<T> null() {
         return ray3<T>{vec3<T>::zero(), vec3<T>::zero()};