Skip to main content

oak_wat/builder/
mod.rs

1use crate::{WatParser, ast::*, language::WatLanguage, lexer::token_type::WatTokenType, parser::element_type::WatElementType};
2use oak_core::{Builder, BuilderCache, GreenNode, OakDiagnostics, OakError, Parser, RedNode, RedTree, SourceText, TextEdit, builder::BuildOutput, source::Source};
3
4/// AST builder for the WebAssembly Text (WAT) format.
5#[derive(Clone, Copy)]
6pub struct WatBuilder<'config> {
7    /// Language configuration
8    config: &'config WatLanguage,
9}
10
11impl<'config> WatBuilder<'config> {
12    /// Creates a new `WatBuilder` with the given configuration.
13    pub const fn new(config: &'config WatLanguage) -> Self {
14        Self { config }
15    }
16}
17
18impl<'config> Builder<WatLanguage> for WatBuilder<'config> {
19    fn build<'a, S: oak_core::source::Source + ?Sized>(&self, source: &S, edits: &[TextEdit], _cache: &'a mut impl BuilderCache<WatLanguage>) -> BuildOutput<WatLanguage> {
20        let parser = WatParser::new(self.config);
21        let mut session = oak_core::parser::session::ParseSession::<WatLanguage>::default();
22        let parse_result = parser.parse(source, edits, &mut session);
23
24        match parse_result.result {
25            Ok(green_tree) => {
26                let source_text = SourceText::new(source.get_text_in((0..source.length()).into()).into_owned());
27                match self.build_root(green_tree.clone(), &source_text) {
28                    Ok(ast_root) => OakDiagnostics { result: Ok(ast_root), diagnostics: parse_result.diagnostics },
29                    Err(build_error) => {
30                        let mut diagnostics = parse_result.diagnostics;
31                        diagnostics.push(build_error.clone());
32                        OakDiagnostics { result: Err(build_error), diagnostics }
33                    }
34                }
35            }
36            Err(parse_error) => OakDiagnostics { result: Err(parse_error), diagnostics: parse_result.diagnostics },
37        }
38    }
39}
40
41impl<'config> WatBuilder<'config> {
42    fn build_root(&self, green_tree: GreenNode<WatLanguage>, source: &SourceText) -> Result<WatRoot, OakError> {
43        let red_root = RedNode::new(&green_tree, 0);
44        let mut items = Vec::new();
45
46        for child in red_root.children() {
47            if let RedTree::Node(n) = child {
48                match n.green.kind {
49                    WatElementType::Module => {
50                        items.push(WatItem::Module(self.build_module(n, source)?));
51                    }
52                    _ => {}
53                }
54            }
55        }
56
57        Ok(WatRoot { items })
58    }
59
60    fn build_module(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatModule, OakError> {
61        let mut name = None;
62        let mut items = Vec::new();
63
64        for child in node.children() {
65            match child {
66                RedTree::Leaf(t) => {
67                    if t.kind == WatTokenType::Identifier {
68                        name = Some(source.get_text_in(t.span.clone()).to_string());
69                    }
70                }
71                RedTree::Node(n) => match n.green.kind {
72                    WatElementType::Func => {
73                        items.push(WatModuleField::Func(self.build_func(n, source)?));
74                    }
75                    WatElementType::Export => {
76                        items.push(WatModuleField::Export(self.build_export(n, source)?));
77                    }
78                    WatElementType::Import => {
79                        items.push(WatModuleField::Import(self.build_import(n, source)?));
80                    }
81                    WatElementType::Type => {
82                        items.push(WatModuleField::Type(self.build_type(n, source)?));
83                    }
84                    WatElementType::Table => {
85                        items.push(WatModuleField::Table(self.build_table(n, source)?));
86                    }
87                    WatElementType::Memory => {
88                        items.push(WatModuleField::Memory(self.build_memory(n, source)?));
89                    }
90                    WatElementType::Global => {
91                        items.push(WatModuleField::Global(self.build_global(n, source)?));
92                    }
93                    _ => {}
94                },
95            }
96        }
97
98        Ok(WatModule { name, items })
99    }
100
101    fn build_func(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatFunc, OakError> {
102        let mut name = None;
103        let mut params = Vec::new();
104        let mut results = Vec::new();
105        let mut locals = Vec::new();
106        let mut body = Vec::new();
107
108        for child in node.children() {
109            match child {
110                RedTree::Leaf(t) => {
111                    if t.kind == WatTokenType::Identifier {
112                        name = Some(source.get_text_in(t.span.clone().into()).to_string());
113                    }
114                }
115                RedTree::Node(n) => match n.green.kind {
116                    WatElementType::Param => {
117                        params.push(self.build_param(n, source)?);
118                    }
119                    WatElementType::Result => {
120                        results.push(self.build_result(n, source)?);
121                    }
122                    WatElementType::Local => {
123                        locals.push(self.build_local(n, source)?);
124                    }
125                    WatElementType::Instruction => {
126                        body.push(self.build_instruction(n, source)?);
127                    }
128                    _ => {}
129                },
130            }
131        }
132
133        Ok(WatFunc { name, params, results, locals, body })
134    }
135
136    fn build_param(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatParam, OakError> {
137        let mut name = None;
138        let mut ty = WatTypeKind::I32; // Default
139
140        for child in node.children() {
141            match child {
142                RedTree::Leaf(t) => {
143                    if t.kind == WatTokenType::Identifier {
144                        name = Some(source.get_text_in(t.span.clone()).to_string());
145                    }
146                    else if let Some(t_kind) = self.parse_type_kind(&source.get_text_in(t.span.clone())) {
147                        ty = t_kind;
148                    }
149                }
150                _ => {}
151            }
152        }
153
154        Ok(WatParam { name, ty })
155    }
156
157    fn build_result(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatResult, OakError> {
158        let mut ty = WatTypeKind::I32; // Default
159
160        for child in node.children() {
161            if let RedTree::Leaf(t) = child {
162                if let Some(t_kind) = self.parse_type_kind(&source.get_text_in(t.span.clone())) {
163                    ty = t_kind;
164                }
165            }
166        }
167
168        Ok(WatResult { ty })
169    }
170
171    fn build_local(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatLocal, OakError> {
172        let mut name = None;
173        let mut ty = WatTypeKind::I32; // Default
174
175        for child in node.children() {
176            match child {
177                RedTree::Leaf(t) => {
178                    if t.kind == WatTokenType::Identifier {
179                        name = Some(source.get_text_in(t.span.clone()).to_string());
180                    }
181                    else if let Some(t_kind) = self.parse_type_kind(&source.get_text_in(t.span.clone())) {
182                        ty = t_kind;
183                    }
184                }
185                _ => {}
186            }
187        }
188
189        Ok(WatLocal { name, ty })
190    }
191
192    fn build_instruction(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatInstruction, OakError> {
193        let mut opcode = String::new();
194        let mut args = Vec::new();
195
196        for child in node.children() {
197            match child {
198                RedTree::Leaf(t) => {
199                    let text = source.get_text_in(t.span.clone()).to_string();
200                    if text != "(" && text != ")" {
201                        if opcode.is_empty() {
202                            opcode = text;
203                        }
204                        else {
205                            args.push(text);
206                        }
207                    }
208                }
209                RedTree::Node(n) => {
210                    if n.green.kind == WatElementType::Instruction {
211                        // For nested instructions, we could handle them if WatInstruction was recursive.
212                        // For now, we'll just treat the inner instruction's opcode as an argument or similar.
213                        // Let's just flatten the text for now as a simple fallback.
214                        args.push(source.get_text_in(n.span().into()).to_string());
215                    }
216                }
217            }
218        }
219
220        match opcode.as_str() {
221            "unreachable" => Ok(WatInstruction::Unreachable),
222            "nop" => Ok(WatInstruction::Nop),
223            "drop" => Ok(WatInstruction::Drop),
224            "select" => Ok(WatInstruction::Select),
225            "return" => Ok(WatInstruction::Return),
226            "local.get" if !args.is_empty() => Ok(WatInstruction::LocalGet(args[0].clone())),
227            "local.set" if !args.is_empty() => Ok(WatInstruction::LocalSet(args[0].clone())),
228            "local.tee" if !args.is_empty() => Ok(WatInstruction::LocalTee(args[0].clone())),
229            "global.get" if !args.is_empty() => Ok(WatInstruction::GlobalGet(args[0].clone())),
230            "global.set" if !args.is_empty() => Ok(WatInstruction::GlobalSet(args[0].clone())),
231            "i32.const" if !args.is_empty() => Ok(WatInstruction::I32Const(args[0].parse().unwrap_or(0))),
232            "i64.const" if !args.is_empty() => Ok(WatInstruction::I64Const(args[0].parse().unwrap_or(0))),
233            "f32.const" if !args.is_empty() => Ok(WatInstruction::F32Const(args[0].parse().unwrap_or(0.0))),
234            "f64.const" if !args.is_empty() => Ok(WatInstruction::F64Const(args[0].parse().unwrap_or(0.0))),
235            "i32.add" => Ok(WatInstruction::I32Add),
236            "i32.sub" => Ok(WatInstruction::I32Sub),
237            "i32.mul" => Ok(WatInstruction::I32Mul),
238            "i64.add" => Ok(WatInstruction::I64Add),
239            "i64.sub" => Ok(WatInstruction::I64Sub),
240            "i64.mul" => Ok(WatInstruction::I64Mul),
241            _ => Ok(WatInstruction::Other(opcode, args)),
242        }
243    }
244
245    fn build_export(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatExport, OakError> {
246        let mut name = String::new();
247        let mut kind = String::new();
248        let mut id = String::new();
249
250        for child in node.children() {
251            if let RedTree::Leaf(t) = child {
252                let text = source.get_text_in(t.span.clone().into()).to_string();
253                if t.kind == WatTokenType::StringLiteral {
254                    name = text.trim_matches('"').to_string();
255                }
256                else if text == "(" || text == ")" || text == "export" {
257                    continue;
258                }
259                else if kind.is_empty() {
260                    kind = text;
261                }
262                else {
263                    id = text;
264                }
265            }
266        }
267
268        Ok(WatExport { name, kind, id })
269    }
270
271    fn build_import(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatImport, OakError> {
272        let mut module = String::new();
273        let mut name = String::new();
274        let mut kind = String::new();
275
276        for child in node.children() {
277            if let RedTree::Leaf(t) = child {
278                let text = source.get_text_in(t.span.clone().into()).to_string();
279                if t.kind == WatTokenType::StringLiteral {
280                    if module.is_empty() {
281                        module = text.trim_matches('"').to_string();
282                    }
283                    else {
284                        name = text.trim_matches('"').to_string();
285                    }
286                }
287                else if text == "(" || text == ")" || text == "import" {
288                    continue;
289                }
290                else if kind.is_empty() {
291                    kind = text;
292                }
293            }
294        }
295
296        Ok(WatImport { module, name, kind })
297    }
298
299    fn build_type(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatType, OakError> {
300        let mut id = None;
301        for child in node.children() {
302            if let RedTree::Leaf(t) = child {
303                if t.kind == WatTokenType::Identifier {
304                    id = Some(source.get_text_in(t.span.clone()).to_string());
305                }
306            }
307        }
308        Ok(WatType { id })
309    }
310
311    fn build_table(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatTable, OakError> {
312        let mut id = None;
313        for child in node.children() {
314            if let RedTree::Leaf(t) = child {
315                if t.kind == WatTokenType::Identifier {
316                    id = Some(source.get_text_in(t.span.clone()).to_string());
317                }
318            }
319        }
320        Ok(WatTable { id, span: node.span() })
321    }
322
323    fn build_memory(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatMemory, OakError> {
324        let mut id = None;
325        for child in node.children() {
326            if let RedTree::Leaf(t) = child {
327                if t.kind == WatTokenType::Identifier {
328                    id = Some(source.get_text_in(t.span.clone()).to_string());
329                }
330            }
331        }
332        Ok(WatMemory { id, span: node.span() })
333    }
334
335    fn build_global(&self, node: RedNode<WatLanguage>, source: &SourceText) -> Result<WatGlobal, OakError> {
336        let mut id = None;
337        let mut ty = WatTypeKind::I32;
338        let mut mutable = false;
339
340        for child in node.children() {
341            if let RedTree::Leaf(t) = child {
342                let text = source.get_text_in(t.span.clone()).to_string();
343                if t.kind == WatTokenType::Identifier {
344                    id = Some(text);
345                }
346                else if text == "mut" {
347                    mutable = true;
348                }
349                else if let Some(k) = self.parse_type_kind(&text) {
350                    ty = k;
351                }
352            }
353        }
354
355        Ok(WatGlobal { id, ty, mutable })
356    }
357
358    fn parse_type_kind(&self, text: &str) -> Option<WatTypeKind> {
359        match text {
360            "i32" => Some(WatTypeKind::I32),
361            "i64" => Some(WatTypeKind::I64),
362            "f32" => Some(WatTypeKind::F32),
363            "f64" => Some(WatTypeKind::F64),
364            _ => None,
365        }
366    }
367}