virtue_next/generate/
impl_for.rs

1use super::{FnBuilder, GenConst, Parent, StreamBuilder, StringOrIdent, generate_item::FnParent};
2use crate::{
3    parse::{GenericConstraints, Generics},
4    prelude::{Delimiter, Result},
5};
6
7#[must_use]
8/// A helper struct for implementing a trait for a given struct or enum.
9pub struct ImplFor<'a, P: Parent> {
10    generator: &'a mut P,
11    outer_attr: Vec<StreamBuilder>,
12    inner_attr: Vec<StreamBuilder>,
13    type_name: StringOrIdent,
14    trait_name: Option<StringOrIdent>,
15    lifetimes: Option<Vec<String>>,
16    trait_generics: Option<Vec<String>>,
17    impl_generics: Vec<String>,
18    consts: Vec<StreamBuilder>,
19    custom_generic_constraints: Option<GenericConstraints>,
20    impl_types: Vec<StreamBuilder>,
21    fns: Vec<(StreamBuilder, StreamBuilder)>,
22}
23
24impl<'a, P: Parent> ImplFor<'a, P> {
25    pub(super) fn new(
26        generator: &'a mut P,
27        type_name: StringOrIdent,
28        trait_name: Option<StringOrIdent>,
29    ) -> Self {
30        Self {
31            generator,
32            outer_attr: Vec::new(),
33            inner_attr: Vec::new(),
34            trait_name,
35            type_name,
36            lifetimes: None,
37            trait_generics: None,
38            impl_generics: vec![],
39            consts: Vec::new(),
40            custom_generic_constraints: None,
41            impl_types: Vec::new(),
42            fns: Vec::new(),
43        }
44    }
45
46    /// Internal helper function to set lifetimes
47    pub(crate) fn with_lifetimes<ITER>(mut self, lifetimes: ITER) -> Self
48    where
49        ITER: IntoIterator,
50        ITER::Item: Into<String>,
51    {
52        self.lifetimes = Some(lifetimes.into_iter().map(Into::into).collect());
53        self
54    }
55
56    /// Make the new lifetimes added by `Generator::impl_for_with_lifetimes` depend on the existing lifetimes from the original derive.
57    ///
58    /// See [`impl_for_with_lifetimes`] for more information.
59    ///
60    /// Calling this method in any other context has no effect.
61    ///
62    /// [`impl_for_with_lifetimes`]: struct.Generator.html#method.impl_for_with_lifetimes
63    pub fn new_lifetimes_depend_on_existing(mut self) -> Self {
64        if let Some(new_lt) = &self.lifetimes {
65            if let Some(generics) = self.generator.generics() {
66                let constraints = self.custom_generic_constraints.get_or_insert_with(|| {
67                    self.generator
68                        .generic_constraints()
69                        .cloned()
70                        .unwrap_or_default()
71                });
72                for old_lt in generics.iter_lifetimes() {
73                    for new_lt in new_lt {
74                        constraints
75                            .push_parsed_constraint(format!("'{}: '{}", new_lt, old_lt.ident))
76                            .expect("Could not ensure new lifetimes depend on existing lifetimes");
77                    }
78                }
79            }
80        }
81        self
82    }
83
84    /// Add generic parameters to the trait implementation.
85    ///```
86    /// # use virtue::prelude::Generator;
87    /// # let mut generator = Generator::with_name("Bar");
88    /// generator.impl_for("Foo")
89    ///          .with_trait_generics(["Baz"]);
90    /// # generator.assert_eq("impl Foo < Baz > for Bar { }");
91    /// # Ok::<_, virtue::Error>(())
92    /// ```
93    ///
94    /// Generates:
95    /// ```ignore
96    /// impl Foo for <struct or enum> {
97    ///     const BAR: u8 = 5;
98    /// }
99    /// ```
100    pub fn with_trait_generics<ITER>(mut self, generics: ITER) -> Self
101    where
102        ITER: IntoIterator,
103        ITER::Item: Into<String>,
104    {
105        self.trait_generics = Some(generics.into_iter().map(Into::into).collect());
106        self
107    }
108
109    /// Add generic parameters to the impl block.
110    ///```
111    /// # use virtue::prelude::Generator;
112    /// # let mut generator = Generator::with_name("Bar");
113    /// generator.impl_for("Foo")
114    ///          .with_impl_generics(["Baz"]);
115    /// # generator.assert_eq("impl < Baz > Foo for Bar { }");
116    /// # Ok::<_, virtue::Error>(())
117    /// ```
118    ///
119    /// Generates:
120    /// ```ignore
121    /// impl<Baz> Foo for Bar { }
122    /// ```
123    pub fn with_impl_generics<ITER>(mut self, generics: ITER) -> Self
124    where
125        ITER: IntoIterator,
126        ITER::Item: Into<String>,
127    {
128        self.impl_generics = generics.into_iter().map(Into::into).collect();
129        self
130    }
131
132    /// Add a outer attribute to the trait implementation
133    pub fn impl_outer_attr(&mut self, attr: impl AsRef<str>) -> Result {
134        let mut builder = StreamBuilder::new();
135        builder.punct('#').group(Delimiter::Bracket, |builder| {
136            builder.push_parsed(attr)?;
137            Ok(())
138        })?;
139        self.outer_attr.push(builder);
140        Ok(())
141    }
142
143    /// Add a inner attribute to the trait implementation
144    pub fn impl_inner_attr(&mut self, attr: impl AsRef<str>) -> Result {
145        let mut builder = StreamBuilder::new();
146        builder
147            .punct('#')
148            .punct('!')
149            .group(Delimiter::Brace, |builder| {
150                builder.push_parsed(attr)?;
151                Ok(())
152            })?;
153        self.inner_attr.push(builder);
154        Ok(())
155    }
156
157    /// Add a const to the trait implementation
158    /// ```
159    /// # use virtue::prelude::Generator;
160    /// # let mut generator = Generator::with_name("Bar");
161    /// generator.impl_for("Foo")
162    ///          .generate_const("BAR", "u8")
163    ///          .with_value(|b| {
164    ///             b.push_parsed("5")?;
165    ///             Ok(())
166    ///          })?;
167    /// # generator.assert_eq("impl Foo for Bar { const BAR : u8 = 5 ; }");
168    /// # Ok::<_, virtue::Error>(())
169    /// ```
170    ///
171    /// Generates:
172    /// ```ignore
173    /// impl Foo for <struct or enum> {
174    ///     const BAR: u8 = 5;
175    /// }
176    pub fn generate_const(
177        &mut self,
178        name: impl Into<String>,
179        ty: impl Into<String>,
180    ) -> GenConst<'_> {
181        GenConst::new(&mut self.consts, name, ty)
182    }
183
184    /// Add a function to the trait implementation.
185    ///
186    /// `generator.impl_for("Foo").generate_fn("bar")` results in code like:
187    ///
188    /// ```ignore
189    /// impl Foo for <struct or enum> {
190    ///     fn bar() {}
191    /// }
192    /// ```
193    ///
194    /// See [`FnBuilder`] for more options, as well as information on how to fill the function body.
195    pub fn generate_fn(&mut self, name: impl Into<String>) -> FnBuilder<'_, ImplFor<'a, P>> {
196        FnBuilder::new(self, name)
197    }
198
199    /// Add a type to the impl
200    ///
201    /// `generator.impl_for("Foo").impl_type("Bar", "u8")` results in code like:
202    ///
203    /// ```ignore
204    /// impl Foo for <struct or enum> {
205    ///     type Bar = u8;
206    /// }
207    /// ```
208    pub fn impl_type(&mut self, name: impl AsRef<str>, value: impl AsRef<str>) -> Result {
209        let mut builder = StreamBuilder::new();
210        builder
211            .ident_str("type")
212            .push_parsed(name)?
213            .punct('=')
214            .push_parsed(value)?
215            .punct(';');
216        self.impl_types.push(builder);
217        Ok(())
218    }
219
220    ///
221    /// Modify the generic constraints of a type.
222    /// This can be used to add additional type constraints to your implementation.
223    ///
224    /// ```ignore
225    /// // Your derive:
226    /// #[derive(YourTrait)]
227    /// pub struct Foo<B> {
228    ///     ...
229    /// }
230    ///
231    /// // With this code:
232    /// generator
233    ///     .impl_for("YourTrait")
234    ///     .modify_generic_constraints(|generics, constraints| {
235    ///         for g in generics.iter_generics() {
236    ///             constraints.push_generic(g, "YourTrait");
237    ///         }
238    ///     })
239    ///
240    /// // will generate:
241    /// impl<B> YourTrait for Foo<B>
242    ///     where B: YourTrait // <-
243    /// {
244    /// }
245    /// ```
246    ///
247    pub fn modify_generic_constraints<CB>(&mut self, cb: CB) -> Result<&mut Self>
248    where
249        CB: FnOnce(&Generics, &mut GenericConstraints) -> Result,
250    {
251        if let Some(generics) = self.generator.generics() {
252            let constraints = self.custom_generic_constraints.get_or_insert_with(|| {
253                self.generator
254                    .generic_constraints()
255                    .cloned()
256                    .unwrap_or_default()
257            });
258            cb(generics, constraints)?;
259        }
260        Ok(self)
261    }
262}
263
264impl<'a, P: Parent> FnParent for ImplFor<'a, P> {
265    fn append(&mut self, fn_definition: StreamBuilder, fn_body: StreamBuilder) -> Result {
266        self.fns.push((fn_definition, fn_body));
267        Ok(())
268    }
269}
270
271impl<P: Parent> Drop for ImplFor<'_, P> {
272    fn drop(&mut self) {
273        if std::thread::panicking() {
274            return;
275        }
276        let mut builder = StreamBuilder::new();
277        for attr in std::mem::take(&mut self.outer_attr) {
278            builder.append(attr);
279        }
280
281        self.generate_impl_definition(&mut builder);
282
283        builder
284            .group(Delimiter::Brace, |builder| {
285                for attr in std::mem::take(&mut self.inner_attr) {
286                    builder.append(attr);
287                }
288                for ty in std::mem::take(&mut self.impl_types) {
289                    builder.append(ty);
290                }
291                for r#const in std::mem::take(&mut self.consts) {
292                    builder.append(r#const);
293                }
294                for (fn_def, fn_body) in std::mem::take(&mut self.fns) {
295                    builder.append(fn_def);
296                    builder
297                        .group(Delimiter::Brace, |body| {
298                            *body = fn_body;
299                            Ok(())
300                        })
301                        .unwrap();
302                }
303                Ok(())
304            })
305            .unwrap();
306
307        self.generator.append(builder);
308    }
309}
310
311impl<P: Parent> ImplFor<'_, P> {
312    fn generate_impl_definition(&mut self, builder: &mut StreamBuilder) {
313        builder.ident_str("impl");
314
315        let impl_generics = self.impl_generics.as_slice();
316        if let Some(lifetimes) = &self.lifetimes {
317            if let Some(generics) = self.generator.generics() {
318                builder.append(generics.impl_generics_with_additional(lifetimes, impl_generics));
319            } else {
320                append_lifetimes_and_generics(builder, lifetimes, impl_generics);
321            }
322        } else if let Some(generics) = self.generator.generics() {
323            builder.append(generics.impl_generics_with_additional(&[], impl_generics));
324        } else if !impl_generics.is_empty() {
325            append_lifetimes_and_generics(builder, &[], impl_generics)
326        }
327        if let Some(t) = &self.trait_name {
328            builder.push_parsed(t.to_string()).unwrap();
329
330            let lifetimes = self.lifetimes.as_deref().unwrap_or_default();
331            let generics = self.trait_generics.as_deref().unwrap_or_default();
332            append_lifetimes_and_generics(builder, lifetimes, generics);
333            builder.ident_str("for");
334        }
335        builder.push_parsed(self.type_name.to_string()).unwrap();
336        if let Some(generics) = &self.generator.generics() {
337            builder.append(generics.type_generics());
338        }
339        match self.custom_generic_constraints.take() {
340            Some(generic_constraints) => {
341                builder.append(generic_constraints.where_clause());
342            }
343            _ => {
344                if let Some(generic_constraints) = &self.generator.generic_constraints() {
345                    builder.append(generic_constraints.where_clause());
346                }
347            }
348        }
349    }
350}
351
352fn append_lifetimes_and_generics(
353    builder: &mut StreamBuilder,
354    lifetimes: &[String],
355    generics: &[String],
356) {
357    if lifetimes.is_empty() && generics.is_empty() {
358        return;
359    }
360
361    builder.punct('<');
362
363    for (idx, lt) in lifetimes.iter().enumerate() {
364        if idx > 0 {
365            builder.punct(',');
366        }
367        builder.lifetime_str(lt);
368    }
369
370    for (idx, r#gen) in generics.iter().enumerate() {
371        if idx > 0 || !lifetimes.is_empty() {
372            builder.punct(',');
373        }
374        builder.push_parsed(r#gen).unwrap();
375    }
376
377    builder.punct('>');
378}
virtue_next/generate/impl_for.rs

virtue_next/generate/
impl_for.rs
