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