pest_typed/
iterators.rs

1// pest-typed. A statically typed version of pest.
2// Copyright (c) 2023 黄博奕
3//
4// Licensed under the Apache License, Version 2.0
5// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
6// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
7// option. All files in the project carrying such notice may not be copied,
8// modified, or distributed except according to those terms.
9
10//! Simulates [`pest::iterators`].
11
12use crate::{
13    predefined_node::{
14        AlwaysFail, AtomicRepeat, CharRange, Empty, Insens, Negative, PeekSlice1, PeekSlice2,
15        Positive, Push, RepeatMin, RepeatMinMax, Skip, Skipped, Str, ANY, DROP, NEWLINE, PEEK,
16        PEEK_ALL, POP, POP_ALL, SOI,
17    },
18    typed_node::{RuleStorage, RuleStruct, Spanned},
19    RuleType, Span, StringArrayWrapper, StringWrapper, TypedNode,
20};
21use alloc::{boxed::Box, collections::VecDeque, string::String, vec::Vec};
22use core::{
23    iter::{once, Iterator},
24    mem::swap,
25};
26#[cfg(feature = "serde")]
27use serde::ser::SerializeStruct;
28
29/// Token.
30#[derive(Clone, Debug, Eq, PartialEq)]
31pub struct ThinToken<R: RuleType> {
32    /// Rule.
33    pub rule: R,
34    /// Start position.
35    pub start: usize,
36    /// End position.
37    pub end: usize,
38    /// Children.
39    pub children: Vec<Self>,
40}
41
42/// Token.
43#[derive(Clone, Debug, Eq, PartialEq)]
44pub struct Token<'i, R: RuleType> {
45    /// Rule.
46    pub rule: R,
47    /// Span.
48    pub span: Span<'i>,
49    /// Children.
50    pub children: Vec<Self>,
51}
52
53impl<R: RuleType> Token<'_, R> {
54    /// To [`ThinToken`].
55    pub fn to_thin(&self) -> ThinToken<R> {
56        let rule = self.rule;
57        let start = self.span.start();
58        let end = self.span.end();
59        let children = self.children.iter().map(|c| c.to_thin()).collect();
60        ThinToken {
61            rule,
62            start,
63            end,
64            children,
65        }
66    }
67}
68
69#[cfg(feature = "serde")]
70impl<'i, R: RuleType> serde::Serialize for Token<'i, R> {
71    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
72    where
73        S: serde::Serializer,
74    {
75        let mut seq = serializer.serialize_struct(self.rule.name(), 3)?;
76
77        seq.serialize_field("type", self.rule.name())?;
78        seq.serialize_field("content", self.span.as_str())?;
79        seq.serialize_field("children", &self.children)?;
80
81        seq.end()
82    }
83}
84
85/// Simulate [`pest::iterators::Pairs`].
86pub trait Pairs<'i, R: RuleType> {
87    /// For each inner pair by reference if this is a container, otherwise for self.
88    fn for_self_or_each_child(&self, f: &mut impl FnMut(Token<'i, R>));
89    /// Collect inner pairs and make them into a [`Vec`].
90    fn self_or_children(&self) -> Vec<Token<'i, R>> {
91        let mut children = Vec::new();
92        self.for_self_or_each_child(&mut |token| children.push(token));
93        children
94    }
95}
96
97/// Simulate [`pest::iterators::Pair`].
98pub trait Pair<'i, R: RuleType>: Spanned<'i, R> + RuleStorage<R> {
99    /// Iterate over all inner children.
100    fn for_each_child(&self, f: impl FnMut(Token<'i, R>));
101    /// Collect inner [Token]s and make them into a [`Vec`].
102    fn children(&self) -> Vec<Token<'i, R>> {
103        let mut children = Vec::new();
104        self.for_each_child(|token| children.push(token));
105        children
106    }
107    /// As [Token]s. Call [`Pair::children`] inside.
108    fn as_token(&self) -> Token<'i, R> {
109        let children = self.children();
110        Token::<R> {
111            rule: self.rule(),
112            span: self.span(),
113            children,
114        }
115    }
116    /// As [Token]s. Call [`Pair::children`] inside.
117    fn as_thin_token(&self) -> ThinToken<R> {
118        self.as_token().to_thin()
119    }
120}
121
122fn iterate_level_order<'i, R: RuleType, E>(
123    p: &impl Pair<'i, R>,
124    mut f: impl FnMut(&Token<'i, R>, usize) -> Result<(), E>,
125) -> Result<(), E> {
126    let mut queue: VecDeque<Token<'i, R>> = VecDeque::new();
127    let mut next = VecDeque::new();
128    queue.push_back(p.as_token());
129    loop {
130        while let Some(p) = queue.pop_front() {
131            f(&p, queue.len())?;
132            next.extend(p.children);
133        }
134        swap(&mut queue, &mut next);
135        if queue.is_empty() {
136            return Ok(());
137        }
138    }
139}
140/// Pre-order traversal
141fn iterate_pre_order<'i, R: RuleType, E>(
142    p: &impl Pair<'i, R>,
143    mut f: impl FnMut(&Token<'i, R>, usize) -> Result<(), E>,
144) -> Result<(), E> {
145    let mut stack: Vec<VecDeque<Token<'i, R>>> = Vec::new();
146
147    let root: Token<'i, R> = p.as_token();
148    stack.push(VecDeque::<Token<'i, R>>::from_iter(once(root)));
149
150    loop {
151        if let Some(parent) = stack.last_mut() {
152            if let Some(first) = parent.pop_front() {
153                f(&first, stack.len() - 1)?;
154                stack.push(first.children.into());
155            } else {
156                stack.pop();
157            }
158        } else {
159            return Ok(());
160        }
161    }
162}
163
164/// Write the tree to.
165fn write_tree_to<'i: 'n, 'n, R: RuleType + 'n>(
166    p: &'n impl Pair<'i, R>,
167    buf: &mut impl core::fmt::Write,
168) -> core::fmt::Result {
169    iterate_pre_order(p, |p, depth| {
170        if p.children.is_empty() {
171            buf.write_fmt(format_args!(
172                "{}{:?} {:?}\n",
173                &"    ".repeat(depth),
174                p.rule,
175                p.span.as_str(),
176            ))
177        } else {
178            buf.write_fmt(format_args!("{}{:?}\n", &"    ".repeat(depth), p.rule))
179        }
180    })
181}
182
183/// A trait to traverse the pair as the root of a tree.
184pub trait PairTree<'i, R: RuleType>: Pair<'i, R> + Sized {
185    /// Level order traversal
186    fn iterate_level_order<E>(
187        &self,
188        f: impl FnMut(&Token<'i, R>, usize) -> Result<(), E>,
189    ) -> Result<(), E> {
190        iterate_level_order(self, f)
191    }
192    /// Pre-order traversal
193    fn iterate_pre_order<E>(
194        &self,
195        f: impl FnMut(&Token<'i, R>, usize) -> Result<(), E>,
196    ) -> Result<(), E> {
197        iterate_pre_order(self, f)
198    }
199
200    /// Write the tree to the `buf`.
201    fn write_tree_to(&self, buf: &mut impl core::fmt::Write) -> core::fmt::Result {
202        write_tree_to(self, buf)
203    }
204
205    /// Format as a tree.
206    fn format_as_tree(&self) -> Result<String, core::fmt::Error> {
207        let mut buf = String::new();
208        self.write_tree_to(&mut buf)?;
209        Ok(buf)
210    }
211}
212
213impl<'i, R: RuleType, T: RuleStruct<'i, R> + Pairs<'i, R> + Pair<'i, R>> PairTree<'i, R> for T {}
214
215macro_rules! impl_empty {
216    ($node:ty $(, $($tt:tt)*)?) => {
217        impl<'i, R: RuleType $(, $($tt)*)?> Pairs<'i, R> for $node {
218            fn for_self_or_each_child(&self, _f: &mut impl FnMut(Token<'i, R>)) {}
219        }
220    };
221}
222
223macro_rules! impl_forward_inner {
224    ($node:ident) => {
225        impl<'i, R: RuleType, T: TypedNode<'i, R> + Pairs<'i, R>> Pairs<'i, R> for $node<T> {
226            fn for_self_or_each_child(&self, f: &mut impl FnMut(Token<'i, R>)) {
227                self.content.for_self_or_each_child(f)
228            }
229        }
230    };
231}
232
233impl_empty!(Str<T>, T: StringWrapper);
234impl_empty!(Insens<'i, T>, T: StringWrapper);
235impl_empty!(PeekSlice2<START, END>, const START: i32, const END: i32);
236impl_empty!(PeekSlice1<START>, const START: i32);
237impl_forward_inner!(Push);
238impl_empty!(Skip<'i, Strings>, Strings: StringArrayWrapper);
239impl_empty!(CharRange<MIN, MAX>, const MIN: char, const MAX: char);
240impl_empty!(Positive<T>, T: TypedNode<'i, R>);
241impl_empty!(Negative<T>, T: TypedNode<'i, R>);
242
243impl<'i, R: RuleType, T1: Pairs<'i, R>, T2: Pairs<'i, R>> Pairs<'i, R> for (T1, T2) {
244    fn for_self_or_each_child(&self, f: &mut impl FnMut(Token<'i, R>)) {
245        self.0.for_self_or_each_child(f);
246        self.1.for_self_or_each_child(f);
247    }
248}
249
250impl<'i, R: RuleType, T: Pairs<'i, R>, const N: usize> Pairs<'i, R> for [T; N] {
251    fn for_self_or_each_child(&self, f: &mut impl FnMut(Token<'i, R>)) {
252        self.as_slice()
253            .iter()
254            .for_each(|n| n.for_self_or_each_child(f))
255    }
256}
257
258impl<'i, R: RuleType, T: TypedNode<'i, R> + Pairs<'i, R>> Pairs<'i, R> for Box<T> {
259    fn for_self_or_each_child(&self, f: &mut impl FnMut(Token<'i, R>)) {
260        self.as_ref().for_self_or_each_child(f)
261    }
262}
263
264impl<'i, R: RuleType, T: TypedNode<'i, R> + Pairs<'i, R>> Pairs<'i, R> for Option<T> {
265    fn for_self_or_each_child(&self, f: &mut impl FnMut(Token<'i, R>)) {
266        if let Some(node) = self {
267            node.for_self_or_each_child(f)
268        }
269    }
270}
271
272impl<'i, R: RuleType, T: Pairs<'i, R>, Skip: Pairs<'i, R>, const SKIP: usize> Pairs<'i, R>
273    for Skipped<T, Skip, SKIP>
274{
275    fn for_self_or_each_child(&self, f: &mut impl FnMut(Token<'i, R>)) {
276        self.skipped.for_self_or_each_child(f);
277        self.matched.for_self_or_each_child(f);
278    }
279}
280
281macro_rules! impl_with_vec {
282    ($name:ident, $(const $args:ident : $t:ty,)*) => {
283        impl<
284                'i,
285                R: RuleType,
286                T: Pairs<'i, R>,
287                $(const $args: $t, )*
288            > Pairs<'i, R> for $name<T, $($args, )*>
289        {
290            fn for_self_or_each_child(&self, f: &mut impl FnMut(Token<'i, R>)) {
291                self.content.iter().for_each(|n| n.for_self_or_each_child(f));
292            }
293        }
294    };
295}
296
297impl_with_vec!(RepeatMinMax, const MIN: usize, const MAX: usize,);
298impl_with_vec!(RepeatMin, const MIN: usize,);
299impl_with_vec!(AtomicRepeat,);
300
301impl_empty!(ANY);
302impl_empty!(SOI);
303impl_empty!(NEWLINE);
304impl_empty!(PEEK<'i>);
305impl_empty!(PEEK_ALL<'i>);
306impl_empty!(POP<'i>);
307impl_empty!(POP_ALL<'i>);
308impl_empty!(DROP);
309
310impl_empty!(AlwaysFail<'i>);
311impl_empty!(Empty<'i>);
312
313/// An iterator that maybe contain another iterator.
314pub struct Maybe<Item, T: Iterator<Item = Item>>(Option<T>);
315impl<Item, T: Iterator<Item = Item>> Iterator for Maybe<Item, T> {
316    type Item = Item;
317
318    fn next(&mut self) -> Option<Self::Item> {
319        match &mut self.0 {
320            Some(inner) => inner.next(),
321            None => None,
322        }
323    }
324}