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

forge_kit/
utils.rs
