shader_sense/symbols/
shader_module_parser.rs

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