Skip to main content

virtue_next/parse/
generics.rs

1use super::utils::assume_ident;
2use super::utils::assume_punct;
3use super::utils::consume_punct_if;
4use super::utils::ident_eq;
5use super::utils::read_tokens_until_punct;
6use crate::Error;
7use crate::Result;
8use crate::generate::StreamBuilder;
9use crate::prelude::Ident;
10use crate::prelude::TokenTree;
11use std::iter::Peekable;
12use std::ops::Deref;
13use std::ops::DerefMut;
14
15/// A generic parameter for a struct or enum.
16///
17/// ```
18/// use std::fmt::Display;
19/// use std::marker::PhantomData;
20///
21/// // Generic will be `Generic::Generic("F")`
22/// struct Foo<F> {
23///     f: PhantomData<F>,
24/// }
25/// // Generics will be `Generic::Generic("F: Display")`
26/// struct Bar<F: Display> {
27///     f: PhantomData<F>,
28/// }
29/// // Generics will be `[Generic::Lifetime("a"), Generic::Generic("F: Display")]`
30/// struct Baz<'a, F> {
31///     f: PhantomData<&'a F>,
32/// }
33/// ```
34#[derive(Debug, Clone)]
35pub struct Generics(pub Vec<Generic>);
36
37impl Generics {
38    pub(crate) fn try_take(
39        input: &mut Peekable<impl Iterator<Item = TokenTree>>
40    ) -> Result<Option<Self>> {
41        let maybe_punct = input.peek();
42        if let Some(TokenTree::Punct(punct)) = maybe_punct
43            && punct.as_char() == '<'
44        {
45            let punct = assume_punct(input.next(), '<');
46            let mut result = Self(Vec::new());
47            loop {
48                match input.peek() {
49                    | Some(TokenTree::Punct(punct)) if punct.as_char() == '\'' => {
50                        result.push(Lifetime::take(input)?.into());
51                        consume_punct_if(input, ',');
52                    },
53                    | Some(TokenTree::Punct(punct)) if punct.as_char() == '>' => {
54                        assume_punct(input.next(), '>');
55                        break;
56                    },
57                    | Some(TokenTree::Ident(ident)) if ident_eq(ident, "const") => {
58                        result.push(ConstGeneric::take(input)?.into());
59                        consume_punct_if(input, ',');
60                    },
61                    | Some(TokenTree::Ident(_)) => {
62                        result.push(SimpleGeneric::take(input)?.into());
63                        consume_punct_if(input, ',');
64                    },
65                    | x => {
66                        return Err(Error::InvalidRustSyntax {
67                            span: x.map_or_else(|| punct.span(), TokenTree::span),
68                            expected: format!("', > or an ident, got {x:?}"),
69                        });
70                    },
71                }
72            }
73            return Ok(Some(result));
74        }
75        Ok(None)
76    }
77
78    /// Returns `true` if any of the generics is a [`Generic::Lifetime`]
79    #[must_use]
80    pub fn has_lifetime(&self) -> bool {
81        self.iter().any(Generic::is_lifetime)
82    }
83
84    /// Returns an iterator which contains only the simple type generics
85    pub fn iter_generics(&self) -> impl Iterator<Item = &SimpleGeneric> {
86        self.iter().filter_map(|g| {
87            match g {
88                | Generic::Generic(s) => Some(s),
89                | _ => None,
90            }
91        })
92    }
93
94    /// Returns an iterator which contains only the lifetimes
95    pub fn iter_lifetimes(&self) -> impl Iterator<Item = &Lifetime> {
96        self.iter().filter_map(|g| {
97            match g {
98                | Generic::Lifetime(s) => Some(s),
99                | _ => None,
100            }
101        })
102    }
103
104    /// Returns an iterator which contains only the const generics
105    pub fn iter_consts(&self) -> impl Iterator<Item = &ConstGeneric> {
106        self.iter().filter_map(|g| {
107            match g {
108                | Generic::Const(s) => Some(s),
109                | _ => None,
110            }
111        })
112    }
113
114    pub(crate) fn impl_generics(&self) -> StreamBuilder {
115        let mut result = StreamBuilder::new();
116        result.punct('<');
117
118        for (idx, generic) in self.iter().enumerate() {
119            if idx > 0 {
120                result.punct(',');
121            }
122
123            generic.append_to_result_with_constraints(&mut result);
124        }
125
126        result.punct('>');
127
128        result
129    }
130
131    pub(crate) fn impl_generics_with_additional(
132        &self,
133        lifetimes: &[String],
134        types: &[String],
135    ) -> StreamBuilder {
136        let mut result = StreamBuilder::new();
137        result.punct('<');
138        let mut is_first = true;
139        for lt in lifetimes {
140            if is_first {
141                is_first = false;
142            } else {
143                result.punct(',');
144            }
145            result.lifetime_str(lt);
146        }
147
148        for generic in self.iter() {
149            if is_first {
150                is_first = false;
151            } else {
152                result.punct(',');
153            }
154            generic.append_to_result_with_constraints(&mut result);
155        }
156        for ty in types {
157            if is_first {
158                is_first = false;
159            } else {
160                result.punct(',');
161            }
162            result.ident_str(ty);
163        }
164
165        result.punct('>');
166
167        result
168    }
169
170    pub(crate) fn type_generics(&self) -> StreamBuilder {
171        let mut result = StreamBuilder::new();
172        result.punct('<');
173
174        for (idx, generic) in self.iter().enumerate() {
175            if idx > 0 {
176                result.punct(',');
177            }
178            if generic.is_lifetime() {
179                result.lifetime(generic.ident().clone());
180            } else {
181                result.ident(generic.ident().clone());
182            }
183        }
184
185        result.punct('>');
186        result
187    }
188}
189
190impl Deref for Generics {
191    type Target = Vec<Generic>;
192
193    fn deref(&self) -> &Self::Target {
194        &self.0
195    }
196}
197
198impl DerefMut for Generics {
199    fn deref_mut(&mut self) -> &mut Self::Target {
200        &mut self.0
201    }
202}
203
204/// A single generic argument on a type
205#[derive(Debug, Clone)]
206#[allow(clippy::enum_variant_names)]
207#[non_exhaustive]
208pub enum Generic {
209    /// A lifetime generic
210    ///
211    /// ```
212    /// # use std::marker::PhantomData;
213    /// struct Foo<'a> {
214    ///     // will be Generic::Lifetime("a")
215    /// #   a: PhantomData<&'a ()>,
216    /// }
217    /// ```
218    Lifetime(Lifetime),
219    /// A simple generic
220    ///
221    /// ```
222    /// # use std::marker::PhantomData;
223    /// struct Foo<F> {
224    ///     // will be Generic::Generic("F")
225    /// #   a: PhantomData<F>,
226    /// }
227    /// ```
228    Generic(SimpleGeneric),
229    /// A const generic
230    ///
231    /// ```
232    /// struct Foo<const N: usize> {
233    ///     // will be Generic::Const("N")
234    /// #   a: [u8; N],
235    /// }
236    /// ```
237    Const(ConstGeneric),
238}
239
240impl Generic {
241    const fn is_lifetime(&self) -> bool {
242        matches!(self, Self::Lifetime(_))
243    }
244
245    /// The ident of this generic
246    #[must_use]
247    pub const fn ident(&self) -> &Ident {
248        match self {
249            | Self::Lifetime(lt) => &lt.ident,
250            | Self::Generic(r#gen) => &r#gen.ident,
251            | Self::Const(r#gen) => &r#gen.ident,
252        }
253    }
254
255    const fn has_constraints(&self) -> bool {
256        match self {
257            | Self::Lifetime(lt) => !lt.constraint.is_empty(),
258            | Self::Generic(r#gen) => !r#gen.constraints.is_empty(),
259            | Self::Const(_) => true, // const generics always have a constraint
260        }
261    }
262
263    fn constraints(&self) -> Vec<TokenTree> {
264        match self {
265            | Self::Lifetime(lt) => lt.constraint.clone(),
266            | Self::Generic(r#gen) => r#gen.constraints.clone(),
267            | Self::Const(r#gen) => r#gen.constraints.clone(),
268        }
269    }
270
271    fn append_to_result_with_constraints(
272        &self,
273        builder: &mut StreamBuilder,
274    ) {
275        match self {
276            | Self::Lifetime(lt) => builder.lifetime(lt.ident.clone()),
277            | Self::Generic(r#gen) => builder.ident(r#gen.ident.clone()),
278            | Self::Const(r#gen) => {
279                builder.ident(r#gen.const_token.clone());
280                builder.ident(r#gen.ident.clone())
281            },
282        };
283        if self.has_constraints() {
284            builder.punct(':');
285            builder.extend(self.constraints());
286        }
287    }
288}
289
290impl From<Lifetime> for Generic {
291    fn from(lt: Lifetime) -> Self {
292        Self::Lifetime(lt)
293    }
294}
295
296impl From<SimpleGeneric> for Generic {
297    fn from(r#gen: SimpleGeneric) -> Self {
298        Self::Generic(r#gen)
299    }
300}
301
302impl From<ConstGeneric> for Generic {
303    fn from(r#gen: ConstGeneric) -> Self {
304        Self::Const(r#gen)
305    }
306}
307
308#[test]
309fn test_generics_try_take() {
310    use crate::token_stream;
311
312    assert!(Generics::try_take(&mut token_stream("")).unwrap().is_none());
313    assert!(
314        Generics::try_take(&mut token_stream("foo"))
315            .unwrap()
316            .is_none()
317    );
318    assert!(
319        Generics::try_take(&mut token_stream("()"))
320            .unwrap()
321            .is_none()
322    );
323
324    let stream = &mut token_stream("struct Foo<'a, T>()");
325    let (data_type, ident) = super::DataType::take(stream).unwrap();
326    assert_eq!(data_type, super::DataType::Struct);
327    assert_eq!(ident, "Foo");
328    let generics = Generics::try_take(stream).unwrap().unwrap();
329    assert_eq!(generics.len(), 2);
330    assert_eq!(generics[0].ident(), "a");
331    assert_eq!(generics[1].ident(), "T");
332
333    let stream = &mut token_stream("struct Foo<A, B>()");
334    let (data_type, ident) = super::DataType::take(stream).unwrap();
335    assert_eq!(data_type, super::DataType::Struct);
336    assert_eq!(ident, "Foo");
337    let generics = Generics::try_take(stream).unwrap().unwrap();
338    assert_eq!(generics.len(), 2);
339    assert_eq!(generics[0].ident(), "A");
340    assert_eq!(generics[1].ident(), "B");
341
342    let stream = &mut token_stream("struct Foo<'a, T: Display>()");
343    let (data_type, ident) = super::DataType::take(stream).unwrap();
344    assert_eq!(data_type, super::DataType::Struct);
345    assert_eq!(ident, "Foo");
346    let generics = Generics::try_take(stream).unwrap().unwrap();
347    dbg!(&generics);
348    assert_eq!(generics.len(), 2);
349    assert_eq!(generics[0].ident(), "a");
350    assert_eq!(generics[1].ident(), "T");
351
352    let stream = &mut token_stream("struct Foo<'a, T: for<'a> Bar<'a> + 'static>()");
353    let (data_type, ident) = super::DataType::take(stream).unwrap();
354    assert_eq!(data_type, super::DataType::Struct);
355    assert_eq!(ident, "Foo");
356    dbg!(&generics);
357    assert_eq!(generics.len(), 2);
358    assert_eq!(generics[0].ident(), "a");
359    assert_eq!(generics[1].ident(), "T");
360
361    let stream = &mut token_stream(
362        "struct Baz<T: for<'a> Bar<'a, for<'b> Bar<'b, for<'c> Bar<'c, u32>>>> {}",
363    );
364    let (data_type, ident) = super::DataType::take(stream).unwrap();
365    assert_eq!(data_type, super::DataType::Struct);
366    assert_eq!(ident, "Baz");
367    let generics = Generics::try_take(stream).unwrap().unwrap();
368    dbg!(&generics);
369    assert_eq!(generics.len(), 1);
370    assert_eq!(generics[0].ident(), "T");
371
372    let stream = &mut token_stream("struct Baz<()> {}");
373    let (data_type, ident) = super::DataType::take(stream).unwrap();
374    assert_eq!(data_type, super::DataType::Struct);
375    assert_eq!(ident, "Baz");
376    assert!(
377        Generics::try_take(stream)
378            .unwrap_err()
379            .is_invalid_rust_syntax()
380    );
381
382    let stream = &mut token_stream("struct Bar<A: FnOnce(&'static str) -> SomeStruct, B>");
383    let (data_type, ident) = super::DataType::take(stream).unwrap();
384    assert_eq!(data_type, super::DataType::Struct);
385    assert_eq!(ident, "Bar");
386    let generics = Generics::try_take(stream).unwrap().unwrap();
387    dbg!(&generics);
388    assert_eq!(generics.len(), 2);
389    assert_eq!(generics[0].ident(), "A");
390    assert_eq!(generics[1].ident(), "B");
391
392    let stream = &mut token_stream("struct Bar<A = ()>");
393    let (data_type, ident) = super::DataType::take(stream).unwrap();
394    assert_eq!(data_type, super::DataType::Struct);
395    assert_eq!(ident, "Bar");
396    let generics = Generics::try_take(stream).unwrap().unwrap();
397    dbg!(&generics);
398    assert_eq!(generics.len(), 1);
399    if let Generic::Generic(generic) = &generics[0] {
400        assert_eq!(generic.ident, "A");
401        assert_eq!(generic.default_value.len(), 1);
402        assert_eq!(generic.default_value[0].to_string(), "()");
403    } else {
404        panic!("Expected simple generic, got {:?}", generics[0]);
405    }
406}
407
408/// a lifetime generic parameter, e.g. `struct Foo<'a> { ... }`
409#[derive(Debug, Clone)]
410pub struct Lifetime {
411    /// The ident of this lifetime
412    pub ident: Ident,
413    /// Any constraints that this lifetime may have
414    pub constraint: Vec<TokenTree>,
415}
416
417impl Lifetime {
418    #[allow(clippy::useless_let_if_seq)]
419    pub(crate) fn take(input: &mut Peekable<impl Iterator<Item = TokenTree>>) -> Result<Self> {
420        let start = assume_punct(input.next(), '\'');
421        let ident = match input.peek() {
422            | Some(TokenTree::Ident(_)) => assume_ident(input.next()),
423            | Some(t) => return Err(Error::ExpectedIdent(t.span())),
424            | None => return Err(Error::ExpectedIdent(start.span())),
425        };
426
427        let mut constraint = Vec::new();
428        if let Some(TokenTree::Punct(p)) = input.peek()
429            && p.as_char() == ':'
430        {
431            assume_punct(input.next(), ':');
432            constraint = read_tokens_until_punct(input, &[',', '>'])?;
433        }
434
435        Ok(Self { ident, constraint })
436    }
437
438    #[cfg(test)]
439    fn is_ident(
440        &self,
441        s: &str,
442    ) -> bool {
443        self.ident == s
444    }
445}
446
447#[test]
448fn test_lifetime_take() {
449    use crate::token_stream;
450    use std::panic::catch_unwind;
451    assert!(
452        Lifetime::take(&mut token_stream("'a"))
453            .unwrap()
454            .is_ident("a")
455    );
456    assert!(catch_unwind(|| Lifetime::take(&mut token_stream("'0"))).is_err());
457    assert!(catch_unwind(|| Lifetime::take(&mut token_stream("'("))).is_err());
458    assert!(catch_unwind(|| Lifetime::take(&mut token_stream("')"))).is_err());
459    assert!(catch_unwind(|| Lifetime::take(&mut token_stream("'0'"))).is_err());
460
461    let stream = &mut token_stream("'a: 'b>");
462    let lifetime = Lifetime::take(stream).unwrap();
463    assert_eq!(lifetime.ident, "a");
464    assert_eq!(lifetime.constraint.len(), 2);
465    assume_punct(stream.next(), '>');
466    assert!(stream.next().is_none());
467}
468
469/// a simple generic parameter, e.g. `struct Foo<F> { .. }`
470#[derive(Debug, Clone)]
471#[non_exhaustive]
472pub struct SimpleGeneric {
473    /// The ident of this generic
474    pub ident: Ident,
475    /// The constraints of this generic, e.g. `F: SomeTrait`
476    pub constraints: Vec<TokenTree>,
477    /// The default value of this generic, e.g. `F = ()`
478    pub default_value: Vec<TokenTree>,
479}
480
481impl SimpleGeneric {
482    pub(crate) fn take(input: &mut Peekable<impl Iterator<Item = TokenTree>>) -> Result<Self> {
483        let ident = assume_ident(input.next());
484        let mut constraints = Vec::new();
485        let mut default_value = Vec::new();
486        if let Some(TokenTree::Punct(punct)) = input.peek() {
487            let punct_char = punct.as_char();
488            if punct_char == ':' {
489                assume_punct(input.next(), ':');
490                constraints = read_tokens_until_punct(input, &['>', ','])?;
491            }
492            if punct_char == '=' {
493                assume_punct(input.next(), '=');
494                default_value = read_tokens_until_punct(input, &['>', ','])?;
495            }
496        }
497        Ok(Self {
498            ident,
499            constraints,
500            default_value,
501        })
502    }
503
504    /// The name of this generic, e.g. `T`
505    ///
506    /// # Errors
507    ///
508    /// Returns an error if the operation fails.
509    #[must_use]
510    pub fn name(&self) -> Ident {
511        self.ident.clone()
512    }
513}
514
515/// a const generic parameter, e.g. `struct Foo<const N: usize> { .. }`
516#[derive(Debug, Clone)]
517pub struct ConstGeneric {
518    /// The `const` token for this generic
519    pub const_token: Ident,
520    /// The ident of this generic
521    pub ident: Ident,
522    /// The "constraints" (type) of this generic, e.g. the `usize` from `const N: usize`
523    pub constraints: Vec<TokenTree>,
524}
525
526impl ConstGeneric {
527    #[allow(clippy::useless_let_if_seq)]
528    pub(crate) fn take(input: &mut Peekable<impl Iterator<Item = TokenTree>>) -> Result<Self> {
529        let const_token = assume_ident(input.next());
530        let ident = assume_ident(input.next());
531        let mut constraints = Vec::new();
532        if let Some(TokenTree::Punct(punct)) = input.peek()
533            && punct.as_char() == ':'
534        {
535            assume_punct(input.next(), ':');
536            constraints = read_tokens_until_punct(input, &['>', ','])?;
537        }
538        Ok(Self {
539            const_token,
540            ident,
541            constraints,
542        })
543    }
544}
545
546/// Constraints on generic types.
547///
548/// ```
549/// # use std::marker::PhantomData;
550/// # use std::fmt::Display;
551///
552/// struct Foo<F>
553///     where F: Display // These are `GenericConstraints`
554/// {
555///     f: PhantomData<F>
556/// }
557#[derive(Debug, Clone, Default)]
558pub struct GenericConstraints {
559    constraints: Vec<TokenTree>,
560}
561
562impl GenericConstraints {
563    pub(crate) fn try_take(
564        input: &mut Peekable<impl Iterator<Item = TokenTree>>
565    ) -> Result<Option<Self>> {
566        match input.peek() {
567            | Some(TokenTree::Ident(ident)) => {
568                if !ident_eq(ident, "where") {
569                    return Ok(None);
570                }
571            },
572            | _ => {
573                return Ok(None);
574            },
575        }
576        input.next();
577        let constraints = read_tokens_until_punct(input, &['{', '('])?;
578        Ok(Some(Self { constraints }))
579    }
580
581    pub(crate) fn where_clause(&self) -> StreamBuilder {
582        let mut result = StreamBuilder::new();
583        result.ident_str("where");
584        result.extend(self.constraints.clone());
585        result
586    }
587
588    /// Push the given constraint onto this stream.
589    ///
590    /// ```ignore
591    /// let mut generic_constraints = GenericConstraints::parse("T: Foo"); // imaginary function
592    /// let mut generic = SimpleGeneric::new("U"); // imaginary function
593    ///
594    /// generic_constraints.push_constraint(&generic, "Bar");
595    ///
596    /// // generic_constraints is now:
597    /// // `T: Foo, U: Bar`
598    /// ```
599    ///
600    /// # Errors
601    ///
602    /// Returns an error if parsing fails.
603    pub fn push_constraint(
604        &mut self,
605        generic: &SimpleGeneric,
606        constraint: impl AsRef<str>,
607    ) -> Result<()> {
608        let mut builder = StreamBuilder::new();
609        let last_constraint_was_comma = self
610            .constraints
611            .last()
612            .is_some_and(|l| matches!(l, TokenTree::Punct(c) if c.as_char() == ','));
613        if !self.constraints.is_empty() && !last_constraint_was_comma {
614            builder.punct(',');
615        }
616        builder.ident(generic.ident.clone());
617        builder.punct(':');
618        builder.push_parsed(constraint)?;
619        self.constraints.extend(builder.stream);
620
621        Ok(())
622    }
623
624    /// Push the given constraint onto this stream.
625    ///
626    /// ```ignore
627    /// let mut generic_constraints = GenericConstraints::parse("T: Foo"); // imaginary function
628    ///
629    /// generic_constraints.push_parsed_constraint("u32: SomeTrait");
630    ///
631    /// // generic_constraints is now:
632    /// // `T: Foo, u32: SomeTrait`
633    /// ```
634    ///
635    /// # Panics
636    ///
637    /// Panics if an internal invariant is violated.
638    ///
639    /// # Errors
640    ///
641    /// Returns an error if parsing fails.
642    pub fn push_parsed_constraint(
643        &mut self,
644        constraint: impl AsRef<str>,
645    ) -> Result<()> {
646        let mut builder = StreamBuilder::new();
647        if !self.constraints.is_empty() {
648            builder.punct(',');
649        }
650        builder.push_parsed(constraint)?;
651        self.constraints.extend(builder.stream);
652
653        Ok(())
654    }
655
656    /// Clear the constraints
657    ///
658    /// # Panics
659    ///
660    /// Panics if an internal invariant is violated.
661    pub fn clear(&mut self) {
662        self.constraints.clear();
663    }
664}
665
666#[test]
667fn test_generic_constraints_try_take() {
668    use super::DataType;
669    use super::StructBody;
670    use super::Visibility;
671    use crate::parse::body::Fields;
672    use crate::token_stream;
673
674    let stream = &mut token_stream("struct Foo where Foo: Bar { }");
675    DataType::take(stream).unwrap();
676    assert!(GenericConstraints::try_take(stream).unwrap().is_some());
677
678    let stream = &mut token_stream("struct Foo { }");
679    DataType::take(stream).unwrap();
680    assert!(GenericConstraints::try_take(stream).unwrap().is_none());
681
682    let stream = &mut token_stream("struct Foo where Foo: Bar(Foo)");
683    DataType::take(stream).unwrap();
684    assert!(GenericConstraints::try_take(stream).unwrap().is_some());
685
686    let stream = &mut token_stream("struct Foo()");
687    DataType::take(stream).unwrap();
688    assert!(GenericConstraints::try_take(stream).unwrap().is_none());
689
690    let stream = &mut token_stream("struct Foo()");
691    assert!(GenericConstraints::try_take(stream).unwrap().is_none());
692
693    let stream = &mut token_stream("{}");
694    assert!(GenericConstraints::try_take(stream).unwrap().is_none());
695
696    let stream = &mut token_stream("");
697    assert!(GenericConstraints::try_take(stream).unwrap().is_none());
698
699    let stream = &mut token_stream("pub(crate) struct Test<T: Encode> {}");
700    assert_eq!(Visibility::Pub, Visibility::try_take(stream).unwrap());
701    let (data_type, ident) = DataType::take(stream).unwrap();
702    assert_eq!(data_type, DataType::Struct);
703    assert_eq!(ident, "Test");
704    let constraints = Generics::try_take(stream).unwrap().unwrap();
705    assert_eq!(constraints.len(), 1);
706    assert_eq!(constraints[0].ident(), "T");
707    let body = StructBody::take(stream).unwrap();
708    if let Some(Fields::Struct(v)) = body.fields {
709        assert!(v.is_empty());
710    } else {
711        panic!("wrong fields {:?}", body.fields);
712    }
713}
714
715#[test]
716fn test_generic_constraints_trailing_comma() {
717    use crate::parse::Attribute;
718    use crate::parse::AttributeLocation;
719    use crate::parse::DataType;
720    use crate::parse::GenericConstraints;
721    use crate::parse::Generics;
722    use crate::parse::StructBody;
723    use crate::parse::Visibility;
724    use crate::token_stream;
725    let source = &mut token_stream("pub struct MyStruct<T> where T: Clone, { }");
726
727    Attribute::try_take(AttributeLocation::Container, source).unwrap();
728    Visibility::try_take(source).unwrap();
729    DataType::take(source).unwrap();
730    Generics::try_take(source).unwrap().unwrap();
731    GenericConstraints::try_take(source).unwrap().unwrap();
732    StructBody::take(source).unwrap();
733}