path: root/engine/src/rendering/shader.rs

use std::any::type_name;
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::fmt::{Debug, Display, Write};
use std::path::Path;
use std::str::Utf8Error;
use std::sync::Arc;

use bitflags::{bitflags, bitflags_match};
use ecs::phase::INIT as INIT_PHASE;
use shader_slang::{
    Blob as SlangBlob,
    ComponentType as SlangComponentType,
    DebugInfoLevel as SlangDebugInfoLevel,
    EntryPoint as SlangEntryPoint,
    GlobalSession as SlangGlobalSession,
    Module as SlangModule,
    Session as SlangSession,
};

use crate::asset::{
    Asset,
    Assets,
    Event as AssetEvent,
    Events as AssetEvents,
    Handle as AssetHandle,
    Id as AssetId,
    Submitter as AssetSubmitter,
    HANDLE_ASSETS_PHASE,
};
use crate::builder;
use crate::ecs::pair::ChildOf;
use crate::ecs::phase::Phase;
use crate::ecs::sole::Single;
use crate::ecs::{declare_entity, pair, Sole};
use crate::mesh::vertex_buffer::VertexLabel;

pub mod cursor;
pub mod reflection;

pub const STD_VERT_IN_SEM_NAME_POSITION: &str = "STD_POSITION";
pub const STD_VERT_IN_SEM_NAME_NORMAL: &str = "STD_NORMAL";
pub const STD_VERT_IN_SEM_NAME_UV_FROM_TOP_LEFT: &str = "STD_UV_FROM_TOP_LEFT";
pub const STD_VERT_IN_SEM_NAME_COLOR: &str = "STD_COLOR";

/// Shader module.
#[derive(Debug)]
pub struct ModuleSource
{
    pub name: Cow<'static, str>,
    pub file_path: Cow<'static, Path>,
    pub source: Cow<'static, str>,
    pub link_entrypoints: EntrypointFlags,
}

impl Asset for ModuleSource
{
    fn visit_dependencies<'this>(
        &'this self,
        _func: &mut dyn FnMut(crate::asset::Dependency<'this>),
    )
    {
    }
}

bitflags! {
    #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
    pub struct EntrypointFlags: usize
    {
        const FRAGMENT = 1 << 0;
        const VERTEX = 1 << 1;
    }
}

#[derive(Clone)]
pub struct Module
{
    inner: SlangModule,
}

impl Module
{
    #[must_use]
    pub fn entry_points(&self) -> impl ExactSizeIterator<Item = EntryPoint> + use<'_>
    {
        self.inner
            .entry_points()
            .map(|entry_point| EntryPoint { inner: entry_point })
    }

    #[must_use]
    pub fn get_entry_point(&self, entry_point: &str) -> Option<EntryPoint>
    {
        let entry_point = self.inner.find_entry_point_by_name(entry_point)?;

        Some(EntryPoint { inner: entry_point })
    }

    #[must_use]
    pub fn file_path(&self) -> &str
    {
        self.inner.file_path()
    }
}

pub struct EntryPoint
{
    inner: SlangEntryPoint,
}

impl EntryPoint
{
    #[must_use]
    pub fn function_reflection(&self) -> self::reflection::Function<'_>
    {
        self::reflection::Function {
            inner: self.inner.function_reflection(),
        }
    }
}

#[derive(Clone)]
pub struct Program
{
    inner: SlangComponentType,
    metadata: ProgramMetadata,
}

impl Program
{
    pub fn into_linked(self) -> Result<Program, Error>
    {
        let linked_program = self.inner.link()?;

        Ok(Program {
            inner: linked_program,
            metadata: self.metadata,
        })
    }

    #[must_use]
    pub fn metadata(&self) -> &ProgramMetadata
    {
        &self.metadata
    }

    pub fn get_entry_point_code(&self, entry_point_index: u32) -> Result<Blob, Error>
    {
        let blob = self.inner.entry_point_code(entry_point_index.into(), 0)?;

        Ok(Blob { inner: blob })
    }

    pub fn reflection(&self, target: u32)
        -> Result<self::reflection::Program<'_>, Error>
    {
        let reflection = self.inner.layout(i64::from(target))?;

        Ok(self::reflection::Program { inner: reflection })
    }
}

impl Debug for Program
{
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result
    {
        formatter
            .debug_struct(type_name::<Self>())
            .finish_non_exhaustive()
    }
}

pub struct Blob
{
    inner: SlangBlob,
}

impl Blob
{
    #[must_use]
    pub fn as_bytes(&self) -> &[u8]
    {
        self.inner.as_slice()
    }

    pub fn as_str(&self) -> Result<&str, Utf8Error>
    {
        self.inner.as_str()
    }
}

builder! {
#[builder(name = SettingsBuilder, derives=(Debug))]
#[derive(Debug)]
#[non_exhaustive]
pub struct Settings
{
    link_entrypoints: EntrypointFlags,
}
}

#[derive(Sole)]
pub struct Context
{
    _global_session: SlangGlobalSession,
    session: SlangSession,
    modules: HashMap<AssetId, Module>,
    programs: HashMap<AssetId, Program>,
}

impl Context
{
    #[must_use]
    pub fn get_module(&self, asset_id: &AssetId) -> Option<&Module>
    {
        self.modules.get(asset_id)
    }

    #[must_use]
    pub fn get_program(&self, asset_id: &AssetId) -> Option<&Program>
    {
        self.programs.get(asset_id)
    }

    #[tracing::instrument(skip_all, fields(module_file = module.file_path()))]
    pub fn compose_into_program(
        &self,
        module: Module,
        link_entrypoints: EntrypointFlags,
    ) -> Result<Program, ComposeProgramError>
    {
        let entry_points = link_entrypoints
            .iter()
            .filter_map(|entrypoint_flag| {
                let entrypoint_name = bitflags_match!(entrypoint_flag, {
                    EntrypointFlags::VERTEX => Some("vertex_main"),
                    EntrypointFlags::FRAGMENT => Some("fragment_main"),
                    _ => None
                })?;

                let Some(entry_point) = module.get_entry_point(entrypoint_name) else {
                    return Some(Err(ComposeProgramError::EntrypointNotFoundError {
                        entrypoint_name,
                    }));
                };

                Some(Ok(entry_point))
            })
            .collect::<Result<Vec<_>, _>>()?;

        let components = entry_points
            .into_iter()
            .map(|entry_point| SlangComponentType::from(entry_point.inner))
            .chain([SlangComponentType::from(module.inner)])
            .collect::<Vec<_>>();

        let program = self
            .session
            .create_composite_component_type(&components)
            .map_err(|err| ComposeProgramError::Other(Error(err)))?;

        let vertex_desc = if link_entrypoints.contains(EntrypointFlags::VERTEX) {
            // TODO: Do not have hard coded target here
            let Ok(program_reflection) = program.layout(0) else {
                unreachable!();
            };

            let program_reflection =
                self::reflection::Program { inner: program_reflection };

            Some(VertexDescription::new(
                &program_reflection
                    .get_entry_point_by_name("vertex_main")
                    .expect("Not possible"),
            )?)
        } else {
            None
        };

        Ok(Program {
            inner: program,
            metadata: ProgramMetadata { vertex_desc },
        })
    }
}

#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct ProgramMetadata
{
    /// If the program has a entry point in the vertex stage, this field will contain a
    /// description of the vertex type passed to the entry point.
    pub vertex_desc: Option<VertexDescription>,
}

#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct VertexDescription
{
    pub inputs: Arc<[VertexInputDescription]>,
}

impl VertexDescription
{
    #[tracing::instrument(skip_all, fields(vs_entry_point_name = vs_entrypoint.name()))]
    pub fn new(
        vs_entrypoint: &self::reflection::EntryPointReflection<'_>,
    ) -> Result<Self, VertexDescriptionError>
    {
        if vs_entrypoint.stage() != self::reflection::Stage::Vertex {
            return Err(VertexDescriptionError::EntrypointNotInVertexStage);
        }

        let inputs = VarInputDfsIter::new(vs_entrypoint).inspect(|var_input| {
            if var_input.type_layout.kind() != self::reflection::TypeKind::Struct {
                return;
            }

            let Some(semantic_name) = var_input.var_layout.semantic_name() else {
                return;
            };

            tracing::warn!(
                "Semantic name '{}' of '{}' will be inherited by it's fields",
                semantic_name,
                var_input.var_layout.name().unwrap_or("<unnamed>")
            );
        });

        let mut seen_inputs = HashSet::<VertexInputSemName>::new();

        Ok(Self {
            inputs: inputs
                .filter(|var_input| {
                    var_input.type_layout.kind() != self::reflection::TypeKind::Struct
                })
                .map(|var_input| {
                    let name = var_input.var_layout.name().unwrap_or("<unnamed>");

                    let semantic_name =
                        var_input.var_layout.semantic_name().ok_or_else(|| {
                            VertexDescriptionError::VertexInputMissingSemanticName {
                                name: name.to_owned(),
                            }
                        })?;

                    let semantic_name =
                        VertexInputSemName::from_semantic_name(semantic_name);

                    if seen_inputs.contains(&semantic_name) {
                        return Err(
                            VertexDescriptionError::VertexInputHasOccupiedSemanticName {
                                name: name.to_owned(),
                            },
                        );
                    }

                    let (
                        self::reflection::TypeKind::Scalar
                        | self::reflection::TypeKind::Vector,
                        Some(scalar_type),
                    ) = (
                        var_input.type_layout.kind(),
                        var_input.type_layout.scalar_type(),
                    )
                    else {
                        return Err(VertexDescriptionError::UnsupportedVertexInputType {
                            name: name.to_owned(),
                        });
                    };

                    seen_inputs.insert(semantic_name.clone());

                    Ok(VertexInputDescription {
                        semantic_name,
                        index: var_input.index,
                        type_kind: var_input.type_layout.kind(),
                        scalar_type,
                    })
                })
                .collect::<Result<Vec<_>, _>>()?
                .into(),
        })
    }
}

#[derive(Debug)]
#[non_exhaustive]
pub struct VertexInputDescription
{
    pub semantic_name: VertexInputSemName,
    pub index: usize,
    pub type_kind: self::reflection::TypeKind,
    pub scalar_type: self::reflection::ScalarType,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum VertexInputSemName
{
    Position,
    Normal,
    UvFromTopLeft,
    Color,
    Other(Box<str>),
}

impl VertexInputSemName
{
    fn from_semantic_name(semantic_name: &str) -> Self
    {
        match semantic_name {
            STD_VERT_IN_SEM_NAME_POSITION => Self::Position,
            STD_VERT_IN_SEM_NAME_NORMAL => Self::Normal,
            STD_VERT_IN_SEM_NAME_UV_FROM_TOP_LEFT => Self::UvFromTopLeft,
            STD_VERT_IN_SEM_NAME_COLOR => Self::Color,
            _ => Self::Other(semantic_name.to_lowercase().into_boxed_str()),
        }
    }

    #[must_use]
    pub fn matches_vertex_label(&self, vertex_label: &VertexLabel) -> bool
    {
        match (self, vertex_label) {
            (Self::Position, VertexLabel::Position)
            | (Self::Normal, VertexLabel::Normal)
            | (Self::UvFromTopLeft, VertexLabel::UvFromTopLeft)
            | (Self::Color, VertexLabel::Color) => true,
            (Self::Other(other), VertexLabel::Other(other_vertex_label)) => {
                **other == *other_vertex_label
            }
            _ => false,
        }
    }
}

impl Display for VertexInputSemName
{
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result
    {
        match self {
            Self::Position => formatter.write_str(STD_VERT_IN_SEM_NAME_POSITION),
            Self::Normal => formatter.write_str(STD_VERT_IN_SEM_NAME_NORMAL),
            Self::UvFromTopLeft => {
                formatter.write_str(STD_VERT_IN_SEM_NAME_UV_FROM_TOP_LEFT)
            }
            Self::Color => formatter.write_str(STD_VERT_IN_SEM_NAME_COLOR),
            Self::Other(other) => {
                for character in other.chars() {
                    formatter.write_char(character.to_ascii_uppercase())?;
                }

                Ok(())
            }
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum VertexDescriptionError
{
    #[error("Entrypoint is not in vertex stage")]
    EntrypointNotInVertexStage,

    #[error("Type of vertex input '{name}' is not supported")]
    UnsupportedVertexInputType
    {
        name: String
    },

    #[error("Vertex input '{name}' is missing a semantic name")]
    VertexInputMissingSemanticName
    {
        name: String
    },

    #[error(
        "Vertex input '{name}' has a semantic name already used by another vertex input"
    )]
    VertexInputHasOccupiedSemanticName
    {
        name: String
    },
}

#[derive(Debug, thiserror::Error)]
pub enum ComposeProgramError
{
    #[error("Shader module does not have a '{entrypoint_name}' entry point")]
    EntrypointNotFoundError
    {
        entrypoint_name: &'static str
    },

    #[error("Failed to create vertex description")]
    VertexDescriptionCreationFailed(#[from] VertexDescriptionError),

    #[error(transparent)]
    Other(Error),
}

#[derive(Debug, thiserror::Error)]
#[error(transparent)]
pub struct Error(#[from] shader_slang::Error);

struct VarInputDfsIter<'a>
{
    stack: Vec<(
        self::reflection::VariableLayout<'a>,
        self::reflection::TypeLayout<'a>,
        usize,
    )>,
}

impl<'a> VarInputDfsIter<'a>
{
    fn new(entry_point: &self::reflection::EntryPointReflection<'a>) -> Self
    {
        Self {
            stack: entry_point
                .parameters()
                .map(|param| {
                    let Some(param_type_layout) = param.type_layout() else {
                        // I do not know in which cases type_layout can return None
                        unimplemented!();
                    };

                    (param, param_type_layout, 0)
                })
                .collect(),
        }
    }
}

impl<'a> Iterator for VarInputDfsIter<'a>
{
    type Item = VarInput<'a>;

    fn next(&mut self) -> Option<Self::Item>
    {
        let (var_layout, type_layout, index) = loop {
            let (var_layout, type_layout, acc_offset) = self.stack.pop()?;

            if let Some(offset) = var_layout.varying_input_offset() {
                break (var_layout, type_layout, acc_offset + offset);
            }
        };

        if type_layout.kind() == self::reflection::TypeKind::Struct {
            self.stack.extend(type_layout.fields().map(|field| {
                let Some(field_type_layout) = field.type_layout() else {
                    // I do not know in which cases type_layout can return
                    // None
                    unimplemented!();
                };

                (field, field_type_layout, index)
            }));
        }

        Some(VarInput { var_layout, type_layout, index })
    }
}

struct VarInput<'a>
{
    var_layout: self::reflection::VariableLayout<'a>,
    type_layout: self::reflection::TypeLayout<'a>,
    index: usize,
}

fn import(
    asset_submitter: &mut AssetSubmitter<'_>,
    file_path: &Path,
    settings: Option<&'_ Settings>,
) -> Result<(), ImportError>
{
    if file_path.extension().is_none_or(|ext| ext != "slang") {
        return Err(ImportError::UnsupportedFileType);
    }

    let file_name = file_path
        .file_name()
        .ok_or(ImportError::NoPathFileName)?
        .to_str()
        .ok_or(ImportError::PathFileNameNotUtf8)?;

    let file_path_canonicalized = file_path
        .canonicalize()
        .map_err(ImportError::CanonicalizePathFailed)?;

    asset_submitter.submit_store(ModuleSource {
        name: file_name.to_owned().into(),
        file_path: file_path_canonicalized.into(),
        source: std::fs::read_to_string(file_path)
            .map_err(ImportError::ReadFileFailed)?
            .into(),
        link_entrypoints: settings
            .map(|settings| settings.link_entrypoints)
            .unwrap_or_default(),
    });

    Ok(())
}

#[derive(Debug, thiserror::Error)]
enum ImportError
{
    #[error("File type is unsupported")]
    UnsupportedFileType,

    #[error("Failed to read file")]
    ReadFileFailed(#[source] std::io::Error),

    #[error("Asset path does not have a file name")]
    NoPathFileName,

    #[error("Asset path file name is not valid UTF8")]
    PathFileNameNotUtf8,

    #[error("Failed to canonicalize asset path")]
    CanonicalizePathFailed(#[source] std::io::Error),
}

declare_entity! {
pub IMPORT_SHADERS_PHASE: (Phase, pair!(ChildOf, { *HANDLE_ASSETS_PHASE }));
}

pub(super) fn prepare(collector: &mut crate::ecs::extension::Collector<'_>)
{
    let Some(global_session) = SlangGlobalSession::new() else {
        tracing::error!("Unable to create global shader-slang session");
        return;
    };

    let session_options = shader_slang::CompilerOptions::default()
        .optimization(shader_slang::OptimizationLevel::None)
        .matrix_layout_row(true)
        .debug_information(SlangDebugInfoLevel::Maximal)
        .no_mangle(true);

    let target_desc = shader_slang::TargetDesc::default()
        .format(shader_slang::CompileTarget::Glsl)
        // .format(shader_slang::CompileTarget::Spirv)
        .profile(global_session.find_profile("glsl_330"));
    // .profile(global_session.find_profile("spirv_1_5"));

    let targets = [target_desc];

    let session_desc = shader_slang::SessionDesc::default()
        .targets(&targets)
        .search_paths([""])
        .options(&session_options);

    let Some(session) = global_session.create_session(&session_desc) else {
        tracing::error!("Failed to create shader-slang session");
        return;
    };

    collector
        .add_sole(Context {
            _global_session: global_session,
            session,
            modules: HashMap::new(),
            programs: HashMap::new(),
        })
        .ok();

    collector.spawn_declared_entity(&IMPORT_SHADERS_PHASE);

    collector.add_system(*INIT_PHASE, initialize);
    collector.add_system(*IMPORT_SHADERS_PHASE, load_modules);
}

fn initialize(mut assets: Single<Assets>) -> Result<(), crate::Error>
{
    let assets = assets.get_mut()?;

    assets.set_importer::<ModuleSource, Settings, ImportError>(import);

    Ok(())
}

#[tracing::instrument(skip_all)]
fn load_modules(
    mut context: Single<Context>,
    assets: Single<Assets>,
    asset_events: Single<AssetEvents>,
) -> Result<(), crate::Error>
{
    let Ok(context) = context.get_mut() else {
        unreachable!();
    };

    let assets = assets.get()?;
    let asset_events = asset_events.get()?;

    for AssetEvent::Loaded(asset_id, asset_label) in asset_events.iter_last_tick() {
        let asset_handle = AssetHandle::<ModuleSource>::from_id(*asset_id);

        if !assets.is_loaded_and_has_type(&asset_handle) {
            continue;
        }

        let Some(module_source) = assets.get(&asset_handle) else {
            unreachable!();
        };

        tracing::debug!(asset_label=?asset_label, "Loading shader module");

        let module = match load_module(&context.session, module_source) {
            Ok(module) => module,
            Err(err) => {
                tracing::error!("Failed to load shader module: {err}");
                continue;
            }
        };

        context.modules.insert(*asset_id, module.clone());

        if !module_source.link_entrypoints.is_empty() {
            assert!(!context.programs.contains_key(asset_id));

            let shader_program = match context
                .compose_into_program(module, module_source.link_entrypoints)
            {
                Ok(shader_program) => shader_program,
                Err(err) => {
                    tracing::error!(
                        "Failed to compose shader into program: {:#}",
                        crate::Error::new(err)
                    );
                    continue;
                }
            };

            let linked_shader_program = match shader_program.into_linked() {
                Ok(linked_shader_program) => linked_shader_program,
                Err(err) => {
                    tracing::error!("Failed to link shader: {err}");
                    continue;
                }
            };

            context.programs.insert(*asset_id, linked_shader_program);
        }
    }

    Ok(())
}

fn load_module(
    session: &SlangSession,
    module_source: &ModuleSource,
) -> Result<Module, Error>
{
    let module = session.load_module_from_source_string(
        &module_source.name,
        &module_source.file_path.to_string_lossy(),
        &module_source.source,
    )?;

    Ok(Module { inner: module })
}