wrflib 0.0.3

// Copyright (c) 2021-present, Cruise LLC
//
// This source code is licensed under the Apache License, Version 2.0,
// found in the LICENSE-APACHE file in the root directory of this source tree.
// You may not use this file except in compliance with the License.

//! Managing [GPU shaders](https://en.wikipedia.org/wiki/Shader).

use crate::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use wrflib_shader_compiler::ty::Ty;
use wrflib_shader_compiler::{Decl, ShaderAst};

/// Contains all information necessary to build a shader.
/// Define a new shader.
///
/// Pass in a [`Geometry`] which gets used for instancing (e.g. a quad or a
/// cube).
///
/// The different [`CodeFragment`]s are appended together (but preserving their
/// filename/line/column information for error messages). They are split out
/// into `base_code_fragments` and `main_code_fragment` purely for
/// convenience. (We could instead have used a single [`slice`] but they are
/// annoying to get through concatenation..)
///
/// TODO(JP): Would be good to instead compile shaders beforehand, ie. during
/// compile time. Should look into that at some point.
pub struct Shader {
    /// The [`Geometry`] that we will draw with, if any. Can be overridden using [`DrawCallProps::gpu_geometry`].
    pub build_geom: Option<fn() -> Geometry>,
    /// A bunch of [`CodeFragment`]s that will get concatenated.
    pub code_to_concatenate: &'static [CodeFragment],
    /// The id of the shader (index into [`Cx::shaders`]), or [`Shader::UNCOMPILED_SHADER_ID`] if uninitialized.
    /// You should never read or modify this manually (see TODO below).
    ///
    /// TODO(JP): This shouldn't be public, but right now that's necessary for using [`Shader::DEFAULT`]. We might want to
    /// switch to a `define_shader` helper function once we can pass function pointers to `const` functions (see
    /// <https://github.com/rust-lang/rust/issues/63997> and <https://github.com/rust-lang/rust/issues/57563>).
    pub shader_id: AtomicUsize,
}

impl Shader {
    /// TODO(JP): We might want to switch to a `define_shader` helper function once we can pass function pointers
    /// to `const` functions (see <https://github.com/rust-lang/rust/issues/63997> and
    /// <https://github.com/rust-lang/rust/issues/57563>).
    ///
    /// We suppress `clippy::declare_interior_mutable_const` here since we don't actually want shader_id in this constant
    /// to be editable.
    #[allow(clippy::declare_interior_mutable_const)]
    pub const DEFAULT: Shader =
        Shader { build_geom: None, code_to_concatenate: &[], shader_id: AtomicUsize::new(Self::UNCOMPILED_SHADER_ID) };

    const UNCOMPILED_SHADER_ID: usize = usize::MAX;
}

/// Contains information of a [`CxShader`] of what instances, instances, textures
/// and so on it contains. That information can then be used to modify a [`Shader`
/// or [`DrawCall`].
#[derive(Debug, Default, Clone)]
pub struct CxShaderMapping {
    /// Contains information about the special "rect_pos" and "rect_size" fields.
    /// See [`RectInstanceProps`].
    pub rect_instance_props: RectInstanceProps,
    /// Special structure for user-level uniforms.
    pub user_uniform_props: UniformProps,
    /// Special structure for reading/editing instance properties.
    pub instance_props: InstanceProps,
    /// Special structure for reading/editing geometry properties.
    pub geometry_props: InstanceProps,
    /// Raw definition of all textures.
    pub textures: Vec<PropDef>,
    /// Raw definition of all geometries.
    pub geometries: Vec<PropDef>,
    /// Raw definition of all instances.
    pub instances: Vec<PropDef>,
    /// Raw definition of all user-level uniforms.
    pub user_uniforms: Vec<PropDef>,
    /// Raw definition of all framework-level uniforms that get set per [`DrawCall`].
    pub draw_uniforms: Vec<PropDef>,
    /// Raw definition of all framework-level uniforms that get set per [`View`].
    pub view_uniforms: Vec<PropDef>,
    /// Raw definition of all framework-level uniforms that get set per [`Pass`].
    pub pass_uniforms: Vec<PropDef>,
}

impl CxShaderMapping {
    pub fn from_shader_ast(shader_ast: ShaderAst) -> Self {
        let mut instances = Vec::new();
        let mut geometries = Vec::new();
        let mut user_uniforms = Vec::new();
        let mut draw_uniforms = Vec::new();
        let mut view_uniforms = Vec::new();
        let mut pass_uniforms = Vec::new();
        let mut textures = Vec::new();
        for decl in shader_ast.decls {
            match decl {
                Decl::Geometry(decl) => {
                    let prop_def = PropDef { name: decl.ident.to_string(), ty: decl.ty_expr.ty.borrow().clone().unwrap() };
                    geometries.push(prop_def);
                }
                Decl::Instance(decl) => {
                    let prop_def = PropDef { name: decl.ident.to_string(), ty: decl.ty_expr.ty.borrow().clone().unwrap() };
                    instances.push(prop_def);
                }
                Decl::Uniform(decl) => {
                    let prop_def = PropDef { name: decl.ident.to_string(), ty: decl.ty_expr.ty.borrow().clone().unwrap() };
                    match decl.block_ident {
                        Some(bi) if bi.with(|string| string == "draw") => {
                            draw_uniforms.push(prop_def);
                        }
                        Some(bi) if bi.with(|string| string == "view") => {
                            view_uniforms.push(prop_def);
                        }
                        Some(bi) if bi.with(|string| string == "pass") => {
                            pass_uniforms.push(prop_def);
                        }
                        None => {
                            user_uniforms.push(prop_def);
                        }
                        _ => (),
                    }
                }
                Decl::Texture(decl) => {
                    let prop_def = PropDef { name: decl.ident.to_string(), ty: decl.ty_expr.ty.borrow().clone().unwrap() };
                    textures.push(prop_def);
                }
                _ => (),
            }
        }

        CxShaderMapping {
            rect_instance_props: RectInstanceProps::construct(&instances),
            user_uniform_props: UniformProps::construct(&user_uniforms),
            instance_props: InstanceProps::construct(&instances),
            geometry_props: InstanceProps::construct(&geometries),
            textures,
            instances,
            geometries,
            pass_uniforms,
            view_uniforms,
            draw_uniforms,
            user_uniforms,
        }
    }
}

/// The raw definition of an input property to a [`Shader`].
#[derive(Debug, Clone, Hash, PartialEq)]
pub struct PropDef {
    pub name: String,
    pub ty: Ty,
}

/// Contains information about the special "rect_pos" and "rect_size" fields.
/// These fields describe the typical rectangles drawn in [`crate::QuadIns`]. It is
/// useful to have generalized access to them, so we can e.g. move a whole bunch
/// of rectangles at the same time, e.g. for alignment in [`CxLayoutBox`].
/// TODO(JP): We might want to consider instead doing bulk moves using [`DrawCall`-
/// or [`View`]-level uniforms.
#[derive(Debug, Default, Clone)]
pub struct RectInstanceProps {
    pub rect_pos: Option<usize>,
    pub rect_size: Option<usize>,
}
impl RectInstanceProps {
    pub fn construct(instances: &[PropDef]) -> RectInstanceProps {
        let mut rect_pos = None;
        let mut rect_size = None;
        let mut slot = 0;
        for inst in instances {
            match inst.name.as_ref() {
                "rect_pos" => rect_pos = Some(slot),
                "rect_size" => rect_size = Some(slot),
                _ => (),
            }
            slot += inst.ty.size(); //sg.get_type_slots(&inst.ty);
        }
        RectInstanceProps { rect_pos, rect_size }
    }
}

/// Represents an "instance" GPU input in a [`Shader`].
///
/// TODO(JP): Merge this into [`NamedProp`].
#[derive(Debug, Clone)]
pub struct InstanceProp {
    pub name: String,
    pub ty: Ty,
    pub offset: usize,
    pub slots: usize,
}

/// Represents all "instance" GPU inputs in a [`Shader`].
///
/// TODO(JP): Merge this into [`NamedProps`].
#[derive(Debug, Default, Clone)]
pub struct InstanceProps {
    pub props: Vec<InstanceProp>,
    pub total_slots: usize,
}

/// Represents a "uniform" GPU input in a [`Shader`].
///
/// TODO(JP): Merge this into [`NamedProp`].
#[derive(Debug, Clone)]
pub struct UniformProp {
    pub name: String,
    pub ty: Ty,
    pub offset: usize,
    pub slots: usize,
}

/// Represents all "uniform" GPU inputs in a [`Shader`].
///
/// TODO(JP): Merge this into [`NamedProps`].
#[derive(Debug, Default, Clone)]
pub struct UniformProps {
    pub props: Vec<UniformProp>,
    pub total_slots: usize,
}

/// A generic representation of any kind of [`Shader`] input (instance/uniform/geometry).
#[derive(Debug, Clone)]
pub struct NamedProp {
    pub name: String,
    pub ty: Ty,
    pub offset: usize,
    pub slots: usize,
}

/// A generic representation of a list of [`Shader`] inputs (instance/uniform/geometry).
#[derive(Debug, Default, Clone)]
pub struct NamedProps {
    pub props: Vec<NamedProp>,
    pub total_slots: usize,
}

impl NamedProps {
    pub fn construct(in_props: &[PropDef]) -> NamedProps {
        let mut offset = 0;
        let out_props = in_props
            .iter()
            .map(|prop| {
                let slots = prop.ty.size();
                let prop = NamedProp { ty: prop.ty.clone(), name: prop.name.clone(), offset, slots };
                offset += slots;
                prop
            })
            .collect();
        NamedProps { props: out_props, total_slots: offset }
    }
}

impl InstanceProps {
    pub fn construct(in_props: &[PropDef]) -> InstanceProps {
        let mut offset = 0;
        let out_props = in_props
            .iter()
            .map(|prop| {
                let slots = prop.ty.size();
                let prop = InstanceProp { ty: prop.ty.clone(), name: prop.name.clone(), offset, slots };
                offset += slots;
                prop
            })
            .collect();
        InstanceProps { props: out_props, total_slots: offset }
    }
}

impl UniformProps {
    pub fn construct(in_props: &[PropDef]) -> UniformProps {
        let mut offset = 0;
        let out_props = in_props
            .iter()
            .map(|prop| {
                let slots = prop.ty.size();
                let prop = UniformProp { ty: prop.ty.clone(), name: prop.name.clone(), offset, slots };
                offset += slots;
                prop
            })
            .collect();
        UniformProps { props: out_props, total_slots: offset }
    }

    pub fn find_zbias_uniform_prop(&self) -> Option<usize> {
        for prop in &self.props {
            if prop.name == "zbias" {
                return Some(prop.offset);
            }
        }
        None
    }
}

/// The actual shader information, which gets stored on [`Cx`]. Once compiled the
/// [`ShaderAst`] will be removed, and the [`CxPlatformShader`] (platform-specific
/// part of the compiled shader) gets set.
#[derive(Default)]
pub struct CxShader {
    pub name: String,
    pub gpu_geometry: Option<GpuGeometry>,
    pub(crate) platform: Option<CxPlatformShader>,
    pub mapping: CxShaderMapping,
    pub shader_ast: Option<ShaderAst>,
}

impl Cx {
    /// Get an individual [`Shader`] from a static [`Shader`].
    ///
    /// For more information on what [`LocationHash`] is used for here, see [`Shader`].
    pub(crate) fn get_shader_id(&mut self, shader: &'static Shader) -> usize {
        let shader_id = shader.shader_id.load(Ordering::Relaxed);
        if shader_id != Shader::UNCOMPILED_SHADER_ID {
            shader_id
        } else {
            // Use the last code fragment as the shader name.
            let main_code_fragment = shader.code_to_concatenate.last().expect("No code fragments found");
            let shader_name = format!("{}:{}:{}", main_code_fragment.filename, main_code_fragment.line, main_code_fragment.col);
            let shader_ast = self.shader_ast_generator.generate_shader_ast(&shader_name, shader.code_to_concatenate);

            let gpu_geometry = shader.build_geom.map(|build_geom| GpuGeometry::new(self, (build_geom)()));

            let shader_id = self.shaders.len();
            self.shaders.push(CxShader {
                name: shader_name,
                gpu_geometry,
                mapping: CxShaderMapping::from_shader_ast(shader_ast.clone()),
                platform: None,
                shader_ast: Some(shader_ast),
            });
            self.shader_recompile_ids.push(shader_id);

            shader.shader_id.store(shader_id, Ordering::Relaxed);

            shader_id
        }
    }
}