Merge pull request #1655 from dtolnay/assocgroup

Group associated functions by self type

Author: dtolnay

macro/src/expand.rs

@@ -3,6 +3,7 @@ use crate::syntax::attrs::{self, OtherAttrs};
 use crate::syntax::cfg::{CfgExpr, ComputedCfg};
 use crate::syntax::file::Module;
 use crate::syntax::instantiate::{ImplKey, NamedImplKey};
+use crate::syntax::map::OrderedMap;
 use crate::syntax::message::Message;
 use crate::syntax::namespace::Namespace;
 use crate::syntax::qualified::QualifiedName;
@@ -70,6 +71,7 @@ fn expand(ffi: Module, doc: Doc, attrs: OtherAttrs, apis: &[Api], types: &Types)
             Api::Include(_) | Api::Impl(_) => {}
             Api::Struct(strct) => {
                 expanded.extend(expand_struct(strct));
+                expanded.extend(expand_associated_functions(&strct.name.rust, types));
                 hidden.extend(expand_struct_nonempty(strct));
                 hidden.extend(expand_struct_operators(strct));
                 forbid.extend(expand_struct_forbid_drop(strct));
@@ -81,19 +83,24 @@ fn expand(ffi: Module, doc: Doc, attrs: OtherAttrs, apis: &[Api], types: &Types)
                     hidden.extend(expand_extern_shared_struct(ety, &ffi));
                 } else if !types.enums.contains_key(ident) {
                     expanded.extend(expand_cxx_type(ety));
+                    expanded.extend(expand_associated_functions(&ety.name.rust, types));
                     hidden.extend(expand_cxx_type_assert_pinned(ety, types));
                 }
             }
             Api::CxxFunction(efn) => {
-                expanded.extend(expand_cxx_function_shim(efn, types));
+                if efn.self_type().is_none() {
+                    expanded.extend(expand_cxx_function_shim(efn, types));
+                }
             }
             Api::RustType(ety) => {
                 expanded.extend(expand_rust_type_impl(ety));
+                expanded.extend(expand_associated_functions(&ety.name.rust, types));
                 hidden.extend(expand_rust_type_layout(ety, types));
             }
             Api::RustFunction(efn) => hidden.extend(expand_rust_function_shim(efn, types)),
             Api::TypeAlias(alias) => {
                 expanded.extend(expand_type_alias(alias));
+                expanded.extend(expand_associated_functions(&alias.name.rust, types));
                 hidden.extend(expand_type_alias_verify(alias, types));
             }
         }
@@ -586,6 +593,85 @@ fn expand_extern_shared_struct(ety: &ExternType, ffi: &Module) -> TokenStream {
     }
 }
 
+fn expand_associated_functions(self_type: &Ident, types: &Types) -> TokenStream {
+    let Some(functions) = types.associated_fn.get(self_type) else {
+        return TokenStream::new();
+    };
+
+    let resolve = types.resolve(self_type);
+    let self_type_cfg_attrs = resolve.attrs.cfg();
+    let elided_lifetime = Lifetime::new("'_", Span::call_site());
+    let mut group_by_lifetimes = OrderedMap::new();
+    let mut tokens = TokenStream::new();
+
+    for efn in functions {
+        match efn.lang {
+            Lang::Cxx | Lang::CxxUnwind => {}
+            Lang::Rust => continue,
+        }
+        let mut impl_lifetimes = Vec::new();
+        let mut self_type_lifetimes = Vec::new();
+        let self_lt_token;
+        let self_gt_token;
+        match &efn.kind {
+            FnKind::Method(receiver) if receiver.ty.generics.lt_token.is_some() => {
+                for lifetime in &receiver.ty.generics.lifetimes {
+                    if lifetime.ident != "_"
+                        && efn
+                            .generics
+                            .lifetimes()
+                            .any(|param| param.lifetime == *lifetime)
+                    {
+                        impl_lifetimes.push(lifetime);
+                    }
+                    self_type_lifetimes.push(lifetime);
+                }
+                self_lt_token = receiver.ty.generics.lt_token;
+                self_gt_token = receiver.ty.generics.gt_token;
+            }
+            _ => {
+                self_type_lifetimes.resize(resolve.generics.lifetimes.len(), &elided_lifetime);
+                self_lt_token = resolve.generics.lt_token;
+                self_gt_token = resolve.generics.gt_token;
+            }
+        }
+        if efn.undeclared_lifetimes().is_empty()
+            && self_type_lifetimes.len() == resolve.generics.lifetimes.len()
+        {
+            group_by_lifetimes
+                .entry((impl_lifetimes, self_type_lifetimes))
+                .or_insert_with(Vec::new)
+                .push(efn);
+        } else {
+            let impl_token = Token![impl](efn.name.rust.span());
+            let impl_lt_token = efn.generics.lt_token;
+            let impl_gt_token = efn.generics.gt_token;
+            let self_type = efn.self_type().unwrap();
+            let function = expand_cxx_function_shim(efn, types);
+            tokens.extend(quote! {
+                #self_type_cfg_attrs
+                #impl_token #impl_lt_token #(#impl_lifetimes),* #impl_gt_token #self_type #self_lt_token #(#self_type_lifetimes),* #self_gt_token {
+                    #function
+                }
+            });
+        }
+    }
+
+    for ((impl_lifetimes, self_type_lifetimes), functions) in &group_by_lifetimes {
+        let functions = functions
+            .iter()
+            .map(|efn| expand_cxx_function_shim(efn, types));
+        tokens.extend(quote! {
+            #self_type_cfg_attrs
+            impl <#(#impl_lifetimes),*> #self_type <#(#self_type_lifetimes),*> {
+                #(#functions)*
+            }
+        });
+    }
+
+    tokens
+}
+
 fn expand_cxx_function_decl(efn: &ExternFn, types: &Types) -> TokenStream {
     let receiver = efn.receiver().into_iter().map(|receiver| {
         if types.is_considered_improper_ctype(&receiver.ty) {
@@ -897,7 +983,6 @@ fn expand_cxx_function_shim(efn: &ExternFn, types: &Types) -> TokenStream {
         Some(self_type) => {
             let elided_generics;
             let resolve = types.resolve(self_type);
-            let self_type_cfg_attrs = resolve.attrs.cfg();
             let self_type_generics = match &efn.kind {
                 FnKind::Method(receiver) if receiver.ty.generics.lt_token.is_some() => {
                     &receiver.ty.generics
@@ -926,21 +1011,15 @@ fn expand_cxx_function_shim(efn: &ExternFn, types: &Types) -> TokenStream {
                     self_type_lifetimes.insert(lifetime);
                 }
             }
-            let impl_lifetimes = generics
-                .lifetimes()
-                .filter(|param| self_type_lifetimes.contains(&param.lifetime));
             let fn_lifetimes = generics
                 .lifetimes()
                 .filter(|param| !self_type_lifetimes.contains(&param.lifetime));
             let lt_token = generics.lt_token;
             let gt_token = generics.gt_token;
             quote_spanned! {ident.span()=>
-                #self_type_cfg_attrs
-                impl #lt_token #(#impl_lifetimes),* #gt_token #self_type #self_type_generics {
-                    #doc
-                    #all_attrs
-                    #visibility #unsafety #fn_token #ident #lt_token #(#fn_lifetimes),* #gt_token #arg_list #ret #fn_body
-                }
+                #doc
+                #all_attrs
+                #visibility #unsafety #fn_token #ident #lt_token #(#fn_lifetimes),* #gt_token #arg_list #ret #fn_body
             }
         }
     }

macro/src/lib.rs

@@ -13,6 +13,7 @@
     clippy::nonminimal_bool,
     clippy::redundant_else,
     clippy::ref_option,
+    clippy::similar_names,
     clippy::single_match_else,
     clippy::struct_field_names,
     clippy::too_many_arguments,

syntax/types.rs

@@ -11,7 +11,7 @@ use crate::syntax::trivial::{self, TrivialReason};
 use crate::syntax::unpin::{self, UnpinReason};
 use crate::syntax::visit::{self, Visit};
 use crate::syntax::{
-    toposort, Api, Atom, Enum, ExternType, Impl, Lifetimes, Pair, Struct, Type, TypeAlias,
+    toposort, Api, Atom, Enum, ExternFn, ExternType, Impl, Lifetimes, Pair, Struct, Type, TypeAlias,
 };
 use indexmap::map::Entry;
 use proc_macro2::Ident;
@@ -30,6 +30,8 @@ pub(crate) struct Types<'a> {
     pub required_unpin: UnorderedMap<&'a Ident, UnpinReason<'a>>,
     pub impls: OrderedMap<ImplKey<'a>, ConditionalImpl<'a>>,
     pub resolutions: UnorderedMap<&'a Ident, Resolution<'a>>,
+    #[cfg_attr(not(proc_macro), expect(dead_code))]
+    pub associated_fn: UnorderedMap<&'a Ident, Vec<&'a ExternFn>>,
     pub struct_improper_ctypes: UnorderedSet<&'a Ident>,
     pub toposorted_structs: Vec<&'a Struct>,
 }
@@ -53,6 +55,7 @@ impl<'a> Types<'a> {
         let mut untrusted = UnorderedMap::new();
         let mut impls = OrderedMap::new();
         let mut resolutions = UnorderedMap::new();
+        let mut associated_fn = UnorderedMap::new();
         let struct_improper_ctypes = UnorderedSet::new();
         let toposorted_structs = Vec::new();
 
@@ -177,6 +180,12 @@ impl<'a> Types<'a> {
                     // Note: duplication of the C++ name is fine because C++ has
                     // function overloading.
                     let self_type = efn.self_type();
+                    if let Some(self_type) = self_type {
+                        associated_fn
+                            .entry(self_type)
+                            .or_insert_with(Vec::new)
+                            .push(efn);
+                    }
                     if !self_type.is_some_and(|self_type| self_type == "Self")
                         && !function_names.insert((self_type, &efn.name.rust))
                     {
@@ -253,6 +262,7 @@ impl<'a> Types<'a> {
             required_unpin,
             impls,
             resolutions,
+            associated_fn,
             struct_improper_ctypes,
             toposorted_structs,
         };