khora_lanes/render_lane/

forward_plus_lane.rs

// Copyright 2025 eraflo
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Forward+ (Tiled Forward) rendering lane implementation.
//!
//! This module implements a Forward+ rendering strategy that uses a compute shader
//! to perform per-tile light culling before the main render pass. This approach
//! significantly reduces the number of lights processed per fragment, making it
//! ideal for scenes with many lights (>20).
//!
//! # Architecture
//!
//! The Forward+ pipeline works in two stages:
//!
//! 1. **Light Culling (Compute Pass)**: The screen is divided into tiles (16x16 pixels).
//!    For each tile, the compute shader determines which lights intersect the tile's
//!    frustum and builds a list of affecting light indices.
//!
//! 2. **Rendering (Render Pass)**: Each fragment looks up its tile's light list and
//!    only evaluates lighting for those specific lights, rather than all lights in the scene.
//!
//! # Performance Characteristics
//!
//! - **O(tiles × lights)** for light culling (compute pass)
//! - **O(fragments × lights_per_tile)** for shading (render pass)
//! - **Suitable for**: Scenes with many lights (>20)
//! - **Break-even point**: ~20 lights (vs standard forward rendering)
//!
//! # SAA Compliance (Symbiotic Adaptive Architecture)
//!
//! This lane integrates with GORNA through:
//! - `estimate_cost()`: Provides accurate cost estimation including compute overhead
//! - Configurable tile size and max lights per tile
//! - Runtime-adjustable configuration via `ForwardPlusTileConfig`

use super::RenderWorld;
use crate::render_lane::{RenderLane, ShaderComplexity};

use khora_core::{
    asset::{AssetUUID, Material},
    renderer::{
        api::{
            buffer::BufferId,
            command::{
                ComputePassDescriptor, LoadOp, Operations, RenderPassColorAttachment,
                RenderPassDescriptor, StoreOp,
            },
            PrimitiveTopology,
        },
        traits::CommandEncoder,
        BindGroupId, ComputePipelineId, ForwardPlusTileConfig, GpuMesh, RenderContext,
        RenderPipelineId,
    },
};
use khora_data::assets::Assets;
use std::sync::RwLock;

// --- Cost Estimation Constants ---

/// Base cost per triangle rendered.
const TRIANGLE_COST: f32 = 0.001;

/// Cost per draw call issued.
const DRAW_CALL_COST: f32 = 0.1;

/// Fixed overhead for the compute pass (tile frustum + dispatch).
const COMPUTE_PASS_OVERHEAD: f32 = 0.5;

/// Cost factor per tile in the light culling pass.
const PER_TILE_COST: f32 = 0.0001;

/// Cost factor per light-tile intersection test.
const LIGHT_TILE_TEST_COST: f32 = 0.00001;

// --- ForwardPlusLane ---

/// GPU resource handles for the Forward+ compute pass.
///
/// These are created during lane initialization and used each frame
/// for light culling and rendering.
#[derive(Debug, Clone, Default)]
pub struct ForwardPlusGpuResources {
    /// Light buffer containing all GpuLight instances.
    pub light_buffer: Option<BufferId>,
    /// Buffer containing per-tile light index lists.
    pub light_index_buffer: Option<BufferId>,
    /// Buffer containing (offset, count) pairs per tile.
    pub light_grid_buffer: Option<BufferId>,
    /// Uniform buffer for culling parameters.
    pub culling_uniforms_buffer: Option<BufferId>,
    /// Bind group for the culling compute shader.
    pub culling_bind_group: Option<BindGroupId>,
    /// Bind group for the forward pass (light data).
    pub forward_bind_group: Option<BindGroupId>,
    /// Compute pipeline for light culling.
    pub culling_pipeline: Option<ComputePipelineId>,
    /// Render pipeline for the Forward+ pass.
    pub render_pipeline: Option<RenderPipelineId>,
}

impl ForwardPlusGpuResources {
    /// Returns true if all required resources are initialized.
    pub fn is_initialized(&self) -> bool {
        self.light_buffer.is_some()
            && self.light_index_buffer.is_some()
            && self.light_grid_buffer.is_some()
            && self.culling_uniforms_buffer.is_some()
            && self.culling_bind_group.is_some()
            && self.culling_pipeline.is_some()
    }
}

/// A rendering lane that implements Forward+ (Tiled Forward) rendering.
///
/// Forward+ divides the screen into tiles and uses a compute shader to determine
/// which lights affect each tile before the main render pass. This significantly
/// reduces per-fragment lighting cost for scenes with many lights.
///
/// # Configuration
///
/// The lane is configured via `ForwardPlusTileConfig`, which controls:
/// - **Tile size**: 16x16 or 32x32 pixels (trade-off between culling granularity and overhead)
/// - **Max lights per tile**: Memory budget for per-tile light lists
/// - **Depth pre-pass**: Optional optimization for depth-bounded light culling
#[derive(Debug, Clone)]
pub struct ForwardPlusLane {
    /// Tile configuration for light culling.
    pub tile_config: ForwardPlusTileConfig,

    /// Shader complexity for cost estimation.
    pub shader_complexity: ShaderComplexity,

    /// Current screen dimensions (for tile count calculation).
    screen_size: (u32, u32),

    /// GPU resources for compute and render passes.
    pub gpu_resources: ForwardPlusGpuResources,
}

impl Default for ForwardPlusLane {
    fn default() -> Self {
        Self {
            tile_config: ForwardPlusTileConfig::default(),
            shader_complexity: ShaderComplexity::SimpleLit,
            screen_size: (1920, 1080),
            gpu_resources: ForwardPlusGpuResources::default(),
        }
    }
}

impl ForwardPlusLane {
    /// Creates a new `ForwardPlusLane` with default settings.
    ///
    /// Default configuration:
    /// - Tile size: 16x16 pixels
    /// - Max lights per tile: 128
    /// - No depth pre-pass
    pub fn new() -> Self {
        Self::default()
    }

    /// Creates a new `ForwardPlusLane` with the specified configuration.
    ///
    /// # Arguments
    ///
    /// * `config` - The tile configuration for light culling
    pub fn with_config(config: ForwardPlusTileConfig) -> Self {
        Self {
            tile_config: config,
            ..Default::default()
        }
    }

    /// Creates a new `ForwardPlusLane` with the specified shader complexity.
    pub fn with_complexity(complexity: ShaderComplexity) -> Self {
        Self {
            shader_complexity: complexity,
            ..Default::default()
        }
    }

    /// Updates the screen size used for tile calculations.
    ///
    /// This should be called when the window is resized to recalculate
    /// tile counts and buffer sizes.
    pub fn set_screen_size(&mut self, width: u32, height: u32) {
        self.screen_size = (width, height);
    }

    /// Calculates the number of tiles in each dimension.
    pub fn tile_count(&self) -> (u32, u32) {
        self.tile_config
            .tile_dimensions(self.screen_size.0, self.screen_size.1)
    }

    /// Calculates the total number of tiles on screen.
    pub fn total_tiles(&self) -> u32 {
        let (tiles_x, tiles_y) = self.tile_count();
        tiles_x * tiles_y
    }

    /// Returns the effective number of lights in the scene.
    ///
    /// This counts all light types (directional, point, spot) that will be
    /// processed by the light culling pass.
    pub fn effective_light_count(&self, render_world: &RenderWorld) -> usize {
        render_world.directional_light_count()
            + render_world.point_light_count()
            + render_world.spot_light_count()
    }

    /// Estimates the cost of the compute pass (light culling).
    fn compute_pass_cost(&self, render_world: &RenderWorld) -> f32 {
        let total_tiles = self.total_tiles() as f32;
        let light_count = self.effective_light_count(render_world) as f32;

        // Compute pass cost = overhead + per-tile cost + light-tile tests
        COMPUTE_PASS_OVERHEAD
            + (total_tiles * PER_TILE_COST)
            + (total_tiles * light_count * LIGHT_TILE_TEST_COST)
    }

    /// Calculates the per-fragment light cost factor.
    ///
    /// For Forward+, this uses sqrt(total_lights) instead of linear scaling
    /// because lights are culled per-tile, so each fragment only processes
    /// a subset of lights.
    fn fragment_light_factor(&self, render_world: &RenderWorld) -> f32 {
        let total_lights = self.effective_light_count(render_world) as f32;

        if total_lights == 0.0 {
            return 1.0;
        }

        // Sublinear scaling: sqrt(lights) because of tile culling
        // Clamped to max_lights_per_tile
        let effective_lights = total_lights
            .sqrt()
            .min(self.tile_config.max_lights_per_tile as f32);

        1.0 + (effective_lights * 0.02)
    }
}

impl RenderLane for ForwardPlusLane {
    fn strategy_name(&self) -> &'static str {
        "ForwardPlus"
    }

    fn get_pipeline_for_material(
        &self,
        material_uuid: Option<AssetUUID>,
        materials: &Assets<Box<dyn Material>>,
    ) -> RenderPipelineId {
        // Verify material exists (same logic as LitForwardLane)
        if let Some(uuid) = material_uuid {
            if materials.get(&uuid).is_none() {
                let _ = uuid;
            }
        }

        // Forward+ uses pipeline ID 2 to differentiate from standard forward (ID 1)
        // This will be the pipeline created with the forward_plus.wgsl shader
        RenderPipelineId(2)
    }

    fn render(
        &self,
        render_world: &RenderWorld,
        encoder: &mut dyn CommandEncoder,
        render_ctx: &RenderContext,
        gpu_meshes: &RwLock<Assets<GpuMesh>>,
        materials: &RwLock<Assets<Box<dyn Material>>>,
    ) {
        // Acquire read locks
        let gpu_mesh_assets = gpu_meshes.read().unwrap();
        let material_assets = materials.read().unwrap();

        // Pre-compute pipelines
        let pipelines: Vec<RenderPipelineId> = render_world
            .meshes
            .iter()
            .map(|mesh| self.get_pipeline_for_material(mesh.material_uuid, &material_assets))
            .collect();

        // --- STAGE 1: Light Culling (Compute Pass) ---
        let (tiles_x, tiles_y) = self.tile_count();

        // Only dispatch compute pass if GPU resources are initialized
        if self.gpu_resources.is_initialized() {
            let compute_pass_desc = ComputePassDescriptor {
                label: Some("Forward+ Light Culling"),
                timestamp_writes: None,
            };

            {
                let mut compute_pass = encoder.begin_compute_pass(&compute_pass_desc);

                // Set compute pipeline
                if let Some(ref pipeline) = self.gpu_resources.culling_pipeline {
                    compute_pass.set_pipeline(pipeline);
                }

                // Bind resources: uniforms, lights, light_index_list, light_grid
                if let Some(ref bind_group) = self.gpu_resources.culling_bind_group {
                    compute_pass.set_bind_group(0, bind_group);
                }

                // Dispatch one workgroup per tile
                compute_pass.dispatch_workgroups(tiles_x, tiles_y, 1);
            }
            // compute_pass dropped here, ending the pass
        }

        // --- STAGE 2: Rendering (Render Pass) ---

        let color_attachment = RenderPassColorAttachment {
            view: render_ctx.color_target,
            resolve_target: None,
            ops: Operations {
                load: LoadOp::Clear(render_ctx.clear_color),
                store: StoreOp::Store,
            },
        };

        let render_pass_desc = RenderPassDescriptor {
            label: Some("Forward+ Render Pass"),
            color_attachments: &[color_attachment],
        };

        let mut render_pass = encoder.begin_render_pass(&render_pass_desc);

        // Track current pipeline to avoid redundant state changes
        let mut current_pipeline: Option<RenderPipelineId> = None;

        // Bind Forward+ lighting data (lights, indices, grid, tile info)
        // This bind group is created from the compute pass outputs
        if let Some(ref forward_bg) = self.gpu_resources.forward_bind_group {
            render_pass.set_bind_group(3, forward_bg);
        }

        // Iterate over all extracted meshes
        for (i, extracted_mesh) in render_world.meshes.iter().enumerate() {
            if let Some(gpu_mesh) = gpu_mesh_assets.get(&extracted_mesh.gpu_mesh_uuid) {
                let pipeline = &pipelines[i];

                // Only bind pipeline if changed
                if current_pipeline != Some(*pipeline) {
                    render_pass.set_pipeline(pipeline);
                    current_pipeline = Some(*pipeline);
                }

                // Bind buffers
                render_pass.set_vertex_buffer(0, &gpu_mesh.vertex_buffer, 0);
                render_pass.set_index_buffer(&gpu_mesh.index_buffer, 0, gpu_mesh.index_format);

                // Issue draw call
                render_pass.draw_indexed(0..gpu_mesh.index_count, 0, 0..1);
            }
        }
    }

    fn estimate_cost(
        &self,
        render_world: &RenderWorld,
        gpu_meshes: &RwLock<Assets<GpuMesh>>,
    ) -> f32 {
        let gpu_mesh_assets = gpu_meshes.read().unwrap();

        let mut total_triangles = 0u32;
        let mut draw_call_count = 0u32;

        for extracted_mesh in &render_world.meshes {
            if let Some(gpu_mesh) = gpu_mesh_assets.get(&extracted_mesh.gpu_mesh_uuid) {
                let triangle_count = match gpu_mesh.primitive_topology {
                    PrimitiveTopology::TriangleList => gpu_mesh.index_count / 3,
                    PrimitiveTopology::TriangleStrip => {
                        if gpu_mesh.index_count >= 3 {
                            gpu_mesh.index_count - 2
                        } else {
                            0
                        }
                    }
                    PrimitiveTopology::LineList
                    | PrimitiveTopology::LineStrip
                    | PrimitiveTopology::PointList => 0,
                };

                total_triangles += triangle_count;
                draw_call_count += 1;
            }
        }

        // Base geometry cost
        let geometry_cost =
            (total_triangles as f32 * TRIANGLE_COST) + (draw_call_count as f32 * DRAW_CALL_COST);

        // Shader complexity multiplier
        let shader_multiplier = self.shader_complexity.cost_multiplier();

        // Compute pass overhead
        let compute_cost = self.compute_pass_cost(render_world);

        // Per-fragment light factor (sublinear for Forward+)
        let light_factor = self.fragment_light_factor(render_world);

        // Total cost
        compute_cost + (geometry_cost * shader_multiplier * light_factor)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use khora_core::renderer::TileSize;

    #[test]
    fn test_forward_plus_lane_creation() {
        let lane = ForwardPlusLane::new();
        assert_eq!(lane.tile_config.tile_size, TileSize::X16);
        assert_eq!(lane.tile_config.max_lights_per_tile, 128);
        assert_eq!(lane.shader_complexity, ShaderComplexity::SimpleLit);
    }

    #[test]
    fn test_forward_plus_lane_with_config() {
        let config = ForwardPlusTileConfig {
            tile_size: TileSize::X32,
            max_lights_per_tile: 256,
            use_depth_prepass: true,
        };
        let lane = ForwardPlusLane::with_config(config);

        assert_eq!(lane.tile_config.tile_size, TileSize::X32);
        assert_eq!(lane.tile_config.max_lights_per_tile, 256);
        assert!(lane.tile_config.use_depth_prepass);
    }

    #[test]
    fn test_tile_count_calculation() {
        let mut lane = ForwardPlusLane::new();
        lane.set_screen_size(1920, 1080);

        let (tiles_x, tiles_y) = lane.tile_count();
        assert_eq!(tiles_x, 120); // 1920 / 16
        assert_eq!(tiles_y, 68); // ceil(1080 / 16)
    }

    #[test]
    fn test_strategy_name() {
        let lane = ForwardPlusLane::new();
        assert_eq!(lane.strategy_name(), "ForwardPlus");
    }

    #[test]
    fn test_pipeline_id() {
        let lane = ForwardPlusLane::new();
        let materials = Assets::<Box<dyn Material>>::new();
        let pipeline = lane.get_pipeline_for_material(None, &materials);
        assert_eq!(pipeline, RenderPipelineId(2));
    }
}