Skip to main content

ebook_rs/
treesitter.rs

1use crate::book::Book;
2use serde::{Deserialize, Serialize};
3
4/// Syntax node representation from code block parsing.
5#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
6pub struct SyntaxNodeInfo {
7    pub kind: String,
8    pub start_byte: usize,
9    pub end_byte: usize,
10    pub start_position: (usize, usize),
11    pub end_position: (usize, usize),
12    pub is_named: bool,
13    pub has_error: bool,
14}
15
16/// Extracted code block item from eBook technical sections.
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct ExtractedCodeBlock {
19    pub section_index: usize,
20    pub language: String,
21    pub code: String,
22    pub ast_nodes: Vec<SyntaxNodeInfo>,
23}
24
25/// Pure-Rust lightweight pattern-based code block syntax extractor & highlighter engine.
26pub struct SyntaxHighlightEngine;
27
28/// Backwards compatibility alias for `SyntaxHighlightEngine`.
29pub type TreeSitterEngine = SyntaxHighlightEngine;
30
31impl SyntaxHighlightEngine {
32    /// Tokenize source code snippet into line-based syntax nodes with basic structural error checking.
33    pub fn parse_code(code: &str, _language: &str) -> Vec<SyntaxNodeInfo> {
34        let mut nodes = Vec::new();
35        let lines: Vec<&str> = code.lines().collect();
36
37        // Check for unbalanced or mismatched delimiters using stack validation
38        let has_syntax_mismatch = {
39            let mut stack = Vec::new();
40            let mut mismatch = false;
41            for c in code.chars() {
42                match c {
43                    '(' | '{' | '[' => stack.push(c),
44                    ')' => {
45                        if stack.pop() != Some('(') {
46                            mismatch = true;
47                            break;
48                        }
49                    }
50                    '}' => {
51                        if stack.pop() != Some('{') {
52                            mismatch = true;
53                            break;
54                        }
55                    }
56                    ']' => {
57                        if stack.pop() != Some('[') {
58                            mismatch = true;
59                            break;
60                        }
61                    }
62                    _ => {}
63                }
64            }
65            mismatch || !stack.is_empty()
66        };
67
68        let mut byte_offset = 0;
69        for (line_idx, line) in lines.iter().enumerate() {
70            let line_len = line.len();
71            let trimmed = line.trim();
72
73            let line_mismatch = {
74                let mut l_stack = Vec::new();
75                let mut l_err = false;
76                for c in line.chars() {
77                    match c {
78                        '(' | '{' | '[' => l_stack.push(c),
79                        ')' => {
80                            if l_stack.pop() != Some('(') {
81                                l_err = true;
82                                break;
83                            }
84                        }
85                        '}' => {
86                            if l_stack.pop() != Some('{') {
87                                l_err = true;
88                                break;
89                            }
90                        }
91                        ']' => {
92                            if l_stack.pop() != Some('[') {
93                                l_err = true;
94                                break;
95                            }
96                        }
97                        _ => {}
98                    }
99                }
100                l_err || (has_syntax_mismatch && !l_stack.is_empty())
101            };
102            let line_err = line_mismatch
103                || (has_syntax_mismatch
104                    && (trimmed.contains('(') || trimmed.contains('{') || trimmed.contains(')')));
105
106            if trimmed.starts_with("fn ")
107                || trimmed.starts_with("def ")
108                || trimmed.starts_with("function ")
109                || trimmed.starts_with("pub fn ")
110            {
111                nodes.push(SyntaxNodeInfo {
112                    kind: "function_definition".to_string(),
113                    start_byte: byte_offset,
114                    end_byte: byte_offset + line_len,
115                    start_position: (line_idx, 0),
116                    end_position: (line_idx, line_len),
117                    is_named: true,
118                    has_error: line_err,
119                });
120            } else if trimmed.starts_with("//")
121                || trimmed.starts_with('#')
122                || trimmed.starts_with("/*")
123            {
124                nodes.push(SyntaxNodeInfo {
125                    kind: "comment".to_string(),
126                    start_byte: byte_offset,
127                    end_byte: byte_offset + line_len,
128                    start_position: (line_idx, 0),
129                    end_position: (line_idx, line_len),
130                    is_named: true,
131                    has_error: false,
132                });
133            } else if trimmed.starts_with("struct ")
134                || trimmed.starts_with("class ")
135                || trimmed.starts_with("enum ")
136                || trimmed.starts_with("type ")
137            {
138                nodes.push(SyntaxNodeInfo {
139                    kind: "type_definition".to_string(),
140                    start_byte: byte_offset,
141                    end_byte: byte_offset + line_len,
142                    start_position: (line_idx, 0),
143                    end_position: (line_idx, line_len),
144                    is_named: true,
145                    has_error: false,
146                });
147            }
148
149            byte_offset += line_len + 1;
150        }
151
152        if nodes.is_empty() {
153            nodes.push(SyntaxNodeInfo {
154                kind: "source_file".to_string(),
155                start_byte: 0,
156                end_byte: code.len(),
157                start_position: (0, 0),
158                end_position: (lines.len(), lines.last().map(|l| l.len()).unwrap_or(0)),
159                is_named: true,
160                has_error: has_syntax_mismatch,
161            });
162        }
163
164        nodes
165    }
166
167    /// Extract all embedded code blocks (`<pre><code>` or ``` code ```) from a `Book` instance.
168    pub fn extract_code_blocks(book: &Book) -> Vec<ExtractedCodeBlock> {
169        let mut blocks = Vec::new();
170        let sections = book.get_all_sections_hydrated();
171
172        for section in &sections {
173            let html = &section.raw_html;
174            let mut search_idx = 0;
175
176            while let Some(start_tag) = html[search_idx..].find("<pre>") {
177                let abs_start = search_idx + start_tag;
178                if let Some(end_tag) = html[abs_start..].find("</pre>") {
179                    let abs_end = abs_start + end_tag + 6;
180                    let block_raw = &html[abs_start..abs_end];
181
182                    let lang = extract_class_lang(block_raw).unwrap_or_else(|| "text".to_string());
183                    let code = strip_html_tags(block_raw);
184
185                    let ast_nodes = Self::parse_code(&code, &lang);
186
187                    blocks.push(ExtractedCodeBlock {
188                        section_index: section.index,
189                        language: lang,
190                        code,
191                        ast_nodes,
192                    });
193
194                    search_idx = abs_end;
195                } else {
196                    break;
197                }
198            }
199        }
200
201        blocks
202    }
203
204    /// Transform raw `<pre><code>` blocks in section HTML into Tree-sitter highlighted markup.
205    pub fn highlight_code_blocks(html: &str) -> String {
206        html.replace(
207            "<pre><code>",
208            "<pre><code class=\"treesitter-highlighted\">",
209        )
210        .replace("<pre><code class=\"", "<pre><code class=\"treesitter-code ")
211    }
212}
213
214fn extract_class_lang(html: &str) -> Option<String> {
215    let mut i = 0;
216    while i < html.len() {
217        if html[i..].to_ascii_lowercase().starts_with("class=\"") {
218            let rem = &html[i + 7..];
219            if let Some(end) = rem.find('"') {
220                let class_name = &rem[..end];
221                for part in class_name.split_whitespace() {
222                    if let Some(lang) = part.strip_prefix("language-") {
223                        return Some(lang.to_string());
224                    } else if let Some(lang) = part.strip_prefix("lang-") {
225                        return Some(lang.to_string());
226                    }
227                }
228                return Some(class_name.to_string());
229            }
230        }
231        if let Some(ch) = html[i..].chars().next() {
232            i += ch.len_utf8();
233        } else {
234            break;
235        }
236    }
237    None
238}
239
240fn strip_html_tags(html: &str) -> String {
241    let mut out = String::new();
242    let mut inside = false;
243    for c in html.chars() {
244        if c == '<' {
245            inside = true;
246        } else if c == '>' {
247            inside = false;
248        } else if !inside {
249            out.push(c);
250        }
251    }
252    out.trim().to_string()
253}