Implement trait generics on impl-for

src/generate/impl_for.rs

@@ -13,6 +13,7 @@ pub struct ImplFor<'a, P: Parent> {
     type_name: StringOrIdent,
     trait_name: Option<StringOrIdent>,
     lifetimes: Option<Vec<String>>,
+    generics: Option<Vec<String>>,
     consts: Vec<StreamBuilder>,
     custom_generic_constraints: Option<GenericConstraints>,
     impl_types: Vec<StreamBuilder>,
@@ -32,6 +33,7 @@ impl<'a, P: Parent> ImplFor<'a, P> {
             trait_name,
             type_name,
             lifetimes: None,
+            generics: None,
             consts: Vec::new(),
             custom_generic_constraints: None,
             impl_types: Vec::new(),
@@ -77,6 +79,31 @@ impl<'a, P: Parent> ImplFor<'a, P> {
         self
     }
 
+    /// Add generic parameters to the trait implementation.
+    ///```
+    /// # use virtue::prelude::Generator;
+    /// # let mut generator = Generator::with_name("Bar");
+    /// generator.impl_for("Foo")
+    ///          .with_trait_generics(["Baz"]);
+    /// # generator.assert_eq("impl Foo < Baz > for Bar { }");
+    /// # Ok::<_, virtue::Error>(())
+    /// ```
+    ///
+    /// Generates:
+    /// ```ignore
+    /// impl Foo for <struct or enum> {
+    ///     const BAR: u8 = 5;
+    /// }
+    /// ```
+    pub fn with_trait_generics<ITER>(mut self, generics: ITER) -> Self
+    where
+        ITER: IntoIterator,
+        ITER::Item: Into<String>,
+    {
+        self.generics = Some(generics.into_iter().map(Into::into).collect());
+        self
+    }
+
     /// Add a outer attribute to the trait implementation
     pub fn impl_outer_attr(&mut self, attr: impl AsRef<str>) -> Result {
         let mut builder = StreamBuilder::new();
@@ -259,16 +286,17 @@ impl<P: Parent> ImplFor<'_, P> {
             if let Some(generics) = self.generator.generics() {
                 builder.append(generics.impl_generics_with_additional_lifetimes(lifetimes));
             } else {
-                append_lifetimes(builder, lifetimes);
+                append_lifetimes_and_generics(builder, lifetimes, &[]);
             }
         } else if let Some(generics) = self.generator.generics() {
             builder.append(generics.impl_generics());
         }
         if let Some(t) = &self.trait_name {
             builder.push_parsed(t.to_string()).unwrap();
-            if let Some(lifetimes) = &self.lifetimes {
-                append_lifetimes(builder, lifetimes);
-            }
+
+            let lifetimes = self.lifetimes.as_deref().unwrap_or_default();
+            let generics = self.generics.as_deref().unwrap_or_default();
+            append_lifetimes_and_generics(builder, lifetimes, generics);
             builder.ident_str("for");
         }
         builder.push_parsed(self.type_name.to_string()).unwrap();
@@ -283,10 +311,30 @@ impl<P: Parent> ImplFor<'_, P> {
     }
 }
 
-fn append_lifetimes(builder: &mut StreamBuilder, lifetimes: &[String]) {
+fn append_lifetimes_and_generics(
+    builder: &mut StreamBuilder,
+    lifetimes: &[String],
+    generics: &[String],
+) {
+    if lifetimes.is_empty() && generics.is_empty() {
+        return;
+    }
+
+    builder.punct('<');
+
     for (idx, lt) in lifetimes.iter().enumerate() {
-        builder.punct(if idx == 0 { '<' } else { ',' });
+        if idx > 0 {
+            builder.punct(',');
+        }
         builder.lifetime_str(lt);
     }
+
+    for (idx, gen) in generics.iter().enumerate() {
+        if idx > 0 || !lifetimes.is_empty() {
+            builder.punct(',');
+        }
+        builder.ident_str(gen);
+    }
+
     builder.punct('>');
 }