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
33impl ShaderModuleParser {
34    pub fn glsl() -> Self {
35        Self::from_shading_language(ShadingLanguage::Glsl)
36    }
37    pub fn hlsl() -> Self {
38        Self::from_shading_language(ShadingLanguage::Hlsl)
39    }
40    pub fn wgsl() -> Self {
41        Self::from_shading_language(ShadingLanguage::Wgsl)
42    }
43    pub fn from_shading_language(shading_language: ShadingLanguage) -> Self {
44        let mut tree_sitter_parser = tree_sitter::Parser::new();
45        tree_sitter_parser
46            .set_language(&match shading_language {
47                // TODO:WGSL: Wgsl is not yet compatible with tree-sitter-language.
48                // A PR is waiting to be merged here https://github.com/tree-sitter-grammars/tree-sitter-wgsl-bevy/pull/19
49                ShadingLanguage::Wgsl => tree_sitter_hlsl::LANGUAGE_HLSL.into(), //tree_sitter_wgsl_bevy::language(),
50                ShadingLanguage::Hlsl => tree_sitter_hlsl::LANGUAGE_HLSL.into(),
51                ShadingLanguage::Glsl => tree_sitter_glsl::LANGUAGE_GLSL.into(),
52            })
53            .expect("Error loading grammar");
54        Self { tree_sitter_parser }
55    }
56    // Create shader module from file.
57    pub fn create_module(
58        &mut self,
59        file_path: &Path,
60        shader_content: &str,
61    ) -> Result<ShaderModule, ShaderError> {
62        match self.tree_sitter_parser.parse(shader_content, None) {
63            Some(tree) => Ok(ShaderModule {
64                file_path: file_path.into(),
65                content: shader_content.into(),
66                tree,
67            }),
68            None => Err(ShaderError::ParseSymbolError(format!(
69                "Failed to parse AST for file {}",
70                file_path.display()
71            ))),
72        }
73    }
74    // Update whole content of symbol tree
75    pub fn update_module(
76        &mut self,
77        module: &mut ShaderModule,
78        new_text: &str,
79    ) -> Result<(), ShaderError> {
80        self.update_module_partial(module, &ShaderRange::whole(&module.content), new_text)
81    }
82    // Update partial content of symbol tree
83    pub fn update_module_partial(
84        &mut self,
85        module: &mut ShaderModule,
86        old_range: &ShaderRange,
87        new_text: &str,
88    ) -> Result<(), ShaderError> {
89        let mut new_shader_content = module.content.clone();
90        let old_start_byte_offset = old_range.start.to_byte_offset(&module.content)?;
91        let old_end_byte_offset = old_range.end.to_byte_offset(&module.content)?;
92        new_shader_content.replace_range(old_start_byte_offset..old_end_byte_offset, new_text);
93
94        let line_count = new_text.lines().count();
95        let tree_sitter_range = tree_sitter::Range {
96            start_byte: old_start_byte_offset,
97            end_byte: old_end_byte_offset,
98            start_point: tree_sitter::Point {
99                row: old_range.start.line as usize,
100                column: old_range.start.pos as usize,
101            },
102            end_point: tree_sitter::Point {
103                row: old_range.end.line as usize,
104                column: old_range.end.pos as usize,
105            },
106        };
107        module.tree.edit(&InputEdit {
108            start_byte: tree_sitter_range.start_byte,
109            old_end_byte: tree_sitter_range.end_byte,
110            new_end_byte: tree_sitter_range.start_byte + new_text.len(),
111            start_position: tree_sitter_range.start_point,
112            old_end_position: tree_sitter_range.end_point,
113            new_end_position: tree_sitter::Point {
114                row: if line_count == 0 {
115                    tree_sitter_range.start_point.row + new_text.len()
116                } else {
117                    new_text.lines().last().as_slice().len()
118                },
119                column: tree_sitter_range.start_point.column + line_count,
120            },
121        });
122        // Update the tree.
123        match self
124            .tree_sitter_parser
125            .parse(&new_shader_content, Some(&module.tree))
126        {
127            Some(new_tree) => {
128                module.tree = new_tree;
129                module.content = new_shader_content.clone();
130                Ok(())
131            }
132            None => Err(ShaderError::ParseSymbolError(format!(
133                "Failed to update AST for file {}.",
134                module.file_path.display()
135            ))),
136        }
137    }
138}