wsg/lib/pipeline_cache.rs

//! # PipelineCache Module
//!
//! Defines `PipelineCache`, the library's shader compilation cache. It owns WGSL shader loading and RenderPipeline
//! creation, storing compiled pipelines in a HashMap keyed by shader_id + texture format to avoid duplicate GPU work.
//! Materials request pipelines through this cache; if a pipeline for the given key exists, it is returned directly
//! via Arc cloning. Otherwise the cache compiles one on-the-fly, caches it, then returns it.
//!
//! ## Interaction with Other Modules
//! - **Material** calls `get_or_create()` during its own construction to obtain a shared RenderPipeline.
//! - **conf::BASIC_SHADER** provides fallback WGSL source when an external file is not found.
//! - **vertex::Vertex** defines the CPU-side layout that `build_pipeline` uses as the vertex buffer contract.
//!
//! ## Technical Points
//! - Pipelines are stored behind `Arc` so multiple Materials share the same compiled object without copying.
//! - wgpu 30 requires `compilation_options` in VertexState/FragmentState and `depth_slice` in color attachments.

use crate::conf::BASIC_SHADER;
use crate::vertex::Vertex;
use std::collections::HashMap;
use std::sync::Arc;

/// Shader pipeline cache: maps (shader_id, format) keys to compiled RenderPipelines.
/// Ensures each unique shader+format combination is compiled at most once; subsequent requests return cached instances.
pub struct PipelineCache {
    /// Cached pipelines keyed by their shader identifier string. Multiple Materials sharing the same ID share one Arc-wrapped pipeline.
    pipelines: HashMap<String, Arc<wgpu::RenderPipeline>>,
    /// Maps shader IDs to file paths on disk for WGSL loading in `load_shader()`.
    shader_paths: HashMap<String, String>,
}

impl PipelineCache {
    /// Creates an empty pipeline cache with no pre-loaded shaders or pipelines.
    /// Called at application startup before any Material creation. Shader paths must be registered via register_shader() first.
    pub fn new() -> Self {
        Self {
            pipelines: HashMap::new(),
            // Maps shader IDs to file paths on disk for WGSL loading in load_shader().
            // When a path exists, it reads from it; otherwise falls back to BASIC_SHADER constant.
            shader_paths: HashMap::new(),
        }
    }
    /// Registers an external WGSL shader file path associated with a given ID.
    /// Inputs: id (unique key for this shader), path (filesystem path to .wgsl file).
    /// Returns Ok(id) on success or Err if the ID is already registered. Called during scene setup to register custom shaders.
    pub fn register_shader(&mut self, id: &str, path: &str) -> Result<String, String> {
        if self.shader_paths.contains_key(id) {
            return Err(format!("ID '{}' already exists.", id));
        }
        self.shader_paths.insert(id.to_string(), path.to_string());
        Ok(id.to_string())
    }

    /// Unregisters a shader by its ID, removing both the path reference and any cached compiled pipeline.
    /// Inputs: id (the shader identifier to remove).
    /// Returns Ok(id) on success or Err if the ID does not exist. Called when a shader should be freed from GPU memory.
    pub fn unregister_shader(&mut self, id: &str) -> Result<String, String> {
        if self.shader_paths.remove(id).is_none() {
            return Err(format!("ID '{}' does not exist.", id));
        }
        // Remove cached pipeline so GPU memory is freed (wgpu drops it automatically)

        Ok(id.to_string())
    }

    /// Retrieves a cached RenderPipeline by shader_id, or creates one on-demand if not present.
    /// Inputs: device (GPU command source), format (surface texture format for fragment output),
    ///         shader_id (unique key into the cache).
    /// Returns an Arc-wrapped RenderPipeline ready for rendering. Called by Material::new().
    pub fn get_or_create(
        &mut self,
        device: &wgpu::Device,
        format: wgpu::TextureFormat,
        shader_id: &str,
    ) -> Arc<wgpu::RenderPipeline> {
        // Step 1: Return cached pipeline if it already exists for this shader_id
        if let Some(pipeline) = self.pipelines.get(shader_id) {
            return pipeline.clone();
        }

        // Step 2: Compile a new pipeline — loads shader and builds the GPU render pipeline
        let path = self
            .shader_paths
            .get(shader_id)
            .map(|s| s.as_str())
            .unwrap_or(shader_id);
        let shader = self.load_shader(device, path);
        let pipeline = Self::build_pipeline(device, format, &shader);

        // Step 3: Cache the new pipeline behind Arc and return it
        let pipeline_arc = Arc::new(pipeline);
        self.pipelines
            .insert(shader_id.to_string(), pipeline_arc.clone());
        pipeline_arc
    }

    /// Loads a WGSL shader module: reads from disk first, falls back to the embedded BASIC_SHADER constant.
    /// Called internally by `get_or_create()` when compiling a new pipeline.
    fn load_shader(&self, device: &wgpu::Device, path: &str) -> wgpu::ShaderModule {
        let source = std::fs::read_to_string(path).unwrap_or_else(|_| {
            println!("Shader not found: {}, falling back to default", path);
            BASIC_SHADER.to_string()
        });

        device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some(path),
            source: wgpu::ShaderSource::Wgsl(source.into()),
        })
    }

    /// Builds a RenderPipeline from a shader module, device, and surface texture format.
    /// Inputs: device (GPU command source), format (output texture format), shader (compiled WGSL module).
    /// Returns a fully configured RenderPipeline ready for draw calls. Called internally by `get_or_create()`.
    fn build_pipeline(
        device: &wgpu::Device,
        format: wgpu::TextureFormat,
        shader: &wgpu::ShaderModule,
    ) -> wgpu::RenderPipeline {
        // Define vertex attribute layout — the contract between CPU vertex data and GPU shader inputs.
        // Must match Vertex struct field offsets exactly.
        let vertex_buffer_layout = wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &[
                wgpu::VertexAttribute {
                    offset: 0,
                    shader_location: 0,
                    format: wgpu::VertexFormat::Float32x3,
                }, // position
                wgpu::VertexAttribute {
                    offset: 12,
                    shader_location: 1,
                    format: wgpu::VertexFormat::Float32x3,
                }, // normal
                wgpu::VertexAttribute {
                    offset: 24,
                    shader_location: 2,
                    format: wgpu::VertexFormat::Float32x2,
                }, // uv
                wgpu::VertexAttribute {
                    offset: 32,
                    shader_location: 3,
                    format: wgpu::VertexFormat::Float32x4,
                }, // color
            ],
        };

        // Pipeline layout — defines bind group bindings (empty here; no uniform buffers used).
        // wgpu 30: `immediate_size` replaces `push_constant_ranges`.
        let render_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("render_pipeline_layout"),
                bind_group_layouts: &[],
                immediate_size: 0, // no var<immediate> used
            });

        // Create the full RenderPipeline — vertex state + fragment state + primitive configuration.
        device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Render Pipeline"),
            layout: Some(&render_pipeline_layout),
            vertex: wgpu::VertexState {
                module: shader,
                // entry_point is now Option<&str> — Some to specify explicitly, None for auto-detection.
                entry_point: Some("vs_main"),
                compilation_options: Default::default(), // required field in wgpu 30
                buffers: &[Some(vertex_buffer_layout)],
            },
            fragment: Some(wgpu::FragmentState {
                module: shader,
                entry_point: Some("fs_main"),
                compilation_options: Default::default(), // required field in wgpu 30
                // targets is now &[Option<ColorTargetState>] — each wrapped in Some.
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
            }),
            primitive: wgpu::PrimitiveState::default(),
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            // multiview → replaced by multiview_mask (NonZeroU32) and cache fields in wgpu 30.
            multiview_mask: None,
            cache: None,
        })
    }

    /// Retrieves a cached RenderPipeline by shader_id without creating one.
    /// Inputs: shader_id (unique key into the cache).
    /// Returns Some(Arc<RenderPipeline>) if found, None otherwise. Called by renderer code for pipeline inspection.
    pub fn get(&self, shader_id: &str) -> Option<&Arc<wgpu::RenderPipeline>> {
        self.pipelines.get(shader_id)
    }
}