forge_kit/
utils.rs

1//! Utility functions for working with the ForgeScript AST
2
3use crate::parser::{AstNode, Span};
4
5/// Pretty-print the AST to a string
6pub fn format_ast(node: &AstNode) -> String {
7    let mut output = String::new();
8    format_ast_impl(node, &mut output, 0);
9    output
10}
11
12fn format_ast_impl(node: &AstNode, output: &mut String, depth: usize) {
13    let indent = "  ".repeat(depth);
14
15    match node {
16        AstNode::Program { body, span } => {
17            output.push_str(&format!(
18                "{}Program ({}..{})\n",
19                indent, span.start, span.end
20            ));
21            for child in body {
22                format_ast_impl(child, output, depth + 1);
23            }
24        }
25        AstNode::Text { content, span } => {
26            output.push_str(&format!(
27                "{}Text ({}..{}): {:?}\n",
28                indent, span.start, span.end, content
29            ));
30        }
31        AstNode::FunctionCall {
32            name,
33            args,
34            modifiers,
35            span,
36            ..
37        } => {
38            output.push_str(&format!(
39                "{}FunctionCall ({}..{}): ${}{}{}{}\n",
40                indent,
41                span.start,
42                span.end,
43                name,
44                if modifiers.silent { " [silent]" } else { "" },
45                if modifiers.negated { " [negated]" } else { "" },
46                if let Some(count) = &modifiers.count {
47                    format!(" [count: {}]", count)
48                } else {
49                    String::new()
50                }
51            ));
52            if let Some(args) = args {
53                for (i, arg) in args.iter().enumerate() {
54                    output.push_str(&format!(
55                        "{}  Arg {} ({}..{}):\n",
56                        indent, i, arg.span.start, arg.span.end
57                    ));
58                    for part in &arg.parts {
59                        format_ast_impl(part, output, depth + 2);
60                    }
61                }
62            }
63        }
64        AstNode::JavaScript { code, span } => {
65            output.push_str(&format!(
66                "{}JavaScript ({}..{}): {:?}\n",
67                indent, span.start, span.end, code
68            ));
69        }
70        AstNode::Escaped { content, span } => {
71            output.push_str(&format!(
72                "{}Escaped ({}..{}): {:?}\n",
73                indent, span.start, span.end, content
74            ));
75        }
76    }
77}
78
79/// Extract all function names from the AST
80pub fn extract_function_names(node: &AstNode) -> Vec<String> {
81    let mut names = Vec::new();
82    extract_function_names_impl(node, &mut names);
83    names
84}
85
86fn extract_function_names_impl(node: &AstNode, names: &mut Vec<String>) {
87    match node {
88        AstNode::Program { body, .. } => {
89            for child in body {
90                extract_function_names_impl(child, names);
91            }
92        }
93        AstNode::FunctionCall { name, args, .. } => {
94            names.push(name.clone());
95            if let Some(args) = args {
96                for arg in args {
97                    for part in &arg.parts {
98                        extract_function_names_impl(part, names);
99                    }
100                }
101            }
102        }
103        _ => {}
104    }
105}
106
107/// Count the number of nodes in the AST
108pub fn count_nodes(node: &AstNode) -> usize {
109    match node {
110        AstNode::Program { body, .. } => 1 + body.iter().map(count_nodes).sum::<usize>(),
111        AstNode::FunctionCall { args, .. } => {
112            1 + args
113                .as_ref()
114                .map(|args| {
115                    args.iter()
116                        .map(|arg| arg.parts.iter().map(count_nodes).sum::<usize>())
117                        .sum()
118                })
119                .unwrap_or(0)
120        }
121        _ => 1,
122    }
123}
124
125/// Get all text nodes from the AST
126pub fn extract_text_nodes(node: &AstNode) -> Vec<(String, Span)> {
127    let mut texts = Vec::new();
128    extract_text_nodes_impl(node, &mut texts);
129    texts
130}
131
132fn extract_text_nodes_impl(node: &AstNode, texts: &mut Vec<(String, Span)>) {
133    match node {
134        AstNode::Program { body, .. } => {
135            for child in body {
136                extract_text_nodes_impl(child, texts);
137            }
138        }
139        AstNode::Text { content, span } => {
140            texts.push((content.clone(), *span));
141        }
142        AstNode::FunctionCall { args, .. } => {
143            if let Some(args) = args {
144                for arg in args {
145                    for part in &arg.parts {
146                        extract_text_nodes_impl(part, texts);
147                    }
148                }
149            }
150        }
151        _ => {}
152    }
153}
154
155/// Find the deepest nesting level in the AST
156pub fn max_nesting_depth(node: &AstNode) -> usize {
157    max_nesting_depth_impl(node, 0)
158}
159
160fn max_nesting_depth_impl(node: &AstNode, current: usize) -> usize {
161    match node {
162        AstNode::Program { body, .. } => body
163            .iter()
164            .map(|n| max_nesting_depth_impl(n, current))
165            .max()
166            .unwrap_or(current),
167        AstNode::FunctionCall { args, .. } => {
168            let next = current + 1;
169            args.as_ref()
170                .map(|args| {
171                    args.iter()
172                        .flat_map(|arg| arg.parts.iter())
173                        .map(|part| max_nesting_depth_impl(part, next))
174                        .max()
175                        .unwrap_or(next)
176                })
177                .unwrap_or(next)
178        }
179        _ => current,
180    }
181}
182
183/// Check if the AST contains any JavaScript expressions
184pub fn contains_javascript(node: &AstNode) -> bool {
185    match node {
186        AstNode::Program { body, .. } => body.iter().any(contains_javascript),
187        AstNode::JavaScript { .. } => true,
188        AstNode::FunctionCall { args, .. } => args
189            .as_ref()
190            .map(|args| {
191                args.iter()
192                    .flat_map(|arg| arg.parts.iter())
193                    .any(contains_javascript)
194            })
195            .unwrap_or(false),
196        _ => false,
197    }
198}
199
200/// Get a slice of the source code for a given span
201pub fn get_source_slice<'a>(source: &'a str, span: Span) -> &'a str {
202    &source[span.start..span.end.min(source.len())]
203}
204
205/// Calculate statistics about the AST
206#[derive(Debug, Clone)]
207pub struct AstStats {
208    pub total_nodes: usize,
209    pub text_nodes: usize,
210    pub function_calls: usize,
211    pub javascript_nodes: usize,
212    pub escaped_nodes: usize,
213    pub max_depth: usize,
214    pub unique_functions: usize,
215}
216
217pub fn calculate_stats(node: &AstNode) -> AstStats {
218    let mut text_nodes = 0;
219    let mut function_calls = 0;
220    let mut javascript_nodes = 0;
221    let mut escaped_nodes = 0;
222
223    count_node_types(
224        node,
225        &mut text_nodes,
226        &mut function_calls,
227        &mut javascript_nodes,
228        &mut escaped_nodes,
229    );
230
231    let function_names = extract_function_names(node);
232    let mut unique = function_names.clone();
233    unique.sort();
234    unique.dedup();
235
236    AstStats {
237        total_nodes: count_nodes(node),
238        text_nodes,
239        function_calls,
240        javascript_nodes,
241        escaped_nodes,
242        max_depth: max_nesting_depth(node),
243        unique_functions: unique.len(),
244    }
245}
246
247fn count_node_types(
248    node: &AstNode,
249    text: &mut usize,
250    funcs: &mut usize,
251    js: &mut usize,
252    esc: &mut usize,
253) {
254    match node {
255        AstNode::Program { body, .. } => {
256            for child in body {
257                count_node_types(child, text, funcs, js, esc);
258            }
259        }
260        AstNode::Text { .. } => *text += 1,
261        AstNode::FunctionCall { args, .. } => {
262            *funcs += 1;
263            if let Some(args) = args {
264                for arg in args {
265                    for part in &arg.parts {
266                        count_node_types(part, text, funcs, js, esc);
267                    }
268                }
269            }
270        }
271        AstNode::JavaScript { .. } => *js += 1,
272        AstNode::Escaped { .. } => *esc += 1,
273    }
274}
275
276/// Flatten the AST into a linear sequence of nodes (depth-first)
277pub fn flatten_ast(node: &AstNode) -> Vec<AstNode> {
278    let mut nodes = Vec::new();
279    flatten_ast_impl(node, &mut nodes);
280    nodes
281}
282
283fn flatten_ast_impl(node: &AstNode, nodes: &mut Vec<AstNode>) {
284    nodes.push(node.clone());
285    match node {
286        AstNode::Program { body, .. } => {
287            for child in body {
288                flatten_ast_impl(child, nodes);
289            }
290        }
291        AstNode::FunctionCall { args, .. } => {
292            if let Some(args) = args {
293                for arg in args {
294                    for part in &arg.parts {
295                        flatten_ast_impl(part, nodes);
296                    }
297                }
298            }
299        }
300        _ => {}
301    }
302}
forge_kit/utils.rs

forge_kit/
utils.rs
