radar-gi/src/components/shader.rs

use std::borrow::Borrow;

use super::{
    snippets::{merge_includes, CodeType, Snippet, Variable},
    CodeComponent,
};
use crate::{
    errors::{Error, Result},
    utils::{find_file, Code, CodeBlock},
};
use glow::HasContext;
use log::{info, warn};
pub use utils::fetchcode;
pub type ShaderType = u32;
#[derive(Debug)]
pub struct Shader {
    target: u32,
    #[cfg(feature = "inparser")]
    parsed: CodeBlock,
    parsed: String,
    #[cfg(feature = "inparser")]
    pub hooks: Vec<String>,
}

impl std::fmt::Display for Shader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.parsed)
    }
}

impl Shader {
    pub fn new(target: ShaderType, code: CodeType) -> Result<Self> {
        let code = match code {
            CodeType::Code(code) => code,
            CodeType::Path(path) => {
                let code = find_file(&path).expect(&format!(
                    "Failed to find file {}",
                    &path.as_os_str().to_str().unwrap()
                ));
                code
            }
        };

        #[cfg(feature = "inparser")]
        {
            let code = merge_includes(code).map_err(|e| Error::InvalidSnippet(e.to_string()))?;
            let parsed = CodeBlock::new(&code)?;
            let hooks = parsed.all_hooks().iter().map(|h| h.name.clone()).collect();
            info!("Shader:{} Hooks: {:?}", target, hooks);
            return Ok(Self {
                target,
                parsed,
                hooks,
            });
        }

        Ok(Self {
            target,
            parsed: code,
        })
    }

    pub fn compile(
        &self,
        version: &str,
        gl: &glow::Context,
    ) -> <glow::Context as HasContext>::Shader {
        let shader = unsafe {
            let shader = gl.create_shader(self.target).expect("Cannot create shader");
            gl.shader_source(
                shader,
                &format!("{}\n{}", &format!("#version {}", version), self.to_string()),
            );
            gl.compile_shader(shader);
            if !gl.get_shader_compile_status(shader) {
                panic!("{}", gl.get_shader_info_log(shader))
            }
            shader
        };
        shader
    }

    #[cfg(feature = "inparser")]
    pub fn define(&mut self, s: &str) {
        self.parsed
            .insert(Code::new(&format!("#define {}\n", s)), 0);
    }

    #[cfg(feature = "inparser")]
    pub fn set_hook(&mut self, hook: &str, code: &Snippet) {
        self.parsed.set_hook(hook, code);
    }

    pub fn target(&self) -> u32 {
        self.target
    }
}

impl CodeComponent for Shader {
    fn add_snippet_after(self, snippet: Snippet) -> Self {
        let new_code = self.to_string() + &snippet.to_string();
        let result = Self::new(self.target, CodeType::Code(new_code)).unwrap();
        result
    }

    fn add_snippet_before(self, snippet: Snippet) -> Self {
        let new_code = snippet.to_string() + &self.to_string();
        let result = Self::new(self.target, CodeType::Code(new_code)).unwrap();
        result
    }
}

mod utils {
    use once_cell::sync::Lazy;
    use std::collections::HashMap;

    use crate::{
        components::{CodeType, Snippet},
        utils::CodeBlock,
    };

    static TYPES: Lazy<HashMap<usize, &'static str>> = Lazy::new(|| {
        [
            (1, "float"),
            (2, "vec2 "),
            (3, "vec3 "),
            (4, "vec4 "),
            (9, "mat3 "),
            (16, "mat4 "),
        ]
        .iter()
        .cloned()
        .collect::<std::collections::HashMap<_, _>>()
    });

    pub fn fetchcode<'a, V: AsRef<[(&'a str, usize)]>>(data_types: V, prefix: &str) -> Snippet {
        let mut header = String::from(
            "
            uniform sampler2D uniforms;
            uniform vec3 uniforms_shape;
            layout(location = 0) in float collection_index;\n
        ",
        );

        for data_type in data_types.as_ref().iter() {
            if data_type.0 != "__unused__" {
                if let Some(type_name) = TYPES.get(&data_type.1) {
                    header.push_str(&format!("out {} {}{};\n", type_name, prefix, data_type.0));
                }
            }
        }

        let mut body = String::from(
            "
            void fetch_uniforms() {
                float rows = uniforms_shape.x;
                float cols = uniforms_shape.y;
                float count = uniforms_shape.z;
                float index = collection_index;
                int index_x = int(mod(index, (floor(cols/(count/4.0))))) * int(count/4.0);
                int index_y = int(floor(index / (floor(cols/(count/4.0)))));
                float size_x = cols - 1.0;
                float size_y = rows - 1.0;
                float ty = 0.0;
                if (size_y > 0.0)
                    ty = float(index_y)/size_y;
                int i = index_x;
                vec4 _uniform;
                    ",
        );

        for data_type in data_types.as_ref().iter() {
            if data_type.0 != "__unused__" {
                let component_count = data_type.1;
                let mut store = 0;
                let mut shift = 0;
                let mut count = component_count;
                while count > 0 {
                    if store == 0 {
                        body.push_str(&format!(
                            "\n    _uniform = texture2D(uniforms, vec2(float(i++)/size_x,ty));\n"
                        ));
                        store = 4;
                    }
                    let (a, b) = match store {
                        4 => (
                            "xyzw",
                            match shift {
                                0 => "xyzw",
                                1 => "yzw",
                                2 => "zw",
                                3 => "w",
                                _ => "",
                            },
                        ),
                        3 => (
                            "yzw",
                            match shift {
                                0 => "yzw",
                                1 => "zw",
                                2 => "w",
                                _ => "",
                            },
                        ),
                        2 => (
                            "zw",
                            match shift {
                                0 => "zw",
                                1 => "w",
                                _ => "",
                            },
                        ),
                        1 => ("w", "w"),
                        _ => ("", ""),
                    };
                    let min_length = std::cmp::min(b.len(), count);
                    body.push_str(&format!(
                        "    {}{}{} = _uniform.{};\n",
                        prefix,
                        data_type.0,
                        &b[0..min_length],
                        &a[0..min_length]
                    ));
                    count -= min_length;
                    shift += min_length;
                    store -= min_length;
                }
            }
        }

        body.push_str("}\n\n");
        let input = header + &body;

        let result =
            Snippet::new("fetch_uniform", CodeType::from_code(input), false, None).unwrap();

        result
    }
}

mod test {
    use super::utils::fetchcode;
    use super::*;

    #[test]
    fn test_shader() {
        let shader = Shader::new(
            glow::VERTEX_SHADER,
            CodeType::from_path("agg-fast-path.vert"),
        )
        .unwrap();

        println!("{}", shader.parsed);
    }

    #[test]
    fn test_fetchcode() {
        let code = fetchcode([("position", 3), ("color", 4)], "a_");
        println!("{}", code);
    }
}