Skip to main content

nargo_compiler/
adapter.rs

1use nargo_ir::{IRModule, JsExpr, JsProgram, JsStmt};
2use nargo_types::Result;
3
4pub struct TsAdapter;
5
6impl TsAdapter {
7    pub fn new() -> Self {
8        Self
9    }
10
11    pub fn transform(&self, ir: &mut IRModule) -> Result<()> {
12        if let Some(script) = &mut ir.script {
13            self.transform_script(script)?;
14        }
15        if let Some(script) = &mut ir.script_server {
16            self.transform_script(script)?;
17        }
18        if let Some(script) = &mut ir.script_client {
19            self.transform_script(script)?;
20        }
21        for test in &mut ir.tests {
22            self.transform_script(&mut test.body)?;
23        }
24        Ok(())
25    }
26
27    fn transform_script(&self, script: &mut JsProgram) -> Result<()> {
28        // 1. Identify state variables and computed macros
29        let mut reactive_vars = Vec::new();
30        let mut computed_vars = Vec::new();
31
32        // 预分配空间,避免频繁扩容
33        reactive_vars.reserve(10);
34        computed_vars.reserve(5);
35
36        for stmt in &mut script.body {
37            if let JsStmt::VariableDecl { kind, id, init, .. } = stmt {
38                if kind == "let" || kind == "var" {
39                    // Check if it's already a reactive call
40                    let is_already_reactive = if let Some(JsExpr::Call { callee, .. }) = init {
41                        if let JsExpr::Identifier(name, _, _) = &**callee {
42                            name == "ref" || name == "signal" || name == "createSignal" || name == "reactive" || name == "computed" || name == "createComputed"
43                        }
44                        else {
45                            false
46                        }
47                    }
48                    else {
49                        false
50                    };
51
52                    if !is_already_reactive {
53                        reactive_vars.push(id.clone());
54                    }
55                }
56                else if kind == "const" {
57                    if let Some(JsExpr::Call { callee, args, span, trivia }) = init {
58                        let mut should_transform_effect = false;
59                        if let JsExpr::Identifier(name, _, _) = &**callee {
60                            if name == "$computed" || name == "computed" || name == "createComputed" {
61                                computed_vars.push(id.clone());
62                            }
63                            else if name == "$effect" || name == "effect" || name == "createEffect" {
64                                should_transform_effect = true;
65                            }
66                        }
67
68                        if should_transform_effect {
69                            *callee = Box::new(JsExpr::Identifier("effect".to_string(), *span, trivia.clone()));
70                            // If it's $effect(expr), wrap in arrow function
71                            if args.len() == 1 && !matches!(args[0], JsExpr::ArrowFunction { .. }) {
72                                let expr = args.pop().unwrap();
73                                args.push(JsExpr::ArrowFunction { params: Vec::new(), body: Box::new(expr), span: *span, trivia: trivia.clone() });
74                            }
75                        }
76                    }
77                }
78            }
79        }
80
81        // 2. Transform declarations
82        for stmt in &mut script.body {
83            if let JsStmt::VariableDecl { kind, id, init, span, trivia } = stmt {
84                if (kind == "let" || kind == "var") && reactive_vars.contains(id) {
85                    *kind = "const".to_string();
86                    let initial_value = init.take().unwrap_or(JsExpr::Literal(nargo_types::NargoValue::Null, *span, trivia.clone()));
87
88                    *init = Some(JsExpr::Call { callee: Box::new(JsExpr::Identifier("ref".to_string(), *span, trivia.clone())), args: vec![initial_value], span: *span, trivia: trivia.clone() });
89                }
90                else if kind == "const" && computed_vars.contains(id) {
91                    if let Some(JsExpr::Call { callee, args, span, trivia }) = init {
92                        if let JsExpr::Identifier(name, _, _) = &**callee {
93                            if name == "$computed" || name == "computed" || name == "createComputed" {
94                                *callee = Box::new(JsExpr::Identifier("computed".to_string(), *span, trivia.clone()));
95                                // Transform args (wrap in arrow function if not already)
96                                if args.len() == 1 {
97                                    if !matches!(args[0], JsExpr::ArrowFunction { .. }) {
98                                        let expr = args.pop().unwrap();
99                                        args.push(JsExpr::ArrowFunction { params: Vec::new(), body: Box::new(expr), span: *span, trivia: trivia.clone() });
100                                    }
101                                }
102                            }
103                        }
104                    }
105                }
106            }
107        }
108
109        // 3. Transform usages: x -> x.value
110        let mut all_reactive = reactive_vars;
111        all_reactive.extend(computed_vars);
112
113        for stmt in &mut script.body {
114            self.transform_stmt(stmt, &all_reactive);
115        }
116
117        Ok(())
118    }
119
120    fn transform_stmt(&self, stmt: &mut JsStmt, reactive_vars: &[String]) {
121        match stmt {
122            JsStmt::Expr(expr, _, _) => self.transform_expr(expr, reactive_vars),
123            JsStmt::VariableDecl { init, .. } => {
124                if let Some(expr) = init {
125                    self.transform_expr(expr, reactive_vars);
126                }
127            }
128            JsStmt::Return(expr, _, _) => {
129                if let Some(e) = expr {
130                    self.transform_expr(e, reactive_vars);
131                }
132            }
133            JsStmt::If { test, consequent, alternate, .. } => {
134                self.transform_expr(test, reactive_vars);
135                self.transform_stmt(consequent, reactive_vars);
136                if let Some(alt) = alternate {
137                    self.transform_stmt(alt, reactive_vars);
138                }
139            }
140            JsStmt::While { test, body, .. } => {
141                self.transform_expr(test, reactive_vars);
142                self.transform_stmt(body, reactive_vars);
143            }
144            JsStmt::For { init, test, update, body, .. } => {
145                if let Some(i) = init {
146                    self.transform_stmt(i, reactive_vars);
147                }
148                if let Some(t) = test {
149                    self.transform_expr(t, reactive_vars);
150                }
151                if let Some(u) = update {
152                    self.transform_expr(u, reactive_vars);
153                }
154                self.transform_stmt(body, reactive_vars);
155            }
156            JsStmt::Block(stmts, _, _) => {
157                for s in stmts {
158                    self.transform_stmt(s, reactive_vars);
159                }
160            }
161            JsStmt::FunctionDecl { body, .. } => {
162                for s in body {
163                    self.transform_stmt(s, reactive_vars);
164                }
165            }
166            _ => {}
167        }
168    }
169
170    fn transform_expr(&self, expr: &mut JsExpr, reactive_vars: &[String]) {
171        match expr {
172            JsExpr::Identifier(name, span, trivia) => {
173                if reactive_vars.contains(name) {
174                    // 直接使用引用,避免不必要的克隆
175                    let name_clone = name.clone();
176                    let span_clone = *span;
177                    let trivia_clone = trivia.clone();
178
179                    *expr = JsExpr::Member { object: Box::new(JsExpr::Identifier(name_clone, span_clone, trivia_clone.clone())), property: Box::new(JsExpr::Identifier("value".to_string(), span_clone, trivia_clone.clone())), computed: false, span: span_clone, trivia: trivia_clone };
180                }
181            }
182            JsExpr::Unary { argument, .. } => self.transform_expr(argument, reactive_vars),
183            JsExpr::Binary { left, right, .. } => {
184                self.transform_expr(left, reactive_vars);
185                self.transform_expr(right, reactive_vars);
186            }
187            JsExpr::Call { callee, args, span, trivia } => {
188                // Handle macros
189                if let JsExpr::Identifier(name, _, _) = &**callee {
190                    if name == "$effect" || name == "effect" || name == "createEffect" {
191                        *callee = Box::new(JsExpr::Identifier("effect".to_string(), *span, trivia.clone()));
192                        if args.len() == 1 && !matches!(args[0], JsExpr::ArrowFunction { .. }) {
193                            let expr = args.pop().unwrap();
194                            args.push(JsExpr::ArrowFunction { params: Vec::new(), body: Box::new(expr), span: *span, trivia: trivia.clone() });
195                        }
196                    }
197                }
198
199                self.transform_expr(callee, reactive_vars);
200                for arg in args {
201                    self.transform_expr(arg, reactive_vars);
202                }
203            }
204            JsExpr::Member { object, property, computed, .. } => {
205                self.transform_expr(object, reactive_vars);
206                if *computed {
207                    self.transform_expr(property, reactive_vars);
208                }
209            }
210            JsExpr::Array(elements, _, _) => {
211                for el in elements {
212                    self.transform_expr(el, reactive_vars);
213                }
214            }
215            JsExpr::Object(properties, _, _) => {
216                for value in properties.values_mut() {
217                    self.transform_expr(value, reactive_vars);
218                }
219            }
220            JsExpr::ArrowFunction { body, .. } => {
221                self.transform_expr(body, reactive_vars);
222            }
223            JsExpr::Conditional { test, consequent, alternate, .. } => {
224                self.transform_expr(test, reactive_vars);
225                self.transform_expr(consequent, reactive_vars);
226                self.transform_expr(alternate, reactive_vars);
227            }
228            JsExpr::TemplateLiteral { expressions, .. } => {
229                for e in expressions {
230                    self.transform_expr(e, reactive_vars);
231                }
232            }
233            JsExpr::TseElement { attributes, children, .. } => {
234                for attr in attributes {
235                    if let Some(val) = &mut attr.value {
236                        self.transform_expr(val, reactive_vars);
237                    }
238                }
239                for child in children {
240                    self.transform_expr(child, reactive_vars);
241                }
242            }
243            _ => {}
244        }
245    }
246}