Skip to main content

virtue_next/generate/
mod.rs

1//! Code to help generate functions.
2//!
3//! The structure is:
4//!
5//! - [`Generator`]
6//!   - `.impl_for()`: [`ImplFor`]
7//!     - `.generate_fn()`: [`FnBuilder`]
8//!       - `.body(|builder| { .. })`: [`StreamBuilder`]
9//!
10//! Afterwards, [`Generator::finish()`] **must** be called to take out the [`TokenStream`] produced.
11//!
12//! [`Generator::finish()`]: struct.Generator.html#method.finish
13//! [`TokenStream`]: ../prelude/struct.TokenStream.html
14
15mod gen_enum;
16mod gen_struct;
17mod generate_item;
18mod generate_mod;
19mod generator;
20mod r#impl;
21mod impl_for;
22mod stream_builder;
23
24use crate::parse::Visibility;
25use crate::{
26    parse::{GenericConstraints, Generics},
27    prelude::{Delimiter, Ident, TokenStream},
28};
29use std::fmt;
30use std::marker::PhantomData;
31
32pub use self::gen_enum::GenEnum;
33pub use self::gen_struct::GenStruct;
34pub use self::generate_item::{FnBuilder, FnSelfArg, GenConst};
35pub use self::generate_mod::GenerateMod;
36pub use self::generator::Generator;
37pub use self::r#impl::Impl;
38pub use self::impl_for::ImplFor;
39pub use self::stream_builder::{PushParseError, StreamBuilder};
40
41/// Helper trait to make it possible to nest several builders. Internal use only.
42#[allow(missing_docs)]
43pub trait Parent {
44    fn append(&mut self, builder: StreamBuilder);
45    fn name(&self) -> &Ident;
46    fn generics(&self) -> Option<&Generics>;
47    fn generic_constraints(&self) -> Option<&GenericConstraints>;
48}
49
50/// Helper enum to differentiate between a [`Ident`] or a [`String`].
51#[allow(missing_docs)]
52pub enum StringOrIdent {
53    String(String),
54    // Note that when this is a `string` this could be much more than a single ident.
55    // Therefor you should never use [`StreamBuilder`]`.ident_str(StringOrIdent.to_string())`, but instead use `.push_parsed(StringOrIdent.to_string())?`.
56    Ident(Ident),
57}
58
59impl fmt::Display for StringOrIdent {
60    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61        match self {
62            Self::String(s) => s.fmt(f),
63            Self::Ident(i) => i.fmt(f),
64        }
65    }
66}
67
68impl From<String> for StringOrIdent {
69    fn from(s: String) -> Self {
70        Self::String(s)
71    }
72}
73impl From<Ident> for StringOrIdent {
74    fn from(i: Ident) -> Self {
75        Self::Ident(i)
76    }
77}
78impl<'a> From<&'a str> for StringOrIdent {
79    fn from(s: &'a str) -> Self {
80        Self::String(s.to_owned())
81    }
82}
83
84/// A path of identifiers, like `mod::Type`.
85pub struct Path(Vec<StringOrIdent>);
86
87impl From<String> for Path {
88    fn from(s: String) -> Self {
89        StringOrIdent::from(s).into()
90    }
91}
92
93impl From<Ident> for Path {
94    fn from(i: Ident) -> Self {
95        StringOrIdent::from(i).into()
96    }
97}
98
99impl From<&str> for Path {
100    fn from(s: &str) -> Self {
101        StringOrIdent::from(s).into()
102    }
103}
104
105impl From<StringOrIdent> for Path {
106    fn from(value: StringOrIdent) -> Self {
107        Self(vec![value])
108    }
109}
110
111impl FromIterator<String> for Path {
112    fn from_iter<T: IntoIterator<Item = String>>(iter: T) -> Self {
113        iter.into_iter().map(StringOrIdent::from).collect()
114    }
115}
116
117impl FromIterator<Ident> for Path {
118    fn from_iter<T: IntoIterator<Item = Ident>>(iter: T) -> Self {
119        iter.into_iter().map(StringOrIdent::from).collect()
120    }
121}
122
123impl<'a> FromIterator<&'a str> for Path {
124    fn from_iter<T: IntoIterator<Item = &'a str>>(iter: T) -> Self {
125        iter.into_iter().map(StringOrIdent::from).collect()
126    }
127}
128
129impl FromIterator<StringOrIdent> for Path {
130    fn from_iter<T: IntoIterator<Item = StringOrIdent>>(iter: T) -> Self {
131        Self(iter.into_iter().collect())
132    }
133}
134
135impl IntoIterator for Path {
136    type Item = StringOrIdent;
137    type IntoIter = std::vec::IntoIter<StringOrIdent>;
138
139    fn into_iter(self) -> Self::IntoIter {
140        self.0.into_iter()
141    }
142}
143
144/// A struct or enum variant field.
145struct Field {
146    name: String,
147    vis: Visibility,
148    ty: String,
149    attributes: Vec<StreamBuilder>,
150}
151
152impl Field {
153    fn new(name: impl Into<String>, vis: Visibility, ty: impl Into<String>) -> Self {
154        Self {
155            name: name.into(),
156            vis,
157            ty: ty.into(),
158            attributes: Vec::new(),
159        }
160    }
161}
162
163/// A builder for struct or enum variant fields.
164pub struct FieldBuilder<'a, P> {
165    fields: &'a mut Vec<Field>,
166    _parent: PhantomData<P>, // Keep this to disallow `pub` on enum fields
167}
168
169impl<P> FieldBuilder<'_, P> {
170    /// Add an attribute to the field.
171    ///
172    /// ```
173    /// # use virtue::prelude::Generator;
174    /// # let mut generator = Generator::with_name("Fooz");
175    /// generator
176    ///     .generate_struct("Foo")
177    ///     .add_field("foo", "u16")
178    ///     .make_pub()
179    ///     .with_attribute("serde", |b| {
180    ///         b.push_parsed("(default)")?;
181    ///         Ok(())
182    ///     })?;
183    /// generator
184    ///     .generate_enum("Bar")
185    ///     .add_value("Baz")
186    ///     .add_field("baz", "bool")
187    ///     .with_attribute("serde", |b| {
188    ///         b.push_parsed("(default)")?;
189    ///         Ok(())
190    ///     })?;
191    /// # generator.assert_eq("struct Foo { # [serde (default)] pub foo : u16 , } \
192    /// enum Bar { Baz { # [serde (default)] baz : bool , } , }");
193    /// # Ok::<_, virtue::Error>(())
194    /// ```
195    ///
196    /// Generates:
197    /// ```ignore
198    /// struct Foo {
199    ///     #[serde(default)]
200    ///     pub bar: u16
201    /// }
202    ///
203    /// enum Bar {
204    ///     Baz {
205    ///         #[serde(default)]
206    ///         baz: bool
207    ///     }
208    /// }
209    /// ```
210    pub fn with_attribute(
211        &mut self,
212        name: impl AsRef<str>,
213        value: impl FnOnce(&mut StreamBuilder) -> crate::Result,
214    ) -> crate::Result<&mut Self> {
215        self.current().with_attribute(name, value)?;
216        Ok(self)
217    }
218
219    /// Add a parsed attribute to the field.
220    ///
221    /// ```
222    /// # use virtue::prelude::Generator;
223    /// # let mut generator = Generator::with_name("Fooz");
224    /// generator
225    ///     .generate_struct("Foo")
226    ///     .add_field("foo", "u16")
227    ///     .make_pub()
228    ///     .with_parsed_attribute("serde(default)")?;
229    /// generator
230    ///     .generate_enum("Bar")
231    ///     .add_value("Baz")
232    ///     .add_field("baz", "bool")
233    ///     .with_parsed_attribute("serde(default)")?;
234    /// # generator.assert_eq("struct Foo { # [serde (default)] pub foo : u16 , } \
235    /// enum Bar { Baz { # [serde (default)] baz : bool , } , }");
236    /// # Ok::<_, virtue::Error>(())
237    /// ```
238    ///
239    /// Generates:
240    /// ```ignore
241    /// struct Foo {
242    ///     #[serde(default)]
243    ///     pub bar: u16
244    /// }
245    ///
246    /// enum Bar {
247    ///     Baz {
248    ///         #[serde(default)]
249    ///         baz: bool
250    ///     }
251    /// }
252    /// ```
253    pub fn with_parsed_attribute(
254        &mut self,
255        attribute: impl AsRef<str>,
256    ) -> crate::Result<&mut Self> {
257        self.current().with_parsed_attribute(attribute)?;
258        Ok(self)
259    }
260
261    /// Add a token stream as an attribute to the field.
262    ///
263    /// ```
264    /// # use virtue::prelude::{Generator, TokenStream};
265    /// # let mut generator = Generator::with_name("Fooz");
266    /// let attribute = "serde(default)".parse::<TokenStream>().unwrap();
267    /// generator
268    ///     .generate_struct("Foo")
269    ///     .add_field("foo", "u16")
270    ///     .make_pub()
271    ///     .with_attribute_stream(attribute);
272    /// # generator.assert_eq("struct Foo { # [serde (default)] pub foo : u16 , }");
273    /// # Ok::<_, virtue::Error>(())
274    /// ```
275    ///
276    /// Generates:
277    /// ```ignore
278    /// struct Foo {
279    ///     #[serde(default)]
280    ///     pub bar: u16
281    /// }
282    /// ```
283    pub fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
284        self.current().with_attribute_stream(attribute);
285        self
286    }
287
288    /// Add a field to the parent type.
289    ///
290    /// ```
291    /// # use virtue::prelude::Generator;
292    /// # let mut generator = Generator::with_name("Fooz");
293    /// generator
294    ///     .generate_struct("Foo")
295    ///     .add_field("foo", "u16")
296    ///     .add_field("bar", "bool");
297    /// # generator.assert_eq("struct Foo { foo : u16 , bar : bool , }");
298    /// # Ok::<_, virtue::Error>(())
299    /// ```
300    ///
301    /// Generates:
302    /// ```
303    /// struct Foo {
304    ///     foo: u16,
305    ///     bar: bool
306    /// }
307    /// ```
308    pub fn add_field(&mut self, name: impl Into<String>, ty: impl Into<String>) -> &mut Self {
309        self.fields.push(Field::new(name, Visibility::Default, ty));
310        self
311    }
312}
313
314// Only allow `pub` on struct fields
315impl<'a, P: Parent> FieldBuilder<'_, GenStruct<'a, P>> {
316    /// Make the field public.
317    pub fn make_pub(&mut self) -> &mut Self {
318        self.current().vis = Visibility::Pub;
319        self
320    }
321}
322
323impl<'a, P> From<&'a mut Vec<Field>> for FieldBuilder<'a, P> {
324    fn from(fields: &'a mut Vec<Field>) -> Self {
325        Self {
326            fields,
327            _parent: PhantomData,
328        }
329    }
330}
331
332impl<P> FieldBuilder<'_, P> {
333    fn current(&mut self) -> &mut Field {
334        // A field is always added before this is called, so the unwrap doesn't fail.
335        self.fields.last_mut().unwrap()
336    }
337}
338
339/// A helper trait to share attribute code between struct and enum generators.
340trait AttributeContainer {
341    fn derives(&mut self) -> &mut Vec<Path>;
342    fn attributes(&mut self) -> &mut Vec<StreamBuilder>;
343
344    fn with_derive(&mut self, derive: impl Into<Path>) -> &mut Self {
345        self.derives().push(derive.into());
346        self
347    }
348
349    fn with_derives<T: Into<Path>>(&mut self, derives: impl IntoIterator<Item = T>) -> &mut Self {
350        self.derives().extend(derives.into_iter().map(Into::into));
351        self
352    }
353
354    fn with_attribute(
355        &mut self,
356        name: impl AsRef<str>,
357        value: impl FnOnce(&mut StreamBuilder) -> crate::Result,
358    ) -> crate::Result<&mut Self> {
359        let mut stream = StreamBuilder::new();
360        value(stream.ident_str(name))?;
361        self.attributes().push(stream);
362        Ok(self)
363    }
364
365    fn with_parsed_attribute(&mut self, attribute: impl AsRef<str>) -> crate::Result<&mut Self> {
366        let mut stream = StreamBuilder::new();
367        stream.push_parsed(attribute)?;
368        self.attributes().push(stream);
369        Ok(self)
370    }
371
372    fn with_attribute_stream(&mut self, attribute: impl Into<TokenStream>) -> &mut Self {
373        let stream = StreamBuilder {
374            stream: attribute.into(),
375        };
376        self.attributes().push(stream);
377        self
378    }
379
380    fn build_derives(&mut self, b: &mut StreamBuilder) -> &mut Self {
381        let derives = std::mem::take(self.derives());
382        if !derives.is_empty() {
383            build_attribute(b, |b| {
384                b.ident_str("derive").group(Delimiter::Parenthesis, |b| {
385                    for (idx, derive) in derives.into_iter().enumerate() {
386                        if idx > 0 {
387                            b.punct(',');
388                        }
389                        for (idx, component) in derive.into_iter().enumerate() {
390                            if idx > 0 {
391                                b.puncts("::");
392                            }
393
394                            match component {
395                                StringOrIdent::String(s) => b.ident_str(s),
396                                StringOrIdent::Ident(i) => b.ident(i),
397                            };
398                        }
399                    }
400                    Ok(())
401                })
402            })
403            .expect("could not build derives");
404        }
405        self
406    }
407
408    fn build_attributes(&mut self, b: &mut StreamBuilder) -> &mut Self {
409        for attr in std::mem::take(self.attributes()) {
410            build_attribute(b, |b| Ok(b.extend(attr.stream))).expect("could not build attribute");
411        }
412        self
413    }
414}
415
416impl AttributeContainer for Field {
417    fn derives(&mut self) -> &mut Vec<Path> {
418        unreachable!("fields cannot have derives")
419    }
420
421    fn attributes(&mut self) -> &mut Vec<StreamBuilder> {
422        &mut self.attributes
423    }
424}
425
426fn build_attribute<T>(b: &mut StreamBuilder, build: T) -> crate::Result
427where
428    T: FnOnce(&mut StreamBuilder) -> crate::Result<&mut StreamBuilder>,
429{
430    b.punct('#').group(Delimiter::Bracket, |b| {
431        build(b)?;
432        Ok(())
433    })?;
434
435    Ok(())
436}