Skip to main content

wasm_wave/
ast.rs

1//! Abstract syntax tree types
2
3use core::str::FromStr;
4
5use alloc::{
6    borrow::Cow,
7    collections::BTreeMap,
8    string::{String, ToString},
9    vec::Vec,
10};
11
12use crate::{
13    lex::Span,
14    parser::{ParserError, ParserErrorKind},
15    strings::{StringPartsIter, unescape},
16    wasm::{WasmType, WasmTypeKind, WasmValue, WasmValueError},
17};
18
19/// A WAVE AST node.
20#[derive(Clone, Debug)]
21pub struct Node {
22    ty: NodeType,
23    span: Span,
24    children: Vec<Node>,
25}
26
27impl Node {
28    pub(crate) fn new(
29        ty: NodeType,
30        span: impl Into<Span>,
31        children: impl IntoIterator<Item = Node>,
32    ) -> Self {
33        Self {
34            ty,
35            span: span.into(),
36            children: Vec::from_iter(children),
37        }
38    }
39
40    /// Returns this node's type.
41    pub fn ty(&self) -> NodeType {
42        self.ty
43    }
44
45    /// Returns this node's span.
46    pub fn span(&self) -> Span {
47        self.span.clone()
48    }
49
50    /// Returns a bool value if this node represents a bool.
51    pub fn as_bool(&self) -> Result<bool, ParserError> {
52        match self.ty {
53            NodeType::BoolTrue => Ok(true),
54            NodeType::BoolFalse => Ok(false),
55            _ => Err(self.error(ParserErrorKind::InvalidType)),
56        }
57    }
58
59    /// Returns a number value of the given type (integer or float) if this node
60    /// can represent a number of that type.
61    pub fn as_number<T: FromStr>(&self, src: &str) -> Result<T, ParserError> {
62        self.ensure_type(NodeType::Number)?;
63        self.slice(src)
64            .parse()
65            .map_err(|_| self.error(ParserErrorKind::InvalidValue))
66    }
67
68    /// Returns a char value if this node represents a valid char.
69    pub fn as_char(&self, src: &str) -> Result<char, ParserError> {
70        self.ensure_type(NodeType::Char)?;
71        let inner = &src[self.span.start + 1..self.span.end - 1];
72        let (ch, len) = if inner.starts_with('\\') {
73            unescape(inner).ok_or_else(|| self.error(ParserErrorKind::InvalidEscape))?
74        } else {
75            let ch = inner.chars().next().unwrap();
76            (ch, ch.len_utf8())
77        };
78        // Verify length
79        if len != inner.len() {
80            return Err(self.error(ParserErrorKind::MultipleChars));
81        }
82        Ok(ch)
83    }
84
85    /// Returns a str value if this node represents a valid string.
86    pub fn as_str<'src>(&self, src: &'src str) -> Result<Cow<'src, str>, ParserError> {
87        let mut parts = self.iter_str(src)?;
88        let Some(first) = parts.next().transpose()? else {
89            return Ok("".into());
90        };
91        match parts.next().transpose()? {
92            // Single part may be borrowed
93            None => Ok(first),
94            // Multiple parts must be collected into a single owned String
95            Some(second) => {
96                let s: String = [Ok(first), Ok(second)]
97                    .into_iter()
98                    .chain(parts)
99                    .collect::<Result<_, _>>()?;
100                Ok(s.into())
101            }
102        }
103    }
104
105    /// Returns an iterator of string "parts" which together form a decoded
106    /// string value if this node represents a valid string.
107    pub fn iter_str<'src>(
108        &self,
109        src: &'src str,
110    ) -> Result<impl Iterator<Item = Result<Cow<'src, str>, ParserError>>, ParserError> {
111        match self.ty {
112            NodeType::String => {
113                let span = self.span.start + 1..self.span.end - 1;
114                Ok(StringPartsIter::new(&src[span.clone()], span.start))
115            }
116            NodeType::MultilineString => {
117                let span = self.span.start + 3..self.span.end - 3;
118                Ok(StringPartsIter::new_multiline(
119                    &src[span.clone()],
120                    span.start,
121                )?)
122            }
123            _ => Err(self.error(ParserErrorKind::InvalidType)),
124        }
125    }
126
127    /// Returns an iterator of value nodes if this node represents a tuple.
128    pub fn as_tuple(&self) -> Result<impl ExactSizeIterator<Item = &Node>, ParserError> {
129        self.ensure_type(NodeType::Tuple)?;
130        Ok(self.children.iter())
131    }
132
133    /// Returns an iterator of value nodes if this node represents a list.
134    pub fn as_list(&self) -> Result<impl ExactSizeIterator<Item = &Node>, ParserError> {
135        self.ensure_type(NodeType::List)?;
136        Ok(self.children.iter())
137    }
138
139    /// Returns an iterator of field name and value node pairs if this node
140    /// represents a record value.
141    pub fn as_record<'this, 'src>(
142        &'this self,
143        src: &'src str,
144    ) -> Result<impl ExactSizeIterator<Item = (&'src str, &'this Node)>, ParserError> {
145        self.ensure_type(NodeType::Record)?;
146        Ok(self
147            .children
148            .chunks(2)
149            .map(|chunk| (chunk[0].as_label(src).unwrap(), &chunk[1])))
150    }
151
152    /// Returns a variant label and optional payload if this node can represent
153    /// a variant value.
154    pub fn as_variant<'this, 'src>(
155        &'this self,
156        src: &'src str,
157    ) -> Result<(&'src str, Option<&'this Node>), ParserError> {
158        match self.ty {
159            NodeType::Label => Ok((self.as_label(src)?, None)),
160            NodeType::VariantWithPayload => {
161                let label = self.children[0].as_label(src)?;
162                let value = &self.children[1];
163                Ok((label, Some(value)))
164            }
165            _ => Err(self.error(ParserErrorKind::InvalidType)),
166        }
167    }
168
169    /// Returns an enum value label if this node represents a label.
170    pub fn as_enum<'src>(&self, src: &'src str) -> Result<&'src str, ParserError> {
171        self.as_label(src)
172    }
173
174    /// Returns an option value if this node represents an option.
175    pub fn as_option(&self) -> Result<Option<&Node>, ParserError> {
176        match self.ty {
177            NodeType::OptionSome => Ok(Some(&self.children[0])),
178            NodeType::OptionNone => Ok(None),
179            _ => Err(self.error(ParserErrorKind::InvalidType)),
180        }
181    }
182
183    /// Returns a result value with optional payload value if this node
184    /// represents a result.
185    pub fn as_result(&self) -> Result<Result<Option<&Node>, Option<&Node>>, ParserError> {
186        let payload = self.children.first();
187        match self.ty {
188            NodeType::ResultOk => Ok(Ok(payload)),
189            NodeType::ResultErr => Ok(Err(payload)),
190            _ => Err(self.error(ParserErrorKind::InvalidType)),
191        }
192    }
193
194    /// Returns an iterator of flag labels if this node represents flags.
195    pub fn as_flags<'this, 'src: 'this>(
196        &'this self,
197        src: &'src str,
198    ) -> Result<impl Iterator<Item = &'src str> + 'this, ParserError> {
199        self.ensure_type(NodeType::Flags)?;
200        Ok(self.children.iter().map(|node| node.as_label(src).unwrap()))
201    }
202
203    fn as_label<'src>(&self, src: &'src str) -> Result<&'src str, ParserError> {
204        self.ensure_type(NodeType::Label)?;
205        let label = self.slice(src);
206        let label = label.strip_prefix('%').unwrap_or(label);
207        Ok(label)
208    }
209
210    /// Converts this node into the given typed value from the given input source.
211    pub fn to_wasm_value<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
212        Ok(match ty.kind() {
213            WasmTypeKind::Bool => V::make_bool(self.as_bool()?),
214            WasmTypeKind::S8 => V::make_s8(self.as_number(src)?),
215            WasmTypeKind::S16 => V::make_s16(self.as_number(src)?),
216            WasmTypeKind::S32 => V::make_s32(self.as_number(src)?),
217            WasmTypeKind::S64 => V::make_s64(self.as_number(src)?),
218            WasmTypeKind::U8 => V::make_u8(self.as_number(src)?),
219            WasmTypeKind::U16 => V::make_u16(self.as_number(src)?),
220            WasmTypeKind::U32 => V::make_u32(self.as_number(src)?),
221            WasmTypeKind::U64 => V::make_u64(self.as_number(src)?),
222            WasmTypeKind::F32 => V::make_f32(self.as_number(src)?),
223            WasmTypeKind::F64 => V::make_f64(self.as_number(src)?),
224            WasmTypeKind::Char => V::make_char(self.as_char(src)?),
225            WasmTypeKind::String => V::make_string(self.as_str(src)?),
226            WasmTypeKind::List => self.to_wasm_list(ty, src)?,
227            WasmTypeKind::Record => self.to_wasm_record(ty, src)?,
228            WasmTypeKind::Tuple => self.to_wasm_tuple(ty, src)?,
229            WasmTypeKind::Variant => self.to_wasm_variant(ty, src)?,
230            WasmTypeKind::Enum => self.to_wasm_enum(ty, src)?,
231            WasmTypeKind::Option => self.to_wasm_option(ty, src)?,
232            WasmTypeKind::Result => self.to_wasm_result(ty, src)?,
233            WasmTypeKind::Flags => self.to_wasm_flags(ty, src)?,
234            other => {
235                return Err(
236                    self.wasm_value_error(WasmValueError::UnsupportedType(other.to_string()))
237                );
238            }
239        })
240    }
241
242    /// Converts this node into the given types.
243    /// See [`crate::untyped::UntypedFuncCall::to_wasm_params`].
244    pub fn to_wasm_params<'types, V: WasmValue + 'static>(
245        &self,
246        types: impl IntoIterator<Item = &'types V::Type>,
247        src: &str,
248    ) -> Result<Vec<V>, ParserError> {
249        let mut types = types.into_iter();
250        let mut values = self
251            .as_tuple()?
252            .map(|node| {
253                let ty = types.next().ok_or_else(|| {
254                    ParserError::with_detail(
255                        ParserErrorKind::InvalidParams,
256                        node.span().clone(),
257                        "more param(s) than expected",
258                    )
259                })?;
260                node.to_wasm_value::<V>(ty, src)
261            })
262            .collect::<Result<Vec<_>, _>>()?;
263        // Handle trailing optional fields
264        for ty in types {
265            if ty.kind() == WasmTypeKind::Option {
266                values.push(V::make_option(ty, None).map_err(|err| self.wasm_value_error(err))?);
267            } else {
268                return Err(ParserError::with_detail(
269                    ParserErrorKind::InvalidParams,
270                    self.span.end - 1..self.span.end,
271                    "missing required param(s)",
272                ));
273            }
274        }
275        Ok(values)
276    }
277
278    fn to_wasm_list<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
279        let element_type = ty.list_element_type().unwrap();
280        let elements = self
281            .as_list()?
282            .map(|node| node.to_wasm_value(&element_type, src))
283            .collect::<Result<Vec<_>, _>>()?;
284        V::make_list(ty, elements).map_err(|err| self.wasm_value_error(err))
285    }
286
287    fn to_wasm_record<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
288        let values = self.as_record(src)?.collect::<BTreeMap<_, _>>();
289        let record_fields = ty.record_fields().collect::<Vec<_>>();
290        let fields = record_fields
291            .iter()
292            .map(|(name, field_type)| {
293                let value = match (values.get(name.as_ref()), field_type.kind()) {
294                    (Some(node), _) => node.to_wasm_value(field_type, src)?,
295                    (None, WasmTypeKind::Option) => V::make_option(field_type, None)
296                        .map_err(|err| self.wasm_value_error(err))?,
297                    _ => {
298                        return Err(
299                            self.wasm_value_error(WasmValueError::MissingField(name.to_string()))
300                        );
301                    }
302                };
303                Ok((name.as_ref(), value))
304            })
305            .collect::<Result<Vec<_>, _>>()?;
306        V::make_record(ty, fields).map_err(|err| self.wasm_value_error(err))
307    }
308
309    fn to_wasm_tuple<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
310        let types = ty.tuple_element_types().collect::<Vec<_>>();
311        let values = self.as_tuple()?;
312        if types.len() != values.len() {
313            return Err(
314                self.wasm_value_error(WasmValueError::WrongNumberOfTupleValues {
315                    want: types.len(),
316                    got: values.len(),
317                }),
318            );
319        }
320        let values = ty
321            .tuple_element_types()
322            .zip(self.as_tuple()?)
323            .map(|(ty, node)| node.to_wasm_value(&ty, src))
324            .collect::<Result<Vec<_>, _>>()?;
325        V::make_tuple(ty, values).map_err(|err| self.wasm_value_error(err))
326    }
327
328    fn to_wasm_variant<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
329        let (label, payload) = self.as_variant(src)?;
330        let payload_type = ty
331            .variant_cases()
332            .find_map(|(case, payload)| (case == label).then_some(payload))
333            .ok_or_else(|| self.wasm_value_error(WasmValueError::UnknownCase(label.into())))?;
334        let payload_value = self.to_wasm_maybe_payload(label, &payload_type, payload, src)?;
335        V::make_variant(ty, label, payload_value).map_err(|err| self.wasm_value_error(err))
336    }
337
338    fn to_wasm_enum<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
339        V::make_enum(ty, self.as_enum(src)?).map_err(|err| self.wasm_value_error(err))
340    }
341
342    fn to_wasm_option<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
343        let payload_type = ty.option_some_type().unwrap();
344        let value = match self.ty {
345            NodeType::OptionSome => {
346                self.to_wasm_maybe_payload("some", &Some(payload_type), self.as_option()?, src)?
347            }
348            NodeType::OptionNone => {
349                self.to_wasm_maybe_payload("none", &None, self.as_option()?, src)?
350            }
351            _ if flattenable(payload_type.kind()) => Some(self.to_wasm_value(&payload_type, src)?),
352            _ => {
353                return Err(self.error(ParserErrorKind::InvalidType));
354            }
355        };
356        V::make_option(ty, value).map_err(|err| self.wasm_value_error(err))
357    }
358
359    fn to_wasm_result<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
360        let (ok_type, err_type) = ty.result_types().unwrap();
361        let value = match self.ty {
362            NodeType::ResultOk => {
363                Ok(self.to_wasm_maybe_payload("ok", &ok_type, self.as_result()?.unwrap(), src)?)
364            }
365            NodeType::ResultErr => Err(self.to_wasm_maybe_payload(
366                "err",
367                &err_type,
368                self.as_result()?.unwrap_err(),
369                src,
370            )?),
371            _ => match ok_type {
372                Some(ty) if flattenable(ty.kind()) => Ok(Some(self.to_wasm_value(&ty, src)?)),
373                _ => return Err(self.error(ParserErrorKind::InvalidType)),
374            },
375        };
376        V::make_result(ty, value).map_err(|err| self.wasm_value_error(err))
377    }
378
379    fn to_wasm_flags<V: WasmValue>(&self, ty: &V::Type, src: &str) -> Result<V, ParserError> {
380        // Sort flags in `WasmType::flags_names` order.
381        let names = ty.flags_names().collect::<Vec<_>>();
382        let mut flags = self.as_flags(src)?.collect::<Vec<_>>();
383        flags.sort_by_cached_key(|flag| names.iter().position(|name| name == flag));
384        V::make_flags(ty, flags).map_err(|err| self.wasm_value_error(err))
385    }
386
387    fn to_wasm_maybe_payload<V: WasmValue>(
388        &self,
389        case: &str,
390        payload_type: &Option<V::Type>,
391        payload: Option<&Node>,
392        src: &str,
393    ) -> Result<Option<V>, ParserError> {
394        match (payload_type.as_ref(), payload) {
395            (Some(ty), Some(node)) => Ok(Some(node.to_wasm_value(ty, src)?)),
396            (None, None) => Ok(None),
397            (Some(_), None) => {
398                Err(self.wasm_value_error(WasmValueError::MissingPayload(case.into())))
399            }
400            (None, Some(_)) => {
401                Err(self.wasm_value_error(WasmValueError::UnexpectedPayload(case.into())))
402            }
403        }
404    }
405
406    fn wasm_value_error(&self, err: WasmValueError) -> ParserError {
407        ParserError::with_source(ParserErrorKind::WasmValueError, self.span(), err)
408    }
409
410    pub(crate) fn slice<'src>(&self, src: &'src str) -> &'src str {
411        &src[self.span()]
412    }
413
414    fn ensure_type(&self, ty: NodeType) -> Result<(), ParserError> {
415        if self.ty == ty {
416            Ok(())
417        } else {
418            Err(self.error(ParserErrorKind::InvalidType))
419        }
420    }
421
422    fn error(&self, kind: ParserErrorKind) -> ParserError {
423        ParserError::new(kind, self.span())
424    }
425}
426
427fn flattenable(kind: WasmTypeKind) -> bool {
428    // TODO: Consider wither to allow flattening an option in a result or vice-versa.
429    !matches!(kind, WasmTypeKind::Option | WasmTypeKind::Result)
430}
431
432/// The type of a WAVE AST [`Node`].
433#[derive(Clone, Copy, Debug, PartialEq)]
434pub enum NodeType {
435    /// Boolean `true`
436    BoolTrue,
437    /// Boolean `false`
438    BoolFalse,
439    /// Number
440    /// May be an integer or float, including `nan`, `inf`, `-inf`
441    Number,
442    /// Char
443    /// Span includes delimiters.
444    Char,
445    /// String
446    /// Span includes delimiters.
447    String,
448    /// Multiline String
449    /// Span includes delimiters.
450    MultilineString,
451    /// Tuple
452    /// Child nodes are the tuple values.
453    Tuple,
454    /// List
455    /// Child nodes are the list values.
456    List,
457    /// Record
458    /// Child nodes are field Label, value pairs, e.g.
459    /// `[<field 1>, <value 1>, <field 2>, <value 2>, ...]`
460    Record,
461    /// Label
462    /// In value position may represent an enum value or variant case (without payload).
463    Label,
464    /// Variant case with payload
465    /// Child nodes are variant case Label and payload value.
466    VariantWithPayload,
467    /// Option `some`
468    /// Child node is the payload value.
469    OptionSome,
470    /// Option `none`
471    OptionNone,
472    /// Result `ok`
473    /// Has zero or one child node for the payload value.
474    ResultOk,
475    /// Result `err`
476    /// Has zero or one child node for the payload value.
477    ResultErr,
478    /// Flags
479    /// Child nodes are flag Labels.
480    Flags,
481}