path: root/engine/src/shader_preprocessor.rs

use std::path::{Path, PathBuf};
use std::string::FromUtf8Error;

const PREINCLUDE_DIRECTIVE: &str = "#preinclude";

/// Preprocessor for shaders written in the OpenGL Shading Language.
pub struct ShaderPreprocessor
{
    read_file: fn(&Path) -> Result<Vec<u8>, std::io::Error>,
    base_dir: PathBuf,
}

impl ShaderPreprocessor
{
    pub fn new(base_dir: PathBuf) -> Self
    {
        Self {
            read_file: |path| std::fs::read(path),
            base_dir,
        }
    }

    pub fn preprocess(
        &self,
        shader_content: String,
        shader_file_path: &Path,
    ) -> Result<String, Error>
    {
        let mut preincludes = shader_content
            .match_indices(PREINCLUDE_DIRECTIVE)
            .peekable();

        if preincludes.peek().is_none() {
            // Shader content contains no preincludes
            return Ok(shader_content);
        };

        let mut preprocessed = shader_content.clone();

        let mut curr = shader_content.find(PREINCLUDE_DIRECTIVE);

        let mut last_start = 0;
        let mut span_line_offset = 0;

        while let Some(preinclude_start) = curr {
            let replacement_job = self.handle_preinclude(
                &preprocessed,
                shader_file_path,
                preinclude_start,
                span_line_offset,
            )?;

            let path = replacement_job.path.clone();

            let mut included = String::from_utf8(
                (self.read_file)(&self.base_dir.join(replacement_job.path.clone()))
                    .map_err(|err| Error::ReadIncludedShaderFailed {
                        source: err,
                        path: replacement_job.path.clone(),
                    })?,
            )
            .map_err(|err| Error::IncludedShaderInvalidUtf8 {
                source: err,
                path: path.clone(),
            })?;

            if let Some(first_line) = included.lines().next() {
                if first_line.starts_with("#version") {
                    included = included
                        .chars()
                        .skip_while(|character| *character != '\n')
                        .collect();
                }
            }

            let included_preprocessed =
                self.preprocess(included, &replacement_job.path)?;

            let start = replacement_job.start_index;
            let end = replacement_job.end_index;

            preprocessed.replace_range(start..end, &included_preprocessed);

            curr = preprocessed[last_start + 1..]
                .find(PREINCLUDE_DIRECTIVE)
                .map(|index| index + 1);

            last_start = preinclude_start + included_preprocessed.len();

            span_line_offset += included_preprocessed.lines().count();
        }

        Ok(preprocessed)
    }

    fn handle_preinclude(
        &self,
        shader_content: &str,
        shader_file_path: &Path,
        preinclude_start_index: usize,
        span_line_offset: usize,
    ) -> Result<ReplacementJob, Error>
    {
        let expect_token = |token: char, index: usize| {
            let token_found = shader_content.chars().nth(index).ok_or_else(|| {
                Error::ExpectedToken {
                    expected: token,
                    span: Span::new(
                        shader_content,
                        &self.base_dir.join(shader_file_path),
                        index,
                        span_line_offset,
                        preinclude_start_index,
                    ),
                }
            })?;

            if token_found != token {
                return Err(Error::InvalidToken {
                    expected: token,
                    found: token_found,
                    span: Span::new(
                        shader_content,
                        &self.base_dir.join(shader_file_path),
                        index,
                        span_line_offset,
                        preinclude_start_index,
                    ),
                });
            }

            Ok(())
        };

        let space_index = preinclude_start_index + PREINCLUDE_DIRECTIVE.len();
        let quote_open_index = space_index + 1;

        expect_token(' ', space_index)?;
        expect_token('"', quote_open_index)?;

        let buf = shader_content[quote_open_index + 1..]
            .chars()
            .take_while(|character| *character != '"')
            .map(|character| character as u8)
            .collect::<Vec<_>>();

        if buf.is_empty() {
            return Err(Error::ExpectedToken {
                expected: '"',
                span: Span::new(
                    shader_content,
                    &self.base_dir.join(shader_file_path),
                    shader_content.len() - 1,
                    span_line_offset,
                    preinclude_start_index,
                ),
            });
        }

        let path_len = buf.len();

        let path = PathBuf::from(String::from_utf8(buf).map_err(|err| {
            Error::PreincludePathInvalidUtf8 {
                source: err,
                span: Span::new(
                    shader_content,
                    &self.base_dir.join(shader_file_path),
                    quote_open_index + 1,
                    span_line_offset,
                    preinclude_start_index,
                ),
            }
        })?);

        Ok(ReplacementJob {
            start_index: preinclude_start_index,
            end_index: quote_open_index + 1 + path_len + 1,
            path,
        })
    }
}

struct ReplacementJob
{
    start_index: usize,
    end_index: usize,
    path: PathBuf,
}

/// Shader preprocessing error.
#[derive(Debug, thiserror::Error)]
pub enum Error
{
    #[error(
        "Invalid token at line {}, column {} of {}. Expected '{}', found '{}'",
        span.line,
        span.column,
        span.file.display(),
        expected,
        found
    )]
    InvalidToken
    {
        expected: char,
        found: char,
        span: Span,
    },

    #[error(
        "Expected token '{}' at line {}, column {} of {}. Found eof",
        expected,
        span.line,
        span.column,
        span.file.display(),
    )]
    ExpectedToken
    {
        expected: char, span: Span
    },

    #[error(
        "Preinclude path at line {}, column {} of {} is invalid UTF-8",
        span.line,
        span.column,
        span.file.display(),
    )]
    PreincludePathInvalidUtf8
    {
        #[source]
        source: FromUtf8Error,
        span: Span,
    },

    #[error("Failed to read included shader")]
    ReadIncludedShaderFailed
    {
        #[source]
        source: std::io::Error,
        path: PathBuf,
    },

    #[error("Included shader is not valid UTF-8")]
    IncludedShaderInvalidUtf8
    {
        #[source]
        source: FromUtf8Error,
        path: PathBuf,
    },
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Span
{
    line: usize,
    column: usize,
    file: PathBuf,
}

impl Span
{
    fn new(
        file_content: &str,
        file_path: &Path,
        char_index: usize,
        line_offset: usize,
        line_start_index: usize,
    ) -> Self
    {
        let line = find_line_of_index(file_content, char_index) + 1
            - line_offset.saturating_sub(1);

        Self {
            line,
            column: char_index - line_start_index + 1,
            file: file_path.to_path_buf(),
        }
    }
}

fn find_line_of_index(text: &str, index: usize) -> usize
{
    text.chars()
        .take(index + 1)
        .enumerate()
        .filter(|(_, character)| *character == '\n')
        .count()
}

#[cfg(test)]
mod tests
{
    use std::ffi::OsStr;
    use std::path::{Path, PathBuf};

    use super::{Error, ShaderPreprocessor};

    #[test]
    fn preprocess_no_directives_is_same()
    {
        assert_eq!(
            ShaderPreprocessor {
                read_file: |_| { unreachable!() },
                base_dir: PathBuf::new()
            }
            .preprocess("#version 330 core\n".to_string(), Path::new("foo.glsl"),)
            .unwrap(),
            "#version 330 core\n".to_string()
        );
    }

    #[test]
    fn preprocess_with_directives_works()
    {
        assert_eq!(
            ShaderPreprocessor {
                read_file: |_| { Ok(b"out vec4 FragColor;".to_vec()) },
                base_dir: PathBuf::new()
            }
            .preprocess(
                concat!(
                    "#version 330 core\n",
                    "\n",
                    "#preinclude \"foo.glsl\"\n",
                    "\n",
                    "void main() {}",
                )
                .to_string(),
                Path::new("foo.glsl"),
            )
            .unwrap(),
            concat!(
                "#version 330 core\n",
                "\n",
                "out vec4 FragColor;\n",
                "\n",
                "void main() {}",
            )
        );

        assert_eq!(
            ShaderPreprocessor {
                read_file: |_| { Ok(b"out vec4 FragColor;".to_vec()) },
                base_dir: PathBuf::new()
            }
            .preprocess(
                concat!(
                    "#version 330 core\n",
                    "\n",
                    "#preinclude \"bar.glsl\"\n",
                    "\n",
                    "in vec3 in_frag_color;\n",
                    "\n",
                    "void main() {}",
                )
                .to_string(),
                Path::new("foo.glsl"),
            )
            .unwrap(),
            concat!(
                "#version 330 core\n",
                "\n",
                "out vec4 FragColor;\n",
                "\n",
                "in vec3 in_frag_color;\n",
                "\n",
                "void main() {}",
            )
        );

        assert_eq!(
            ShaderPreprocessor {
                read_file: |path| {
                    if path == OsStr::new("bar.glsl") {
                        Ok(b"out vec4 FragColor;".to_vec())
                    } else {
                        Ok(concat!(
                            "uniform sampler2D input_texture;\n",
                            "in vec2 in_texture_coords;"
                        )
                        .as_bytes()
                        .to_vec())
                    }
                },
                base_dir: PathBuf::new()
            }
            .preprocess(
                concat!(
                    "#version 330 core\n",
                    "\n",
                    "#preinclude \"bar.glsl\"\n",
                    "\n",
                    "in vec3 in_frag_color;\n",
                    "\n",
                    "#preinclude \"foo.glsl\"\n",
                    "\n",
                    "void main() {}",
                )
                .to_string(),
                Path::new("foo.glsl"),
            )
            .unwrap(),
            concat!(
                "#version 330 core\n",
                "\n",
                "out vec4 FragColor;\n",
                "\n",
                "in vec3 in_frag_color;\n",
                "\n",
                "uniform sampler2D input_texture;\n",
                "in vec2 in_texture_coords;\n",
                "\n",
                "void main() {}",
            )
        );
    }

    #[test]
    fn preprocess_invalid_shader_does_not_work()
    {
        let path = Path::new("foo.glsl");

        let res = ShaderPreprocessor {
            read_file: |_| Ok(b"out vec4 FragColor;".to_vec()),
            base_dir: PathBuf::new(),
        }
        .preprocess(
            concat!(
                "#version 330 core\n",
                "\n",
                // Missing "
                "#preinclude foo.glsl\"\n",
                "\n",
                "void main() {}",
            )
            .to_string(),
            path,
        );

        let Err(Error::InvalidToken { expected, found, span }) = res else {
            panic!(
                "Expected result to be Err(Error::InvalidToken {{ ... }}), is {res:?}"
            );
        };

        assert_eq!(expected, '"');
        assert_eq!(found, 'f');
        assert_eq!(span.line, 3);
        assert_eq!(span.column, 13);
        assert_eq!(span.file, path);
    }

    #[test]
    fn preprocess_error_has_correct_span()
    {
        let path = Path::new("foo.glsl");

        let res = ShaderPreprocessor {
            read_file: |path| {
                if path == OsStr::new("bar.glsl") {
                    Ok(concat!(
                        "out vec4 FragColor;\n",
                        "in vec2 in_texture_coords;\n",
                        "in float foo;"
                    )
                    .as_bytes()
                    .to_vec())
                } else {
                    Ok(b"uniform sampler2D input_texture;".to_vec())
                }
            },
            base_dir: PathBuf::new(),
        }
        .preprocess(
            concat!(
                "#version 330 core\n",
                "\n",
                "#preinclude \"bar.glsl\"\n",
                "\n",
                "in vec3 in_frag_color;\n",
                "\n",
                "#preinclude\"foo.glsl\"\n",
                "\n",
                "void main() {}",
            )
            .to_string(),
            Path::new("foo.glsl"),
        );

        let Err(Error::InvalidToken { expected, found, span }) = res else {
            panic!(
                "Expected result to be Err(Error::InvalidToken {{ ... }}), is {res:?}"
            );
        };

        assert_eq!(expected, ' ');
        assert_eq!(found, '"');
        assert_eq!(span.line, 7);
        assert_eq!(span.column, 12);
        assert_eq!(span.file, path);
    }

    #[test]
    fn preprocess_included_shader_with_include_works()
    {
        assert_eq!(
            ShaderPreprocessor {
                read_file: |path| {
                    if path == OsStr::new("bar.glsl") {
                        Ok(concat!(
                            "#preinclude \"foo.glsl\"\n",
                            "\n",
                            "out vec4 FragColor;"
                        )
                        .as_bytes()
                        .to_vec())
                    } else {
                        Ok(concat!(
                            "uniform sampler2D input_texture;\n",
                            "in vec2 in_texture_coords;"
                        )
                        .as_bytes()
                        .to_vec())
                    }
                },
                base_dir: PathBuf::new()
            }
            .preprocess(
                concat!(
                    "#version 330 core\n",
                    "\n",
                    "#preinclude \"bar.glsl\"\n",
                    "\n",
                    "in vec3 in_frag_color;\n",
                    "\n",
                    "void main() {}",
                )
                .to_string(),
                Path::new("foo.glsl"),
            )
            .unwrap(),
            concat!(
                "#version 330 core\n",
                "\n",
                "uniform sampler2D input_texture;\n",
                "in vec2 in_texture_coords;\n",
                "\n",
                "out vec4 FragColor;\n",
                "\n",
                "in vec3 in_frag_color;\n",
                "\n",
                "void main() {}",
            )
        );
    }

    #[test]
    fn preprocess_included_shader_with_include_error_span_is_correct()
    {
        let res = ShaderPreprocessor {
            read_file: |path| {
                if path == OsStr::new("bar.glsl") {
                    Ok(concat!(
                        // ' instead of "
                        "#preinclude 'foo.glsl\"\n",
                        "\n",
                        "out vec4 FragColor;"
                    )
                    .as_bytes()
                    .to_vec())
                } else {
                    Ok(concat!(
                        "uniform sampler2D input_texture;\n",
                        "in vec2 in_texture_coords;"
                    )
                    .as_bytes()
                    .to_vec())
                }
            },
            base_dir: PathBuf::new(),
        }
        .preprocess(
            concat!(
                "#version 330 core\n",
                "\n",
                "#preinclude \"bar.glsl\"\n",
                "\n",
                "in vec3 in_frag_color;\n",
                "\n",
                "void main() {}",
            )
            .to_string(),
            Path::new("foo.glsl"),
        );

        let Err(Error::InvalidToken { expected, found, span }) = res else {
            panic!(
                "Expected result to be Err(Error::InvalidToken {{ ... }}), is {res:?}"
            );
        };

        assert_eq!(expected, '"');
        assert_eq!(found, '\'');
        assert_eq!(span.line, 1);
        assert_eq!(span.column, 13);
        assert_eq!(span.file, Path::new("bar.glsl"));
    }
}