#116: Add const support

la10736 · May 8, 2021 · 7b8742b · 7b8742b
1 parent 75a133a
commit 7b8742b
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 5 deletions.
diff --git a/src/render/fixture.rs b/src/render/fixture.rs
@@ -361,6 +361,20 @@ mod should {
             "fn partial_1<T: IntoIterator>(mut i: T) where T::Item: Copy {}",
         ]
     )]
+    #[case::not_remove_const_generics("fn test<const N:usize>(v: [u32; N]) -> [i32; N] {}",
+        vec![
+            "fn default<const N:usize>() -> [i32; N] {}",
+            "fn partial_1<const N:usize>(v: [u32; N]) -> [i32; N] {}",
+        ]
+    )]
+    #[case::remove_const_generics("fn test<const N:usize>(a: i32, v: [u32; N]) {}",
+        vec![
+            "fn default() {}",
+            "fn partial_1(a:i32) {}",
+            "fn partial_2<const N:usize>(a:i32, v: [u32; N]) {}",
+        ]
+    )]
+
     fn clean_generics(#[case] code: &str, #[case] expected: Vec<&str>) {
         let (item_fn, out) = parse_fixture(code);
         let n_args = item_fn.sig.inputs.iter().count();

diff --git a/src/render/mod.rs b/src/render/mod.rs
@@ -284,9 +284,17 @@ fn generics_clean_up<'a>(
     struct Used(std::collections::HashSet<proc_macro2::Ident>);
     impl<'ast> syn::visit::Visit<'ast> for Used {
         fn visit_type_path(&mut self, i: &'ast syn::TypePath) {
-            if i.qself.is_none() && i.path.leading_colon.is_none() && i.path.segments.len() == 1 {
-                self.0
-                    .insert(i.path.segments.first().unwrap().ident.clone());
+            if let Some(id) = i.path.get_ident() {
+                self.0.insert(id.clone());
+            }
+        }
+        fn visit_type_array(&mut self, i: &'ast syn::TypeArray) {
+            let ep = match &i.len {
+                syn::Expr::Path(ep) => ep,
+                _ => return,
+            };
+            if let Some(id) = ep.path.get_ident() {
+                self.0.insert(id.clone());
             }
         }
     }
@@ -298,8 +306,9 @@ fn generics_clean_up<'a>(
         .params
         .into_iter()
         .filter(|p| match p {
-            syn::GenericParam::Type(tp) if !outs.0.contains(&tp.ident) => false,
-            _ => true,
+            syn::GenericParam::Type(syn::TypeParam { ident, .. })
+            | syn::GenericParam::Const(syn::ConstParam { ident, .. }) => outs.0.contains(ident),
+            syn::GenericParam::Lifetime(_) => true,
         })
         .collect();
     result.where_clause.as_mut().map(|mut w| {