Skip to main content

shader_sense/symbols/
shader_module_parser.rs

1//! Entry point for tree-sitter parsing.
2use std::path::Path;
3
4use tree_sitter::InputEdit;
5
6use crate::{position::ShaderRange, shader::ShadingLanguage, shader_error::ShaderError};
7
8use super::shader_module::ShaderModule;
9
10/// Handle the creation and update of internal tree_sitter AST stored into a [`ShaderModule`]
11/// ```
12/// use shader_sense::symbols::shader_module_parser::ShaderModuleParser;
13/// use shader_sense::position::{ShaderPosition, ShaderRange};
14/// use std::path::Path;
15/// let shader_path = Path::new("./test/hlsl/ok.hlsl");
16/// let shader_content = std::fs::read_to_string(shader_path).unwrap();
17/// let mut shader_module_parser = ShaderModuleParser::hlsl();
18/// let mut shader_module = shader_module_parser.create_module(shader_path, &shader_content).unwrap();
19/// // Here we can simply insert a new text at the position 3, 4.
20/// shader_module_parser.update_module_partial(
21///     &mut shader_module,
22///     &ShaderRange::new(
23///         ShaderPosition::new(3, 4),
24///         ShaderPosition::new(3, 4)
25///     ),
26///     &String::from("inserted text")
27/// ).unwrap();
28/// ```
29pub struct ShaderModuleParser {
30    tree_sitter_parser: tree_sitter::Parser,
31}
32
33pub fn get_tree_sitter_language(shading_language: ShadingLanguage) -> tree_sitter::Language {
34    match shading_language {
35        ShadingLanguage::Wgsl => tree_sitter_wgsl_bevy::LANGUAGE.into(),
36        ShadingLanguage::Hlsl => tree_sitter_hlsl::LANGUAGE_HLSL.into(),
37        ShadingLanguage::Glsl => tree_sitter_glsl::LANGUAGE_GLSL.into(),
38    }
39}
40
41impl ShaderModuleParser {
42    pub fn glsl() -> Self {
43        Self::from_shading_language(ShadingLanguage::Glsl)
44    }
45    pub fn hlsl() -> Self {
46        Self::from_shading_language(ShadingLanguage::Hlsl)
47    }
48    pub fn wgsl() -> Self {
49        Self::from_shading_language(ShadingLanguage::Wgsl)
50    }
51    pub fn from_shading_language(shading_language: ShadingLanguage) -> Self {
52        let mut tree_sitter_parser = tree_sitter::Parser::new();
53        tree_sitter_parser
54            .set_language(&get_tree_sitter_language(shading_language))
55            .expect("Error loading grammar");
56        Self { tree_sitter_parser }
57    }
58    // Create shader module from file.
59    pub fn create_module(
60        &mut self,
61        file_path: &Path,
62        shader_content: &str,
63    ) -> Result<ShaderModule, ShaderError> {
64        match self.tree_sitter_parser.parse(shader_content, None) {
65            Some(tree) => Ok(ShaderModule {
66                file_path: file_path.into(),
67                content: shader_content.into(),
68                tree,
69            }),
70            None => Err(ShaderError::ParseSymbolError(format!(
71                "Failed to parse AST for file {}",
72                file_path.display()
73            ))),
74        }
75    }
76    // Update whole content of symbol tree
77    pub fn update_module(
78        &mut self,
79        module: &mut ShaderModule,
80        new_text: &str,
81    ) -> Result<(), ShaderError> {
82        self.update_module_partial(module, &ShaderRange::whole(&module.content), new_text)
83    }
84    // Update partial content of symbol tree
85    pub fn update_module_partial(
86        &mut self,
87        module: &mut ShaderModule,
88        old_range: &ShaderRange,
89        new_text: &str,
90    ) -> Result<(), ShaderError> {
91        let mut new_shader_content = module.content.clone();
92        let old_start_byte_offset = old_range.start.to_byte_offset(&module.content)?;
93        let old_end_byte_offset = old_range.end.to_byte_offset(&module.content)?;
94        new_shader_content.replace_range(old_start_byte_offset..old_end_byte_offset, new_text);
95
96        let line_count = new_text.lines().count();
97        let tree_sitter_range = tree_sitter::Range {
98            start_byte: old_start_byte_offset,
99            end_byte: old_end_byte_offset,
100            start_point: tree_sitter::Point {
101                row: old_range.start.line as usize,
102                column: old_range.start.pos as usize,
103            },
104            end_point: tree_sitter::Point {
105                row: old_range.end.line as usize,
106                column: old_range.end.pos as usize,
107            },
108        };
109        module.tree.edit(&InputEdit {
110            start_byte: tree_sitter_range.start_byte,
111            old_end_byte: tree_sitter_range.end_byte,
112            new_end_byte: tree_sitter_range.start_byte + new_text.len(),
113            start_position: tree_sitter_range.start_point,
114            old_end_position: tree_sitter_range.end_point,
115            new_end_position: tree_sitter::Point {
116                row: if line_count == 0 {
117                    tree_sitter_range.start_point.row + new_text.len()
118                } else {
119                    new_text.lines().last().as_slice().len()
120                },
121                column: tree_sitter_range.start_point.column + line_count,
122            },
123        });
124        // Update the tree.
125        match self
126            .tree_sitter_parser
127            .parse(&new_shader_content, Some(&module.tree))
128        {
129            Some(new_tree) => {
130                module.tree = new_tree;
131                module.content = new_shader_content.clone();
132                Ok(())
133            }
134            None => Err(ShaderError::ParseSymbolError(format!(
135                "Failed to update AST for file {}.",
136                module.file_path.display()
137            ))),
138        }
139    }
140}