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