refactor: introduce NameGenerator in suggest_name

5 files changed, 192 insertions, 81 deletions

diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_delegate_trait.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_delegate_trait.rs
index c22d19574fb..a55323eb59d 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_delegate_trait.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/generate_delegate_trait.rs
@@ -584,7 +584,7 @@ fn resolve_name_conflicts(
 
             for old_strukt_param in old_strukt_params.generic_params() {
                 // Get old name from `strukt`
-                let mut name = SmolStr::from(match &old_strukt_param {
+                let name = SmolStr::from(match &old_strukt_param {
                     ast::GenericParam::ConstParam(c) => c.name()?.to_string(),
                     ast::GenericParam::LifetimeParam(l) => {
                         l.lifetime()?.lifetime_ident_token()?.to_string()
@@ -593,8 +593,19 @@ fn resolve_name_conflicts(
                 });
 
                 // The new name cannot be conflicted with generics in trait, and the renamed names.
-                name = suggest_name::for_unique_generic_name(&name, old_impl_params);
-                name = suggest_name::for_unique_generic_name(&name, &params);
+                let param_list_to_names = |param_list: &GenericParamList| {
+                    param_list.generic_params().flat_map(|param| match param {
+                        ast::GenericParam::TypeParam(t) => t.name().map(|name| name.to_string()),
+                        p => Some(p.to_string()),
+                    })
+                };
+                let existing_names = param_list_to_names(old_impl_params)
+                    .chain(param_list_to_names(&params))
+                    .collect_vec();
+                let mut name_generator = suggest_name::NameGenerator::new_with_names(
+                    existing_names.iter().map(|s| s.as_str()),
+                );
+                let name = name_generator.suggest_name(&name);
                 match old_strukt_param {
                     ast::GenericParam::ConstParam(c) => {
                         if let Some(const_ty) = c.ty() {
@@ -1213,9 +1224,9 @@ struct S<T> {
     b : B<T>,
 }
 
-impl<T0> Trait<T0> for S<T0> {
-    fn f(&self, a: T0) -> T0 {
-        <B<T0> as Trait<T0>>::f(&self.b, a)
+impl<T1> Trait<T1> for S<T1> {
+    fn f(&self, a: T1) -> T1 {
+        <B<T1> as Trait<T1>>::f(&self.b, a)
     }
 }
 "#,
@@ -1527,12 +1538,12 @@ where
     b : B<T, T1>,
 }
 
-impl<T, T2, T10> Trait<T> for S<T2, T10>
+impl<T, T2, T3> Trait<T> for S<T2, T3>
 where
-    T10: AnotherTrait
+    T3: AnotherTrait
 {
     fn f(&self, a: T) -> T {
-        <B<T2, T10> as Trait<T>>::f(&self.b, a)
+        <B<T2, T3> as Trait<T>>::f(&self.b, a)
     }
 }"#,
         );
@@ -1589,12 +1600,12 @@ where
     b : B<T>,
 }
 
-impl<T, T0> Trait<T> for S<T0>
+impl<T, T2> Trait<T> for S<T2>
 where
-    T0: AnotherTrait
+    T2: AnotherTrait
 {
     fn f(&self, a: T) -> T {
-        <B<T0> as Trait<T>>::f(&self.b, a)
+        <B<T2> as Trait<T>>::f(&self.b, a)
     }
 }"#,
         );
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/introduce_named_generic.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/introduce_named_generic.rs
index a734a6cc2bc..bf6ac1719f3 100644
--- a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/introduce_named_generic.rs
+++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/introduce_named_generic.rs
@@ -1,6 +1,7 @@
 use ide_db::syntax_helpers::suggest_name;
+use itertools::Itertools;
 use syntax::{
-    ast::{self, edit_in_place::GenericParamsOwnerEdit, make, AstNode, HasGenericParams},
+    ast::{self, edit_in_place::GenericParamsOwnerEdit, make, AstNode, HasGenericParams, HasName},
     ted,
 };
 
@@ -33,8 +34,18 @@ pub(crate) fn introduce_named_generic(acc: &mut Assists, ctx: &AssistContext<'_>
             let impl_trait_type = edit.make_mut(impl_trait_type);
             let fn_ = edit.make_mut(fn_);
             let fn_generic_param_list = fn_.get_or_create_generic_param_list();
-            let type_param_name =
-                suggest_name::for_impl_trait_as_generic(&impl_trait_type, &fn_generic_param_list);
+
+            let existing_names = fn_generic_param_list
+                .generic_params()
+                .flat_map(|param| match param {
+                    ast::GenericParam::TypeParam(t) => t.name().map(|name| name.to_string()),
+                    p => Some(p.to_string()),
+                })
+                .collect_vec();
+            let type_param_name = suggest_name::NameGenerator::new_with_names(
+                existing_names.iter().map(|s| s.as_str()),
+            )
+            .for_impl_trait_as_generic(&impl_trait_type);
 
             let type_param = make::type_param(make::name(&type_param_name), Some(type_bound_list))
                 .clone_for_update();
@@ -116,7 +127,7 @@ fn foo<$0B: Bar
         check_assist(
             introduce_named_generic,
             r#"fn foo<B>(bar: $0impl Bar) {}"#,
-            r#"fn foo<B, $0B0: Bar>(bar: B0) {}"#,
+            r#"fn foo<B, $0B1: Bar>(bar: B1) {}"#,
         );
     }
 
@@ -125,7 +136,7 @@ fn foo<$0B: Bar
         check_assist(
             introduce_named_generic,
             r#"fn foo<B, B0, B1, B3>(bar: $0impl Bar) {}"#,
-            r#"fn foo<B, B0, B1, B3, $0B2: Bar>(bar: B2) {}"#,
+            r#"fn foo<B, B0, B1, B3, $0B4: Bar>(bar: B4) {}"#,
         );
     }
 
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/completions/pattern.rs b/src/tools/rust-analyzer/crates/ide-completion/src/completions/pattern.rs
index 2a06fc40175..8f38e02ed76 100644
--- a/src/tools/rust-analyzer/crates/ide-completion/src/completions/pattern.rs
+++ b/src/tools/rust-analyzer/crates/ide-completion/src/completions/pattern.rs
@@ -48,11 +48,12 @@ pub(crate) fn complete_pattern(
 
     // Suggest name only in let-stmt and fn param
     if pattern_ctx.should_suggest_name {
+        let mut name_generator = suggest_name::NameGenerator::new();
         if let Some(suggested) = ctx
             .expected_type
             .as_ref()
             .map(|ty| ty.strip_references())
-            .and_then(|ty| suggest_name::for_type(&ty, ctx.db, ctx.edition))
+            .and_then(|ty| name_generator.for_type(&ty, ctx.db, ctx.edition))
         {
             acc.suggest_name(ctx, &suggested);
         }
diff --git a/src/tools/rust-analyzer/crates/ide-db/src/syntax_helpers/suggest_name.rs b/src/tools/rust-analyzer/crates/ide-db/src/syntax_helpers/suggest_name.rs
index e60deb3bf59..2679cbef61b 100644
--- a/src/tools/rust-analyzer/crates/ide-db/src/syntax_helpers/suggest_name.rs
+++ b/src/tools/rust-analyzer/crates/ide-db/src/syntax_helpers/suggest_name.rs
@@ -1,12 +1,14 @@
 //! This module contains functions to suggest names for expressions, functions and other items
 
+use std::{collections::hash_map::Entry, str::FromStr};
+
 use hir::Semantics;
 use itertools::Itertools;
-use rustc_hash::FxHashSet;
+use rustc_hash::FxHashMap;
 use stdx::to_lower_snake_case;
 use syntax::{
     ast::{self, HasName},
-    match_ast, AstNode, Edition, SmolStr,
+    match_ast, AstNode, Edition, SmolStr, SmolStrBuilder,
 };
 
 use crate::RootDatabase;
@@ -62,71 +64,131 @@ const USELESS_METHODS: &[&str] = &[
     "into_future",
 ];
 
-/// Suggest a name for given type.
+/// Generator for new names
 ///
-/// The function will strip references first, and suggest name from the inner type.
+/// The generator keeps track of existing names and suggests new names that do
+/// not conflict with existing names.
 ///
-/// - If `ty` is an ADT, it will suggest the name of the ADT.
-///   + If `ty` is wrapped in `Box`, `Option` or `Result`, it will suggest the name from the inner type.
-/// - If `ty` is a trait, it will suggest the name of the trait.
-/// - If `ty` is an `impl Trait`, it will suggest the name of the first trait.
+/// The generator will try to resolve conflicts by adding a numeric suffix to
+/// the name, e.g. `a`, `a1`, `a2`, ...
 ///
-/// If the suggested name conflicts with reserved keywords, it will return `None`.
-pub fn for_type(ty: &hir::Type, db: &RootDatabase, edition: Edition) -> Option<String> {
-    let ty = ty.strip_references();
-    name_of_type(&ty, db, edition)
-}
-
-/// Suggest a unique name for generic parameter.
-///
-/// `existing_params` is used to check if the name conflicts with existing
-/// generic parameters.
+/// # Examples
+/// ```rust
+/// let mut generator = NameGenerator::new();
+/// assert_eq!(generator.suggest_name("a"), "a");
+/// assert_eq!(generator.suggest_name("a"), "a1");
 ///
-/// The function checks if the name conflicts with existing generic parameters.
-/// If so, it will try to resolve the conflict by adding a number suffix, e.g.
-/// `T`, `T0`, `T1`, ...
-pub fn for_unique_generic_name(name: &str, existing_params: &ast::GenericParamList) -> SmolStr {
-    let param_names = existing_params
-        .generic_params()
-        .map(|param| match param {
-            ast::GenericParam::TypeParam(t) => t.name().unwrap().to_string(),
-            p => p.to_string(),
-        })
-        .collect::<FxHashSet<_>>();
-    let mut name = name.to_owned();
-    let base_len = name.len();
-    let mut count = 0;
-    while param_names.contains(&name) {
-        name.truncate(base_len);
-        name.push_str(&count.to_string());
-        count += 1;
-    }
-
-    name.into()
+/// assert_eq!(generator.suggest_name("b2"), "b2");
+/// assert_eq!(generator.suggest_name("b"), "b3");
+/// ```
+#[derive(Debug, Default)]
+pub struct NameGenerator {
+    pool: FxHashMap<SmolStr, usize>,
 }
 
-/// Suggest name of impl trait type
-///
-/// `existing_params` is used to check if the name conflicts with existing
-/// generic parameters.
-///
-/// # Current implementation
-///
-/// In current implementation, the function tries to get the name from the first
-/// character of the name for the first type bound.
-///
-/// If the name conflicts with existing generic parameters, it will try to
-/// resolve the conflict with `for_unique_generic_name`.
-pub fn for_impl_trait_as_generic(
-    ty: &ast::ImplTraitType,
-    existing_params: &ast::GenericParamList,
-) -> SmolStr {
-    let c = ty
-        .type_bound_list()
-        .and_then(|bounds| bounds.syntax().text().char_at(0.into()))
-        .unwrap_or('T');
-
-    for_unique_generic_name(c.encode_utf8(&mut [0; 4]), existing_params)
+impl NameGenerator {
+    /// Create a new empty generator
+    pub fn new() -> Self {
+        Self { pool: FxHashMap::default() }
+    }
+
+    /// Create a new generator with existing names. When suggesting a name, it will
+    /// avoid conflicts with existing names.
+    pub fn new_with_names<'a>(existing_names: impl Iterator<Item = &'a str>) -> Self {
+        let mut generator = Self::new();
+        existing_names.for_each(|name| generator.insert(name));
+        generator
+    }
+
+    /// Suggest a name without conflicts. If the name conflicts with existing names,
+    /// it will try to resolve the conflict by adding a numeric suffix.
+    pub fn suggest_name(&mut self, name: &str) -> SmolStr {
+        let (prefix, suffix) = Self::split_numeric_suffix(name);
+        let prefix = SmolStr::new(prefix);
+        let suffix = suffix.unwrap_or(0);
+
+        match self.pool.entry(prefix.clone()) {
+            Entry::Vacant(entry) => {
+                entry.insert(suffix);
+                SmolStr::from_str(name).unwrap()
+            }
+            Entry::Occupied(mut entry) => {
+                let count = entry.get_mut();
+                *count = (*count + 1).max(suffix);
+
+                let mut new_name = SmolStrBuilder::new();
+                new_name.push_str(&prefix);
+                new_name.push_str(count.to_string().as_str());
+                new_name.finish()
+            }
+        }
+    }
+
+    /// Suggest a name for given type.
+    ///
+    /// The function will strip references first, and suggest name from the inner type.
+    ///
+    /// - If `ty` is an ADT, it will suggest the name of the ADT.
+    ///   + If `ty` is wrapped in `Box`, `Option` or `Result`, it will suggest the name from the inner type.
+    /// - If `ty` is a trait, it will suggest the name of the trait.
+    /// - If `ty` is an `impl Trait`, it will suggest the name of the first trait.
+    ///
+    /// If the suggested name conflicts with reserved keywords, it will return `None`.
+    pub fn for_type(
+        &mut self,
+        ty: &hir::Type,
+        db: &RootDatabase,
+        edition: Edition,
+    ) -> Option<SmolStr> {
+        let name = name_of_type(ty, db, edition)?;
+        Some(self.suggest_name(&name))
+    }
+
+    /// Suggest name of impl trait type
+    ///
+    /// # Current implementation
+    ///
+    /// In current implementation, the function tries to get the name from the first
+    /// character of the name for the first type bound.
+    ///
+    /// If the name conflicts with existing generic parameters, it will try to
+    /// resolve the conflict with `for_unique_generic_name`.
+    pub fn for_impl_trait_as_generic(&mut self, ty: &ast::ImplTraitType) -> SmolStr {
+        let c = ty
+            .type_bound_list()
+            .and_then(|bounds| bounds.syntax().text().char_at(0.into()))
+            .unwrap_or('T');
+
+        self.suggest_name(&c.to_string())
+    }
+
+    /// Insert a name into the pool
+    fn insert(&mut self, name: &str) {
+        let (prefix, suffix) = Self::split_numeric_suffix(name);
+        let prefix = SmolStr::new(prefix);
+        let suffix = suffix.unwrap_or(0);
+
+        match self.pool.entry(prefix) {
+            Entry::Vacant(entry) => {
+                entry.insert(suffix);
+            }
+            Entry::Occupied(mut entry) => {
+                let count = entry.get_mut();
+                *count = (*count).max(suffix);
+            }
+        }
+    }
+
+    /// Remove the numeric suffix from the name
+    ///
+    /// # Examples
+    /// `a1b2c3` -> `a1b2c`
+    fn split_numeric_suffix(name: &str) -> (&str, Option<usize>) {
+        let pos =
+            name.rfind(|c: char| !c.is_numeric()).expect("Name cannot be empty or all-numeric");
+        let (prefix, suffix) = name.split_at(pos + 1);
+        (prefix, suffix.parse().ok())
+    }
 }
 
 /// Suggest name of variable for given expression
@@ -290,9 +352,10 @@ fn var_name_from_pat(pat: &ast::Pat) -> Option<ast::Name> {
 
 fn from_type(expr: &ast::Expr, sema: &Semantics<'_, RootDatabase>) -> Option<String> {
     let ty = sema.type_of_expr(expr)?.adjusted();
+    let ty = ty.remove_ref().unwrap_or(ty);
     let edition = sema.scope(expr.syntax())?.krate().edition(sema.db);
 
-    for_type(&ty, sema.db, edition)
+    name_of_type(&ty, sema.db, edition)
 }
 
 fn name_of_type(ty: &hir::Type, db: &RootDatabase, edition: Edition) -> Option<String> {
@@ -925,4 +988,29 @@ fn main() {
             "bar",
         );
     }
+
+    #[test]
+    fn conflicts_with_existing_names() {
+        let mut generator = NameGenerator::new();
+        assert_eq!(generator.suggest_name("a"), "a");
+        assert_eq!(generator.suggest_name("a"), "a1");
+        assert_eq!(generator.suggest_name("a"), "a2");
+        assert_eq!(generator.suggest_name("a"), "a3");
+
+        assert_eq!(generator.suggest_name("b"), "b");
+        assert_eq!(generator.suggest_name("b2"), "b2");
+        assert_eq!(generator.suggest_name("b"), "b3");
+        assert_eq!(generator.suggest_name("b"), "b4");
+        assert_eq!(generator.suggest_name("b3"), "b5");
+
+        // ---------
+        let mut generator = NameGenerator::new_with_names(["a", "b", "b2", "c4"].into_iter());
+        assert_eq!(generator.suggest_name("a"), "a1");
+        assert_eq!(generator.suggest_name("a"), "a2");
+
+        assert_eq!(generator.suggest_name("b"), "b3");
+        assert_eq!(generator.suggest_name("b2"), "b4");
+
+        assert_eq!(generator.suggest_name("c"), "c5");
+    }
 }
diff --git a/src/tools/rust-analyzer/crates/syntax/src/lib.rs b/src/tools/rust-analyzer/crates/syntax/src/lib.rs
index b68374848b9..7fb80b2e7d8 100644
--- a/src/tools/rust-analyzer/crates/syntax/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/syntax/src/lib.rs
@@ -65,7 +65,7 @@ pub use rowan::{
     TokenAtOffset, WalkEvent,
 };
 pub use rustc_lexer::unescape;
-pub use smol_str::{format_smolstr, SmolStr, ToSmolStr};
+pub use smol_str::{format_smolstr, SmolStr, SmolStrBuilder, ToSmolStr};
 
 /// `Parse` is the result of the parsing: a syntax tree and a collection of
 /// errors.