proclet/
op.rs

1use crate::{
2    prelude::*, Error, IntoTokens, Match, Punct, PunctExt, Span, SpanExt, ToTokenStream,
3    TokenTreeExt,
4};
5use std::{borrow::Cow, iter::FusedIterator, marker::PhantomData, mem};
6
7#[cfg(doc)]
8use crate::Parser;
9
10/// Function for OpParser to use to match operators.
11pub trait MatchOpFn: Clone + Fn(&str, Option<char>) -> Match<Cow<'static, str>> {}
12impl<T> MatchOpFn for T where T: Clone + Fn(&str, Option<char>) -> Match<Cow<'static, str>> {}
13
14/// Convenience function for calling [`Op::new_static`]. It can be used for parsing a specific op.
15pub const fn op<S: Span>(op: &'static str) -> Op<S> {
16    Op::new_static(op)
17}
18
19/// An operator. These can be parsed from tokens, and can also be used as parsers
20/// to parse one specific operator.
21#[derive(Clone, Debug)]
22pub struct Op<S: Span> {
23    str: Cow<'static, str>,
24    spans: Option<Box<[S]>>,
25}
26
27impl<S: Span> Op<S> {
28    /// Create a new operator. `str` can't be empty.
29    #[inline]
30    pub fn new(str: impl Into<Cow<'static, str>>) -> Self {
31        let str = str.into();
32        assert!(!str.is_empty());
33        Self { str, spans: None }
34    }
35
36    /// Create a new operator (const). `str` can't be empty.
37    #[inline]
38    pub const fn new_static(str: &'static str) -> Self {
39        assert!(!str.is_empty());
40        Self {
41            str: Cow::Borrowed(str),
42            spans: None,
43        }
44    }
45
46    /// Create a new operator with one specific span used for the whole operator.
47    /// `str` can't be empty.
48    #[inline]
49    pub fn with_span(str: impl Into<Cow<'static, str>>, span: S) -> Self {
50        let str = str.into();
51        assert!(!str.is_empty());
52        let spans = Some(vec![span; str.chars().count()].into_boxed_slice());
53        Self { str, spans }
54    }
55
56    /// Create a new operator with a specific span for each character of the op.
57    /// The length of `spans` must be the same as the number of characters in `str`.
58    /// `str` can't be empty.
59    #[inline]
60    pub fn with_spans(str: impl Into<Cow<'static, str>>, spans: Vec<S>) -> Self {
61        let str = str.into();
62        assert!(!str.is_empty());
63        let spans = spans.into_boxed_slice();
64        assert_eq!(str.chars().count(), spans.len());
65        Self {
66            str,
67            spans: Some(spans),
68        }
69    }
70
71    /// Get this op as a string.
72    #[inline]
73    pub fn as_str(&self) -> &str {
74        &self.str
75    }
76
77    /// Get an iterator over the `Punct`s that make up this op.
78    #[inline]
79    pub fn puncts(&self) -> Puncts<S::Punct> {
80        Puncts::new(self.as_str(), self.spans.as_deref())
81    }
82
83    #[inline]
84    fn alloced_spans(&mut self) -> &mut [S] {
85        if self.spans.is_none() {
86            self.spans = Some(vec![S::call_site(); self.str.chars().count()].into_boxed_slice());
87        }
88        self.spans.as_mut().unwrap()
89    }
90
91    /// Get the spans of this op, or `None` if the default span is used.
92    #[inline]
93    pub fn spans(&self) -> Option<&[S]> {
94        self.spans.as_ref().map(|s| s as _)
95    }
96
97    /// Set the spans of this op. The lengths of `spans` must be the same as the number of
98    /// characters in the string representation of this op.
99    #[inline]
100    pub fn set_spans(&mut self, spans: &[S]) {
101        // rust-analyzer incorrectly marks this code as an error without this mut
102        #[allow(unused_mut)]
103        for (mut s, span) in self.alloced_spans().iter_mut().zip(spans.iter().cycle()) {
104            *s = *span;
105        }
106    }
107
108    /// Get the span of this op. If the op has more than one span, this will currently
109    /// return the first span. This may change in the future if proc-macro* exposes the
110    /// ability to merge spans.
111    #[inline]
112    pub fn span(&self) -> Option<S> {
113        self.spans.as_ref().map(|s| s[0])
114    }
115
116    /// Set a single span for this op.
117    #[inline]
118    pub fn set_span(&mut self, span: S) {
119        // rust-analyzer incorrectly marks this code as an error without this mut
120        #[allow(unused_mut)]
121        for mut s in self.alloced_spans().iter_mut() {
122            *s = span;
123        }
124    }
125}
126
127impl<S: SpanExt> Op<S> {
128    /// Split this op using the provided [`OpParser`].
129    #[inline]
130    pub fn split(
131        &self,
132        parser: &OpParser<S::Punct, impl MatchOpFn>,
133    ) -> ParseOps<S, Puncts<S::Punct>, impl MatchOpFn> {
134        parser.parse_ops(self.puncts())
135    }
136}
137
138impl<T: crate::TokenTreeExt> crate::Parser<T> for Op<T::Span> {
139    type Output<'p, 'b> = Op<T::Span> where Self: 'p;
140
141    #[inline]
142    fn parse<'p, 'b>(
143        &'p self,
144        buf: &mut &'b crate::TokenBuf<T>,
145    ) -> Result<Self::Output<'p, 'b>, Error<T::Span>> {
146        OpParser::<T::Punct, _>::new(|str, next| {
147            if self.str == str {
148                Match::Complete(self.str.clone())
149            } else if let Some(rest) = self.str.strip_prefix(str) {
150                if rest.chars().next() == next {
151                    Match::NeedMore
152                } else {
153                    Match::NoMatch
154                }
155            } else {
156                Match::NoMatch
157            }
158        })
159        .parse(buf)
160    }
161}
162
163impl<T: TokenTreeExt> IntoTokens<T> for Op<T::Span> {
164    #[inline]
165    fn into_tokens(self) -> impl Iterator<Item = T> {
166        self.into_iter().flat_map(|p| p.into_tokens())
167    }
168}
169
170impl<S: SpanExt> ToTokenStream<S::TokenStream> for Op<S> {
171    #[inline]
172    fn extend_token_stream(&self, token_stream: &mut S::TokenStream) {
173        token_stream.extend(self.puncts().map(S::TokenTree::from))
174    }
175}
176
177impl<S: Span> From<&'static str> for Op<S> {
178    #[inline]
179    fn from(value: &'static str) -> Self {
180        Self::new(value)
181    }
182}
183
184impl<S: SpanExt> crate::Parse<S::TokenTree> for Op<S> {
185    /// Generic op parser. This doesn't check against valid ops.
186    #[inline]
187    fn parse(buf: &mut &crate::TokenBuf<S::TokenTree>) -> Result<Self, Error<S>> {
188        let mut str = String::new();
189        let mut spans = Vec::new();
190        buf.parse_prefix(|token| {
191            if let Some(punct) = token.punct() {
192                str.push(punct.as_char());
193                spans.push(punct.span());
194                if punct.spacing().is_joint() {
195                    Match::Partial((str.len(), spans.len()))
196                } else {
197                    Match::Complete((str.len(), spans.len()))
198                }
199            } else {
200                Match::NoMatch
201            }
202        })
203        .map(|(strlen, spanslen)| {
204            str.truncate(strlen);
205            spans.truncate(spanslen);
206            Op::with_spans(str, spans)
207        })
208        .map_err(|mut e| {
209            e.set_message("expected operator");
210            e
211        })
212    }
213}
214
215impl<S: SpanExt> IntoIterator for Op<S> {
216    type Item = S::Punct;
217    type IntoIter = Puncts<'static, S::Punct>;
218
219    #[inline]
220    fn into_iter(self) -> Self::IntoIter {
221        Puncts::new(self.str, self.spans.map(|s| s.into_vec()))
222    }
223}
224
225#[derive(Clone)]
226pub struct ParseOps<S: SpanExt, I: Iterator<Item = S::Punct>, F: MatchOpFn>(
227    I,
228    OpParserInstance<S::Punct, F>,
229);
230
231impl<S: SpanExt, I: Iterator<Item = S::Punct>, F: MatchOpFn> FusedIterator for ParseOps<S, I, F> {}
232
233impl<S: SpanExt, I: Iterator<Item = S::Punct>, F: MatchOpFn> Iterator for ParseOps<S, I, F> {
234    type Item = Result<Op<S>, Error<S>>;
235
236    #[inline]
237    fn next(&mut self) -> Option<Self::Item> {
238        for punct in self.0.by_ref() {
239            if let Some(x) = self.1.apply(punct) {
240                return Some(x);
241            }
242        }
243        self.1.finish()
244    }
245}
246
247/// Iterator over `Punct`s.
248#[derive(Clone, Debug)]
249pub struct Puncts<'a, P: Punct> {
250    str: Cow<'a, str>,
251    spans: Option<Cow<'a, [P::Span]>>,
252    stri: usize,
253    spansi: usize,
254}
255
256impl<'a, P: Punct> Puncts<'a, P> {
257    /// Create a new iterator that will generate `Punct`s from the provided characters and spans.
258    /// Every `Punct` except for the last one will have `Joint` spacing, and the last one will be `Alone`.
259    /// If `spans` is `None`, the puncts will use the default span.
260    #[inline]
261    pub fn new(str: impl Into<Cow<'a, str>>, spans: Option<impl Into<Cow<'a, [P::Span]>>>) -> Self {
262        Self {
263            str: str.into(),
264            spans: spans.map(|s| s.into()),
265            stri: 0,
266            spansi: 0,
267        }
268    }
269}
270
271impl<'a, P: PunctExt> FusedIterator for Puncts<'a, P> where Puncts<'a, P>: Iterator {}
272
273impl<'a, P: PunctExt> Iterator for Puncts<'a, P> {
274    type Item = P;
275
276    #[inline]
277    fn next(&mut self) -> Option<Self::Item> {
278        if let Some(ch) = self.str[self.stri..].chars().next() {
279            let span = if let Some(spans) = &self.spans {
280                spans[self.spansi]
281            } else {
282                P::Span::call_site()
283            };
284            self.stri += ch.len_utf8();
285            self.spansi += 1;
286            Some(P::with_span(
287                ch,
288                if self.str.len() > self.stri {
289                    P::Spacing::Joint
290                } else {
291                    P::Spacing::Alone
292                },
293                span,
294            ))
295        } else {
296            None
297        }
298    }
299}
300
301/// Parser for specific sets of operators. Ops can be parsed via [`Parser`] or manually from `Punct`s.
302#[derive(Clone)]
303pub struct OpParser<P: PunctExt, F: MatchOpFn>(F, PhantomData<fn() -> P>);
304
305impl<P: PunctExt, F: MatchOpFn> OpParser<P, F> {
306    /// Create a new `OpParser` using `match_op` to define the valid operators.
307    /// Usually you wouldn't call this directly, but create the `OpParser` using
308    /// a function generated by the `define_ops` macro in the `proclet-utils` crate.
309    /// If you want to parse a single Op using [`Parser`], use [`Op`] instead.
310    ///
311    /// See [`OpParser::match_op`] for details of how `match_op` should work.
312    #[inline]
313    pub const fn new(match_op: F) -> Self {
314        Self(match_op, PhantomData)
315    }
316
317    /// Create a new `OpParserInstance` for this `OpParser`. You only need this if you
318    /// don't use `Parser::parse` or `parse_ops`.
319    #[inline]
320    pub fn create(&self) -> OpParserInstance<P, F> {
321        OpParserInstance::new(self.0.clone())
322    }
323
324    /// Use this `OpParser` to parse operators from an iterator of `Punct`s.
325    #[inline]
326    pub fn parse_ops<I: Iterator<Item = P>>(&self, puncts: I) -> ParseOps<P::Span, I, F> {
327        ParseOps(puncts, self.create())
328    }
329
330    /// Check if `str` is a valid op.
331    ///
332    /// For example, if `+` and `+=` are valid ops and `+-` is invalid:
333    ///
334    /// - `match_op("+", Some('='))` returns `Match::Partial("+")`
335    /// - `match_op("+", Some('-'))` returns `Match::Complete("+")`
336    /// - `match_op("+", None)` returns `Match::Complete("+")`
337    /// - `match_op("+=", None)` returns `Match::Complete("+=")`
338    /// - `match_op("+-", None)` returns `Match::None`
339    ///
340    /// If `-=` is a valid op, and `-` and `-+` are invalid:
341    ///
342    /// - `match_op("-", Some('='))` returns `Match::NeedMore`
343    /// - `match_op("-", Some('+'))` returns `Match::None`
344    /// - `match_op("-", None)` returns `Match::None`
345    /// - `match_op("-=", None)` returns `Match::Complete("-=")`
346    #[inline]
347    pub fn match_op(&self, str: &str, next: Option<char>) -> Match<Cow<'static, str>> {
348        (self.0)(str, next)
349    }
350}
351
352impl<P: PunctExt, F: MatchOpFn> crate::Parser<P::TokenTree> for OpParser<P, F> {
353    type Output<'p, 'b> = Op<P::Span> where Self:'p;
354
355    #[inline]
356    fn parse<'p, 'b>(
357        &'p self,
358        buf: &mut &'b crate::TokenBuf<P::TokenTree>,
359    ) -> Result<Self::Output<'p, 'b>, Error<P::Span>> {
360        let mut string = String::new();
361        let mut spans = Vec::new();
362        buf.parse_prefix_next(move |token, next| {
363            if let Some(punct) = token.punct() {
364                let next = if punct.spacing().is_joint() {
365                    next.and_then(|next| next.punct().map(|next| next.as_char()))
366                } else {
367                    None
368                };
369                string.push(punct.as_char());
370                spans.push(punct.span());
371
372                match self.match_op(&string, next) {
373                    Match::Complete(str) => {
374                        string.clear();
375                        let op = Op::with_spans(str, mem::take(&mut spans));
376                        Match::Complete(op)
377                    }
378                    Match::Partial(_) | Match::NeedMore => Match::NeedMore,
379                    Match::NoMatch => Match::NoMatch,
380                }
381            } else {
382                Match::NoMatch
383            }
384        })
385        .map_err(|mut e| {
386            e.set_message("expected operator");
387            e
388        })
389    }
390}
391
392/// An instance for doing manual parsing of operators.
393#[derive(Clone)]
394pub struct OpParserInstance<P: PunctExt, F> {
395    str: String,
396    spans: Vec<P::Span>,
397    puncts: Vec<P>,
398    next: Option<P>,
399    match_op: F,
400}
401
402impl<P: PunctExt, F: MatchOpFn> OpParserInstance<P, F> {
403    /// Create a new `OpParserInstance`. Usually you'd do this by calling [`OpParser::create`] instead.
404    #[inline]
405    const fn new(match_op: F) -> Self {
406        Self {
407            str: String::new(),
408            spans: Vec::new(),
409            puncts: Vec::new(),
410            next: None,
411            match_op,
412        }
413    }
414
415    /// Clear the state to remove any accumulated op.
416    #[inline]
417    pub fn clear(&mut self) {
418        self.str.clear();
419        self.spans.clear();
420        self.puncts.clear();
421        self.next = None;
422    }
423
424    /// Add a `Punct` to the currently accumulating op. This may return a new [`Op`], or `None`
425    /// if more puncts are needed. If `punct` has alone spacing and the accumulated op isn't
426    /// valid, it returns an error. Call [`OpParserInstance::finish`] to finish parsing.
427    #[inline]
428    #[allow(clippy::type_complexity)]
429    pub fn apply(&mut self, punct: P) -> Option<Result<Op<P::Span>, Error<P::Span>>> {
430        let (mut punct, mut next_ch) = if let Some(next) = mem::take(&mut self.next) {
431            let next_ch = next.spacing().is_joint().then_some(punct.as_char());
432            self.next = Some(punct);
433            (next, next_ch)
434        } else if punct.spacing().is_alone() {
435            (punct, None)
436        } else {
437            self.next = Some(punct);
438            return None;
439        };
440
441        loop {
442            self.str.push(punct.as_char());
443            self.spans.push(punct.span());
444            self.puncts.push(punct);
445
446            match (self.match_op)(&self.str, next_ch) {
447                Match::Complete(str) => {
448                    self.str.clear();
449                    self.puncts.clear();
450                    return Some(Ok(Op::with_spans(str, mem::take(&mut self.spans))));
451                }
452                Match::Partial(_) | Match::NeedMore => {
453                    if self.next.as_ref().unwrap().spacing().is_alone() {
454                        punct = mem::take(&mut self.next).unwrap();
455                        next_ch = None;
456                        continue;
457                    } else {
458                        return None;
459                    }
460                }
461                Match::NoMatch => {
462                    let err = Error::with_span(*self.spans.first().unwrap(), "invalid op");
463                    self.str.clear();
464                    self.spans.clear();
465                    self.puncts.clear();
466                    return Some(Err(err));
467                }
468            }
469        }
470    }
471
472    /// Finish parsing the currently accumulated op.
473    #[inline]
474    #[allow(clippy::type_complexity)]
475    pub fn finish(&mut self) -> Option<Result<Op<P::Span>, Error<P::Span>>> {
476        mem::take(&mut self.next).map(|punct| {
477            self.str.push(punct.as_char());
478            let m = (self.match_op)(&self.str, None);
479            self.str.clear();
480            if let Match::Complete(str) | Match::Partial(str) = m {
481                self.spans.push(punct.span());
482                self.puncts.clear();
483                Ok(Op::with_spans(str, mem::take(&mut self.spans)))
484            } else {
485                let err = Error::with_span(*self.spans.first().unwrap(), "invalid op");
486                self.spans.clear();
487                self.puncts.clear();
488                Err(err)
489            }
490        })
491    }
492}