radar-gi/src/components/snippets/mod.rs

use regex::{Captures, Regex};
use std::{
    cell::{Ref, RefCell},
    collections::{HashMap, HashSet},
    ops::Add,
    rc::Rc,
    sync::atomic::{AtomicUsize, Ordering},
};
pub use util::Variable;

use crate::{
    errors::{Error, Result},
    utils::{find_file, Code, CodeBlock, Function, Hook, ShareVariable, SnippetCode},
};
mod util;
pub use util::merge_includes;
use util::*;

use super::CodeComponent;

static SNIP_ID: AtomicUsize = AtomicUsize::new(0);

#[derive(Debug, Clone)]
pub enum InputType {
    Vertex(usize, Variable),
    Other(Variable),
}

pub enum CodeType<S: std::borrow::Borrow<str>> {
    Code(S),
    Path(std::path::PathBuf),
}

#[derive(Debug, Clone)]
pub struct Snippet {
    id: usize,
    name: &'static str,
    parsed: CodeBlock,
    link: Option<Box<Snippet>>,
    alias: Option<HashMap<String, String>>,
    chained: String,
    main: Option<String>,
}

impl Snippet {
    pub fn new<S: std::borrow::Borrow<str>>(
        name: &'static str,
        code: CodeType<S>,
        mangling: bool,
        main: Option<String>,
    ) -> Result<Self> {
        let code = match code {
            CodeType::Code(code) => code.borrow().to_string(),
            CodeType::Path(path) => {
                let code = find_file(path).expect("Failed to find file");
                code
            }
        };

        let merged_code = merge_includes(code)
            .map(|code| {
                let clean_code = remove_comments(&code);
                let clean_code = remove_version(&clean_code);
                clean_code
            })
            .map_err(|v| {
                let err_msg = format!("Failed to merge includes cause of {:?}", v);
                Error::InvalidSnippet(err_msg)
            })?;

        let id = SNIP_ID.fetch_add(1, Ordering::SeqCst);

        let mut parsed_code = CodeBlock::new(&merged_code)?;
        let call = if let Some(main) = main {
            Some(main)
        } else {
            parsed_code
                .all_functions()
                .first()
                .map(|f| f.borrow().name.clone())
        };

        let alias = if mangling {
            let alias = parsed_code.mangling(id.to_string());
            Some(alias)
        } else {
            None
        };

        let chained = format!(
            "// START {} //\n{}\n// END {} //\n",
            name, parsed_code, name
        );

        Ok(Self {
            id,
            name,
            parsed: parsed_code,
            chained,
            alias,
            link: None,
            main: call,
        })
    }

    pub fn all_hooks(&self) -> HashSet<Hook> {
        self.parsed.all_hooks()
    }

    pub fn chain(mut self, snippet: Snippet) -> Self {
        self.chained = format!(
            "// START {} //\n{}\n// END {} //\n{}\n\n",
            snippet.name, snippet.chained, snippet.name, self.chained,
        );

        self.link = Some(Box::new(snippet));
        self
    }

    pub fn call(&self, paras: &Vec<String>) -> Option<String> {
        (self.main.as_ref()).map(|name| {
            if let Some(link) = &self.link {
                let call_name = self.alias.as_ref().map(|a| a.get(name).unwrap()).unwrap_or(name);
                let c = link.call(paras);
                if let Some(c) = c {
                    return format!("{}({})", call_name, c);
                } else {
                    return format!("{}()", call_name);
                }
            } else {
                let call_name = self.alias.as_ref().map(|a| a.get(name).unwrap()).unwrap_or(name);
                // let call_name = self.alias.get(name).unwrap();
                format!("{}({})", call_name, paras.join(", "))
            }
        })
    }

    pub fn prepare_code(&self) -> &str {
        &self.chained
    }

    pub fn find_symbol(&self, name: &str) -> Option<&String> {
        self.alias.as_ref().map(|a| a.get(name)).flatten()
    }

    pub fn find_variable(&self, name: &str) -> Option<&Rc<RefCell<ShareVariable>>> {
        self.find_symbol(name)
            .map(|v| self.parsed.find_variable(v))
            .flatten()
    }
}

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

impl Add for Snippet {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        let mut raw_code = self.parsed.to_string();
        let code = rhs.parsed.to_string();

        raw_code.push_str(&code);
        Snippet::new(self.name, CodeType::Code(raw_code),false, None).unwrap()
    }
}

impl CodeComponent for Snippet {
    fn add_snippet_after(self, snippet: Snippet) -> Self {
        self + snippet
    }

    fn add_snippet_before(self, snippet: Snippet) -> Self {
        snippet + self
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_snippet_new() {
        let code = r#"
            #version 450
            #include "math/constants.glsl"
            void main() {
                gl_Position = vec4(0.0);
            }
        "#;

        let snippet = Snippet::new("polar", CodeType::Code(code), true, None).unwrap();

        let snippet2 = Snippet::new("polar2", CodeType::Code(code),true, None).unwrap();

        let snippet3 = snippet.clone() + snippet2.clone();

        let snippet = snippet.chain(snippet2.chain(snippet3));

        println!("{}", snippet.prepare_code());
    }
}