Skip to main content

soma_codemode/
normalize.rs

1pub fn normalize_user_code(code: &str) -> String {
2    let code = strip_code_fences(code.trim()).trim();
3    if code.is_empty() {
4        return "async () => {}".to_string();
5    }
6    if let Some(inner) = code.strip_prefix("export default ") {
7        return normalize_export_default(inner);
8    }
9    if let Some(name) = function_declaration_name(code) {
10        return format!("async () => {{\n{code}\nreturn {name}();\n}}");
11    }
12    if is_bare_arrow_expression(code) || is_bare_function_expression(code) {
13        return code.trim_end_matches(';').trim().to_string();
14    }
15    if let Some((before, after)) = code.rsplit_once(';') {
16        let trailing = after.trim();
17        if !trailing.is_empty() && looks_like_expression(trailing) {
18            return format!("async () => {{\n{before};\nreturn ({trailing})\n}}");
19        }
20    } else if looks_like_expression(code) && !code.trim_start().starts_with("return ") {
21        return format!("async () => {{\nreturn ({code})\n}}");
22    }
23    format!("async () => {{\n{code}\n}}")
24}
25
26fn normalize_export_default(inner: &str) -> String {
27    let inner = inner.trim().trim_end_matches(';').trim();
28    if inner.starts_with("async function") || inner.starts_with("function") {
29        return format!("async () => {{\nreturn ({inner})();\n}}");
30    }
31    if inner.starts_with("class") {
32        return format!("async () => {{\nreturn ({inner});\n}}");
33    }
34    normalize_user_code(inner)
35}
36
37fn strip_code_fences(code: &str) -> &str {
38    let trimmed = code.trim();
39    for lang in ["javascript", "typescript", "tsx", "jsx", "js", "ts", ""] {
40        let prefix = if lang.is_empty() {
41            "```\n".to_string()
42        } else {
43            format!("```{lang}\n")
44        };
45        if trimmed.starts_with(&prefix) && trimmed.ends_with("```") {
46            let without_prefix = &trimmed[prefix.len()..trimmed.len() - 3];
47            return without_prefix.trim();
48        }
49    }
50    trimmed
51}
52
53fn function_declaration_name(code: &str) -> Option<&str> {
54    let trimmed = code.trim_start();
55    let rest = trimmed
56        .strip_prefix("async function ")
57        .or_else(|| trimmed.strip_prefix("function "))?;
58    let end = rest.find(|ch: char| !(ch == '_' || ch == '$' || ch.is_ascii_alphanumeric()))?;
59    let name = &rest[..end];
60    (!name.is_empty()).then_some(name)
61}
62
63fn is_bare_function_expression(code: &str) -> bool {
64    let trimmed = code.trim_start();
65    trimmed.starts_with("async function") || trimmed.starts_with("function")
66}
67
68fn is_bare_arrow_expression(code: &str) -> bool {
69    let trimmed = code.trim();
70    trimmed.starts_with("async ") && trimmed.contains("=>")
71        || trimmed.starts_with('(') && trimmed.contains("=>")
72}
73
74fn looks_like_expression(code: &str) -> bool {
75    let trimmed = code.trim();
76    !(trimmed.starts_with("let ")
77        || trimmed.starts_with("const ")
78        || trimmed.starts_with("var ")
79        || trimmed.starts_with("if ")
80        || trimmed.starts_with("for ")
81        || trimmed.starts_with("while ")
82        || trimmed.starts_with("throw ")
83        || trimmed.ends_with('}'))
84}