From 5a4f7c72dd617d638832c36b396a0db5cdebdb95 Mon Sep 17 00:00:00 2001
From: NightEule5 <24661563+NightEule5@users.noreply.github.com>
Date: Fri, 10 Nov 2023 07:45:36 -0700
Subject: [PATCH] Implement trait generics on impl-for

---
 src/generate/impl_for.rs | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/src/generate/impl_for.rs b/src/generate/impl_for.rs
index b9eb421..0962560 100644
--- a/src/generate/impl_for.rs
+++ b/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: Vec::new(),
             consts: Vec::new(),
             custom_generic_constraints: None,
             impl_types: Vec::new(),
@@ -77,6 +79,23 @@ impl<'a, P: Parent> ImplFor<'a, P> {
         self
     }
 
+    /// Add generic parameters to the trait.
+    ///
+    /// `generator.impl_for("Foo").with_trait_generic(["Bar", "Baz"])` results in code like:
+    ///
+    /// ```ignore
+    /// impl Foo<Bar, Baz> for <struct or enum> {
+    /// }
+    /// ```
+    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();
@@ -266,9 +285,9 @@ impl<P: Parent> ImplFor<'_, P> {
         }
         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_generics(builder, lifetimes, generics);
             builder.ident_str("for");
         }
         builder.push_parsed(self.type_name.to_string()).unwrap();
@@ -283,10 +302,18 @@ impl<P: Parent> ImplFor<'_, P> {
     }
 }
 
-fn append_lifetimes(builder: &mut StreamBuilder, lifetimes: &[String]) {
+fn append_generics(builder: &mut StreamBuilder, lifetimes: &[String], generics: &[String]) {
     for (idx, lt) in lifetimes.iter().enumerate() {
         builder.punct(if idx == 0 { '<' } else { ',' });
         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('>');
 }