space-crush/glc/src/shader.rs

use std::ffi::CString;
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::ptr::{null, null_mut};
use std::str::from_utf8;

use crate::{
    error::Error,
    matrix::Matrix4f,
    misc::{AsEnum, Bindable},
    vector::{Vector2f, Vector4f},
};

/* Programm */

#[derive(Default)]
pub struct Program {
    id: gl::GLuint,
}

impl Program {
    pub fn builder() -> Builder {
        Builder::default()
    }

    pub fn from_shaders<I>(iter: I) -> Result<Self, Error>
    where
        I: IntoIterator<Item = Shader>,
    {
        Self::from_shaders_result(iter.into_iter().map(Ok))
    }

    pub fn from_shaders_result<I, E>(iter: I) -> Result<Self, E>
    where
        I: IntoIterator<Item = Result<Shader, E>>,
        E: From<Error>,
    {
        let mut builder = Self::builder();

        for shader in iter.into_iter() {
            builder = builder.add_shader(shader?);
        }

        let program = builder.build()?;

        Ok(program)
    }

    pub fn id(&self) -> gl::GLuint {
        self.id
    }

    pub fn uniform<T>(&self, location: T, uniform: Uniform<'_>) -> Result<(), Error>
    where
        T: IntoUniformLocation,
    {
        let location = IntoUniformLocation::into_location(location, &self)?;

        match uniform {
            Uniform::Matrix4f(v) => gl::uniform_matrix_4fv(location, 1, gl::FALSE, v.as_ptr()),
            Uniform::Vector4f(v) => gl::uniform_4fv(location, 1, v.as_ptr()),
            Uniform::Vector2f(v) => gl::uniform_2fv(location, 1, v.as_ptr()),
            Uniform::Float(v) => gl::uniform_1f(location, v),
            Uniform::Texture(i) => gl::uniform_1i(location, i),
            Uniform::TextureVec(i) => gl::uniform_1iv(location, i.len() as _, i.as_ptr()),
        }

        Ok(())
    }

    pub fn uniform_location<T>(&self, name: T) -> Result<gl::GLint, Error>
    where
        T: Into<String>,
    {
        let name = Into::<String>::into(name);
        let name = CString::new(name)?;
        let location = Error::checked(|| gl::get_uniform_location(self.id, name.as_ptr()))?;

        if location >= 0 {
            Ok(location)
        } else {
            Err(Error::ShaderUnknownUniform(name.into_string().unwrap()))
        }
    }

    pub fn uniform_block_index<T>(&self, name: T) -> Result<gl::GLuint, Error>
    where
        T: Into<String>,
    {
        let name = Into::<String>::into(name);
        let name = CString::new(name)?;
        let location = Error::checked(|| gl::get_uniform_block_index(self.id, name.as_ptr()))?;

        if location != gl::INVALID_INDEX {
            Ok(location)
        } else {
            Err(Error::ShaderUnknownUniform(name.into_string().unwrap()))
        }
    }

    pub fn uniform_block_binding<T>(&self, index: T, binding: gl::GLuint) -> Result<(), Error>
    where
        T: IntoUniformBlockIndex,
    {
        let index = IntoUniformBlockIndex::into_index(index, &self)?;

        gl::uniform_block_binding(self.id, index, binding);

        Ok(())
    }
}

impl Drop for Program {
    fn drop(&mut self) {
        gl::delete_program(self.id);
    }
}

impl Bindable for Program {
    fn bind(&self) {
        gl::use_program(self.id);
    }

    fn unbind(&self) {
        gl::use_program(0);
    }
}

/* Shader */

pub struct Shader {
    id: gl::GLuint,
}

impl Shader {
    pub fn from_string<F>(type_: Type, mut source: String, f: F) -> Result<Self, Error>
    where
        F: Fn(&str) -> Result<String, Error>,
    {
        while let Some(p) = source.find("#pragma include ") {
            let s = &source[p..];
            let e = s.find('\n').unwrap_or_else(|| s.len());
            let s = &s[..e];
            let c = f(&s[16..])?;

            source = source.replace(s, &c);
        }

        let id = gl::create_shader(type_.as_enum());
        let id = Error::err_if(&0, id)?;

        let source = CString::new(source)?;
        let ptr = source.as_ptr();
        let ptr: *const *const _ = &ptr;
        gl::shader_source(id, 1, ptr as *const *const _, null());
        gl::compile_shader(id);

        let mut success = 1;
        gl::get_shader_iv(id, gl::COMPILE_STATUS, &mut success);

        if success != 1 {
            let mut len = 0;
            gl::get_shader_iv(id, gl::INFO_LOG_LENGTH, &mut len);

            let mut buffer = Vec::<u8>::with_capacity(len as usize + 1);
            buffer.resize(len as usize, 0);

            gl::get_shader_info_log(
                id,
                len,
                null_mut(),
                buffer.as_mut_ptr() as *mut gl::types::GLchar,
            );

            let msg = from_utf8(&buffer)?;
            let code = source.into_string().unwrap();
            let code = code
                .lines()
                .enumerate()
                .map(|(i, s)| format!("{:03}: {}\n", i + 1, s))
                .collect();

            return Err(Error::ShaderCompile {
                code,
                error: msg.into(),
            });
        }

        Ok(Self { id })
    }

    pub fn from_reader<R, F>(type_: Type, reader: &mut R, f: F) -> Result<Self, Error>
    where
        R: Read,
        F: Fn(&str) -> Result<String, Error>,
    {
        let mut source = String::new();
        reader.read_to_string(&mut source)?;

        Self::from_string(type_, source, f)
    }

    pub fn from_file<P, F>(type_: Type, path: P, f: F) -> Result<Self, Error>
    where
        P: AsRef<Path>,
        F: Fn(&str) -> Result<String, Error>,
    {
        let mut file = File::open(&path)?;

        match Self::from_reader(type_, &mut file, f) {
            Ok(v) => Ok(v),
            Err(Error::ShaderCompile { code, error }) => Err(Error::ShaderCompile {
                code: format!("{}\n{}", path.as_ref().display(), code),
                error,
            }),
            Err(err) => Err(err),
        }
    }

    pub fn id(&self) -> gl::GLuint {
        self.id
    }
}

impl Drop for Shader {
    fn drop(&mut self) {
        gl::delete_shader(self.id);
    }
}

/* Builder */

#[derive(Default)]
pub struct Builder {
    shaders: Vec<Shader>,
    transform_feedback_varyings: Option<(TransformFeedbackVaryingsMode, &'static [&'static str])>,
}

impl Builder {
    pub fn add_shader(mut self, shader: Shader) -> Self {
        self.shaders.push(shader);

        self
    }

    pub fn set_transform_feedback_varyings(
        mut self,
        mode: TransformFeedbackVaryingsMode,
        value: &'static [&'static str],
    ) -> Self {
        self.transform_feedback_varyings = Some((mode, value));

        self
    }

    pub fn build(self) -> Result<Program, Error> {
        let id = gl::create_program();
        let id = Error::err_if(&0, id)?;

        for shader in &self.shaders {
            gl::attach_shader(id, shader.id());
        }

        if let Some((mode, tfv)) = self.transform_feedback_varyings {
            let tfv = tfv
                .iter()
                .map(|v: &&str| CString::new(*v))
                .collect::<Result<Vec<_>, _>>()?;
            let tfv = tfv
                .iter()
                .map(|v: &CString| v.as_ptr() as *const gl::GLchar)
                .collect::<Vec<_>>();
            let ptr: *const *const gl::GLchar = tfv.as_slice().as_ptr();

            Error::checked(|| {
                gl::transform_feedback_varyings(id, tfv.len() as _, ptr, mode.as_enum())
            })?;
        }

        gl::link_program(id);

        for shader in &self.shaders {
            gl::detach_shader(id, shader.id());
        }

        let mut success = 1;
        gl::get_program_iv(id, gl::LINK_STATUS, &mut success);

        if success != 1 {
            let mut len = 0;
            gl::get_program_iv(id, gl::INFO_LOG_LENGTH, &mut len);

            let mut buffer = Vec::<u8>::with_capacity(len as usize + 1);
            buffer.resize(len as usize, 0);

            gl::get_program_info_log(
                id,
                len,
                null_mut(),
                buffer.as_mut_ptr() as *mut gl::types::GLchar,
            );

            let msg = from_utf8(&buffer).map_err(Error::Utf8Error)?;

            return Err(Error::ShaderLink(msg.into()));
        }

        Ok(Program { id })
    }
}

/* Type */

#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Type {
    Vertex,
    Compute,
    Fragment,
    Geometry,
    TesslationControl,
    TesslationEvaluation,
}

impl AsEnum for Type {
    fn as_enum(&self) -> gl::GLenum {
        match self {
            Self::Vertex => gl::VERTEX_SHADER,
            Self::Compute => gl::COMPUTE_SHADER,
            Self::Fragment => gl::FRAGMENT_SHADER,
            Self::Geometry => gl::GEOMETRY_SHADER,
            Self::TesslationControl => gl::TESS_CONTROL_SHADER,
            Self::TesslationEvaluation => gl::TESS_EVALUATION_SHADER,
        }
    }
}

/* TransformFeedbackVaryingsMode */

pub enum TransformFeedbackVaryingsMode {
    Interleaved,
    Separate,
}

impl AsEnum for TransformFeedbackVaryingsMode {
    fn as_enum(&self) -> gl::GLenum {
        match self {
            Self::Interleaved => gl::INTERLEAVED_ATTRIBS,
            Self::Separate => gl::SEPARATE_ATTRIBS,
        }
    }
}

/* Uniform */

pub enum Uniform<'a> {
    Matrix4f(&'a Matrix4f),
    Vector4f(&'a Vector4f),
    Vector2f(&'a Vector2f),
    Float(f32),
    Texture(gl::GLint),
    TextureVec(&'a [gl::GLint]),
}

/* IntoUniformLocation */

pub trait IntoUniformLocation {
    fn into_location(self, program: &Program) -> Result<gl::GLint, Error>;
}

impl IntoUniformLocation for gl::GLint {
    fn into_location(self, _program: &Program) -> Result<gl::GLint, Error> {
        Ok(self)
    }
}

impl IntoUniformLocation for &str {
    fn into_location(self, program: &Program) -> Result<gl::GLint, Error> {
        program.uniform_location(self)
    }
}

impl IntoUniformLocation for String {
    fn into_location(self, program: &Program) -> Result<gl::GLint, Error> {
        program.uniform_location(self)
    }
}

/* IntoUniformBlockIndex */

pub trait IntoUniformBlockIndex {
    fn into_index(self, program: &Program) -> Result<gl::GLuint, Error>;
}

impl IntoUniformBlockIndex for gl::GLuint {
    fn into_index(self, _program: &Program) -> Result<gl::GLuint, Error> {
        Ok(self)
    }
}

impl IntoUniformBlockIndex for &str {
    fn into_index(self, program: &Program) -> Result<gl::GLuint, Error> {
        program.uniform_block_index(self)
    }
}

impl IntoUniformBlockIndex for String {
    fn into_index(self, program: &Program) -> Result<gl::GLuint, Error> {
        program.uniform_block_index(self)
    }
}