Refactor the codebase and add aov support.

source/Algorithm/PathTracing.c

@@ -3,15 +3,157 @@
 #include "Lighting/LightEvaluation.h"
 #include "Algorithm/BSDF.h"
 
-// TODO: Split the diffuse and specular into different Monte Carlo, so we can decide the sample count for each one
-vec4s path_trace(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_t max_depth)
+static inline bool aov_wants_lighting(aov_flags_t flags)
+{
+    return (flags & (AOV_BEAUTY | AOV_DIRECT | AOV_INDIRECT)) != 0;
+}
+
+static inline bool aov_wants_surface(aov_flags_t flags)
+{
+    return (flags & (AOV_AlBEDO | AOV_NORMAL | AOV_DEPTH | AOV_POSITION)) != 0;
+}
+
+static inline vec4s vec4_add_rgb(vec4s base, vec3s rgb)
+{
+    return glms_vec4_add(base, glms_vec4(rgb, 0.0f));
+}
+
+static inline void aov_init_output(aov_output_t* out)
+{
+    *out = (aov_output_t){0};
+    out->beauty = (vec4s){0.0f, 0.0f, 0.0f, 1.0f};
+    out->direct = (vec4s){0.0f, 0.0f, 0.0f, 1.0f};
+    out->indirect = (vec4s){0.0f, 0.0f, 0.0f, 1.0f};
+}
+
+static inline shading_context_t make_shading_context(const scene_t* scene,
+                                                     vec3s wo,
+                                                     vec3s throughput,
+                                                     uint32_t sample_index,
+                                                     uint32_t bounce_depth,
+                                                     const hit_result_t* hit,
+                                                     float cone_width,
+                                                     float spread_angle)
+{
+    return (shading_context_t){
+        .camera_position = scene->camera.position,
+        .camera_direction = glms_vec3_normalize(glms_vec3_sub(hit->point, scene->camera.position)),
+
+        .position = hit->point,
+        .normal = hit->normal,
+        .tangent = hit->tangent,
+        .uv = hit->uv,
+        .wo = wo,
+        .throughput = throughput,
+
+        .sample_index = sample_index,
+        .bounce_depth = bounce_depth,
+
+        .bvh_tree = &scene->bvh_tree,
+        .triangles = &scene->triangles,
+        .lights = &scene->lights,
+        .textures = &scene->textures,
+
+        .triangle_id = hit->triangle_id,
+        .cone_width = cone_width,
+        .spread_angle = spread_angle,
+    };
+}
+
+static void trace_surface_aovs_only(const scene_t* scene,
+                                   ray_t ray,
+                                   uint32_t sample_index,
+                                   aov_flags_t aov_flags,
+                                   aov_output_t* out)
+{
+    hit_result_t closest_hit = ray_intersect_scene_closest(&ray, scene);
+    if (!closest_hit.hit)
+    {
+        return;
+    }
+
+    const material_t* hit_material = &scene->materials.buffer[scene->triangles.buffer[closest_hit.triangle_id].material_id];
+
+    float cone_width = ray.width + closest_hit.distance * ray.spread_angle;
+    shading_context_t shading_context = make_shading_context(scene,
+                                                            ray.direction,
+                                                            glms_vec3_one(),
+                                                            sample_index,
+                                                            0,
+                                                            &closest_hit,
+                                                            cone_width,
+                                                            ray.spread_angle);
+
+    render_material_aov(hit_material, &shading_context, out);
+
+    if (aov_flags & AOV_POSITION)
+    {
+        out->position = glms_vec4(closest_hit.point, 1.0f);
+    }
+    if (aov_flags & AOV_DEPTH)
+    {
+        out->depth = closest_hit.distance;
+    }
+}
+
+static inline void aov_accumulate_env(aov_output_t* out, aov_flags_t flags, vec3s env_contrib, uint16_t depth)
+{
+    if (flags & AOV_BEAUTY)
+    {
+        out->beauty = vec4_add_rgb(out->beauty, env_contrib);
+    }
+
+    // Environment visible from camera => direct; reached after bounces => indirect.
+    if (depth == 0)
+    {
+        if (flags & AOV_DIRECT)
+        {
+            out->direct = vec4_add_rgb(out->direct, env_contrib);
+        }
+    }
+    else
+    {
+        if (flags & AOV_INDIRECT)
+        {
+            out->indirect = vec4_add_rgb(out->indirect, env_contrib);
+        }
+    }
+}
+
+static inline void aov_accumulate_nee(aov_output_t* out, aov_flags_t flags, vec3s L, uint16_t depth)
+{
+    // Next-event estimation at camera-visible vertex => direct.
+    // NEE at later vertices is reached via bounces => indirect.
+    if (depth == 0)
+    {
+        if (flags & AOV_DIRECT)
+        {
+            out->direct = vec4_add_rgb(out->direct, L);
+        }
+    }
+    else
+    {
+        if (flags & AOV_INDIRECT)
+        {
+            out->indirect = vec4_add_rgb(out->indirect, L);
+        }
+    }
+
+    if (flags & AOV_BEAUTY)
+    {
+        out->beauty = vec4_add_rgb(out->beauty, L);
+    }
+}
+
+static void trace_lighting_aovs(const scene_t* scene,
+                               ray_t ray,
+                               uint32_t sample_index,
+                               uint16_t max_depth,
+                               aov_flags_t aov_flags,
+                               aov_output_t* out)
 {
-    vec4s accumulated_color = (vec4s){0.0f, 0.0f, 0.0f, 1.0f};
     vec3s throughput = glms_vec3_one();
-
     ray_t active_ray = ray;
-
-    // PDF of the direction that generated the current ray segment (used for MIS on env hits).
     float last_bsdf_pdf = 0.0f;
 
     uint16_t depth = 0;
@@ -21,7 +163,6 @@ vec4s path_trace(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_
 
         if (!closest_hit.hit)
         {
-            // Set bvh to null indicate that the ray is not hit anything
             light_shading_context_t light_context =
             {
                 .wo = active_ray.direction,
@@ -30,7 +171,6 @@ vec4s path_trace(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_
             };
             path_output sky_output = evaluate_bsdf_sky(&scene->lights, &light_context, throughput, sample_index);
 
-            // MIS for BSDF-sampled environment hit (for depth==0 camera ray, use weight 1).
             float w = 1.0f;
             if (depth > 0)
             {
@@ -43,52 +183,48 @@ vec4s path_trace(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_
             }
 
             vec3s env_contrib = glms_vec3_scale(sky_output.direct_lighting, w);
-            accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(env_contrib, 0.0f));
+            aov_accumulate_env(out, aov_flags, env_contrib, depth);
             break;
         }
 
         uint8_t material_id = scene->triangles.buffer[closest_hit.triangle_id].material_id;
         const material_t* hit_material = &scene->materials.buffer[material_id];
-        
-        // Calculate ray cone width at the hit point
+
         float current_cone_width = active_ray.width + closest_hit.distance * active_ray.spread_angle;
+        shading_context_t shading_context = make_shading_context(scene,
+                                                                active_ray.direction,
+                                                                throughput,
+                                                                sample_index,
+                                                                depth,
+                                                                &closest_hit,
+                                                                current_cone_width,
+                                                                active_ray.spread_angle);
 
-        shading_context_t shading_context =
+        // First-hit surface AOVs are still cheap; record them if requested.
+        if (depth == 0 && aov_wants_surface(aov_flags))
         {
-            .camera_position = scene->camera.position,
-            .camera_direction = glms_vec3_normalize(glms_vec3_sub(closest_hit.point, scene->camera.position)),
-
-            .position = closest_hit.point,
-            .normal = closest_hit.normal,
-            .tangent = closest_hit.tangent,
-            .uv = closest_hit.uv,
-            .wo = active_ray.direction,
-            .throughput = throughput,
-
-            .sample_index = sample_index,
-            .bounce_depth = depth,
-
-            .bvh_tree = &scene->bvh_tree,
-            .triangles = &scene->triangles,
-            .lights = &scene->lights,
-            .textures = &scene->textures,
-
-            .triangle_id = closest_hit.triangle_id,
-            .cone_width = current_cone_width,
-            .spread_angle = active_ray.spread_angle,
-        };
+            render_material_aov(hit_material, &shading_context, out);
+            if (aov_flags & AOV_POSITION)
+            {
+                out->position = glms_vec4(closest_hit.point, 1.0f);
+            }
+            if (aov_flags & AOV_DEPTH)
+            {
+                out->depth = closest_hit.distance;
+            }
+        }
 
         path_output material_output = render_material(hit_material, &shading_context);
         if (glms_vec3_isinf(material_output.direct_lighting) || glms_vec3_isnan(material_output.direct_lighting))
         {
-            goto end_path_trace;
+            break;
         }
 
-        accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(material_output.direct_lighting, 0.0f));
+        aov_accumulate_nee(out, aov_flags, material_output.direct_lighting, depth);
 
         if (material_output.pdf < FLT_EPSILON)
         {
-            goto end_path_trace;
+            break;
         }
 
         last_bsdf_pdf = material_output.pdf;
@@ -96,78 +232,52 @@ vec4s path_trace(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_
         throughput = glms_vec3_mul(throughput, material_output.bsdf);
         if (glms_vec3_isinf(throughput) || glms_vec3_isnan(throughput))
         {
-            goto end_path_trace;
+            break;
         }
 
-        // We do Russian roulette to decide whether to continue tracing or terminate the path
         if (depth > 1)
         {
             float q = fminf(glms_vec3_max(throughput), 0.95f);
             float rr = sobol_sample(sample_index, sobol_get_dimension(depth, PRNG_TERMINATE));
             if (rr > q)
             {
-                goto end_path_trace;
+                break;
             }
-            // Keep the energy of the path by scaling the throughput
+
             throughput = glms_vec3_scale(throughput, 1.0f / q);
         }
 
-        switch (material_output.state)
+        if (material_output.state != PS_SUCCESS)
         {
-            //case PATH_THROUGH:
-            //    active_ray = ray_create(BIAS_RAY_ORIGION(closest_hit.point, glms_vec3_negate(closest_hit.normal)), active_ray.direction);
-            //    continue;
-            case PS_SUCCESS:
-                vec3s origin = offset_ray_origin(closest_hit.point, closest_hit.normal, shading_context.wo);
-                active_ray = ray_create(origin, material_output.wi, current_cone_width, material_output.spread_angle);
-                depth++;
-                break;
-            default:
-                goto end_path_trace;
+            break;
         }
-    }
 
-end_path_trace:
-    return accumulated_color;
+        vec3s origin = offset_ray_origin(closest_hit.point, closest_hit.normal, shading_context.wo);
+        active_ray = ray_create(origin, material_output.wi, current_cone_width, material_output.spread_angle);
+        depth++;
+    }
 }
 
-// How to handle multi-bounced aov like indirect lighting?
-// Maybe we should move aov to path_trace and split accumulated_color into direct/indirect diffuse/specular before returning.
-void render_aov(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_t max_depth, aov_output_t* aov_output)
+void path_trace_aov(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_t max_depth, aov_flags_t aov_flags, aov_output_t* out)
 {
-    hit_result_t closest_hit = ray_intersect_scene_closest(&ray, scene);
-    if (!closest_hit.hit)
+    if (out == NULL)
     {
         return;
     }
 
-    const material_t* hit_material = &scene->materials.buffer[scene->triangles.buffer[closest_hit.triangle_id].material_id];
-    shading_context_t shading_context =
+    aov_init_output(out);
+
+    // Fast path: surface-only AOVs (single intersection + material aov evaluation).
+    if (!aov_wants_lighting(aov_flags))
     {
-        .camera_position = scene->camera.position,
-        .camera_direction = glms_vec3_normalize(glms_vec3_sub(closest_hit.point, scene->camera.position)),
+        if (!aov_wants_surface(aov_flags))
+        {
+            return;
+        }
 
-        .position = closest_hit.point,
-        .normal = closest_hit.normal,
-        .tangent = closest_hit.tangent,
-        .uv = closest_hit.uv,
-        .wo = ray.direction,
-        .throughput = 1.0f,
+        trace_surface_aovs_only(scene, ray, sample_index, aov_flags, out);
+        return;
+    }
 
-        .sample_index = sample_index,
-        .bounce_depth = 0,
-
-        .bvh_tree = &scene->bvh_tree,
-        .triangles = &scene->triangles,
-        .lights = &scene->lights,
-        .textures = &scene->textures,
-
-        .triangle_id = closest_hit.triangle_id,
-        .cone_width = ray.width + closest_hit.distance * ray.spread_angle,
-        .spread_angle = ray.spread_angle,
-    };
-
-    render_material_aov(hit_material, &shading_context, aov_output);
-    aov_output->position = glms_vec4(closest_hit.point, 1.0f);
-    aov_output->depth = closest_hit.distance;
+    trace_lighting_aovs(scene, ray, sample_index, max_depth, aov_flags, out);
 }