rust-analyzer/crates/ide_db/src/helpers/import_assets.rs

//! Look up accessible paths for items.
use hir::{
    AsAssocItem, AssocItem, AssocItemContainer, Crate, ItemInNs, MacroDef, ModPath, Module,
    ModuleDef, Name, PathResolution, PrefixKind, ScopeDef, Semantics, SemanticsScope, Type,
};
use itertools::Itertools;
use rustc_hash::FxHashSet;
use syntax::{ast, AstNode};

use crate::{
    items_locator::{self, AssocItemSearch, DEFAULT_QUERY_SEARCH_LIMIT},
    RootDatabase,
};

use super::item_name;

#[derive(Debug)]
pub enum ImportCandidate {
    // A path, qualified (`std::collections::HashMap`) or not (`HashMap`).
    Path(PathImportCandidate),
    /// A trait associated function (with no self parameter) or associated constant.
    /// For 'test_mod::TestEnum::test_function', `ty` is the `test_mod::TestEnum` expression type
    /// and `name` is the `test_function`
    TraitAssocItem(TraitImportCandidate),
    /// A trait method with self parameter.
    /// For 'test_enum.test_method()', `ty` is the `test_enum` expression type
    /// and `name` is the `test_method`
    TraitMethod(TraitImportCandidate),
}

#[derive(Debug)]
pub struct TraitImportCandidate {
    pub receiver_ty: Type,
    pub name: NameToImport,
}

#[derive(Debug)]
pub struct PathImportCandidate {
    pub qualifier: Qualifier,
    pub name: NameToImport,
}

#[derive(Debug)]
pub enum Qualifier {
    Absent,
    FirstSegmentUnresolved(ast::NameRef, ModPath),
}

#[derive(Debug)]
pub enum NameToImport {
    Exact(String),
    Fuzzy(String),
}

impl NameToImport {
    pub fn text(&self) -> &str {
        match self {
            NameToImport::Exact(text) => text.as_str(),
            NameToImport::Fuzzy(text) => text.as_str(),
        }
    }
}

#[derive(Debug)]
pub struct ImportAssets<'a> {
    import_candidate: ImportCandidate,
    module_with_candidate: Module,
    scope: SemanticsScope<'a>,
}

impl<'a> ImportAssets<'a> {
    pub fn for_method_call(
        method_call: &ast::MethodCallExpr,
        sema: &'a Semantics<RootDatabase>,
    ) -> Option<Self> {
        let scope = sema.scope(method_call.syntax());
        Some(Self {
            import_candidate: ImportCandidate::for_method_call(sema, method_call)?,
            module_with_candidate: scope.module()?,
            scope,
        })
    }

    pub fn for_exact_path(
        fully_qualified_path: &ast::Path,
        sema: &'a Semantics<RootDatabase>,
    ) -> Option<Self> {
        let syntax_under_caret = fully_qualified_path.syntax();
        if syntax_under_caret.ancestors().find_map(ast::Use::cast).is_some() {
            return None;
        }
        let scope = sema.scope(syntax_under_caret);
        Some(Self {
            import_candidate: ImportCandidate::for_regular_path(sema, fully_qualified_path)?,
            module_with_candidate: scope.module()?,
            scope,
        })
    }

    pub fn for_fuzzy_path(
        module_with_candidate: Module,
        qualifier: Option<ast::Path>,
        fuzzy_name: String,
        sema: &Semantics<RootDatabase>,
        scope: SemanticsScope<'a>,
    ) -> Option<Self> {
        Some(Self {
            import_candidate: ImportCandidate::for_fuzzy_path(qualifier, fuzzy_name, sema)?,
            module_with_candidate,
            scope,
        })
    }

    pub fn for_fuzzy_method_call(
        module_with_method_call: Module,
        receiver_ty: Type,
        fuzzy_method_name: String,
        scope: SemanticsScope<'a>,
    ) -> Option<Self> {
        Some(Self {
            import_candidate: ImportCandidate::TraitMethod(TraitImportCandidate {
                receiver_ty,
                name: NameToImport::Fuzzy(fuzzy_method_name),
            }),
            module_with_candidate: module_with_method_call,
            scope,
        })
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct LocatedImport {
    pub import_path: ModPath,
    pub item_to_import: ItemInNs,
    pub original_item: ItemInNs,
    pub original_path: Option<ModPath>,
}

impl LocatedImport {
    pub fn new(
        import_path: ModPath,
        item_to_import: ItemInNs,
        original_item: ItemInNs,
        original_path: Option<ModPath>,
    ) -> Self {
        Self { import_path, item_to_import, original_item, original_path }
    }

    pub fn original_item_name(&self, db: &RootDatabase) -> Option<Name> {
        match self.original_item {
            ItemInNs::Types(module_def_id) | ItemInNs::Values(module_def_id) => {
                ModuleDef::from(module_def_id).name(db)
            }
            ItemInNs::Macros(macro_def_id) => MacroDef::from(macro_def_id).name(db),
        }
    }
}

impl<'a> ImportAssets<'a> {
    pub fn import_candidate(&self) -> &ImportCandidate {
        &self.import_candidate
    }

    pub fn search_for_imports(
        &self,
        sema: &Semantics<RootDatabase>,
        prefix_kind: PrefixKind,
    ) -> Vec<LocatedImport> {
        let _p = profile::span("import_assets::search_for_imports");
        self.search_for(sema, Some(prefix_kind))
    }

    /// This may return non-absolute paths if a part of the returned path is already imported into scope.
    pub fn search_for_relative_paths(&self, sema: &Semantics<RootDatabase>) -> Vec<LocatedImport> {
        let _p = profile::span("import_assets::search_for_relative_paths");
        self.search_for(sema, None)
    }

    fn search_for(
        &self,
        sema: &Semantics<RootDatabase>,
        prefixed: Option<PrefixKind>,
    ) -> Vec<LocatedImport> {
        let items_with_candidate_name = match self.name_to_import() {
            NameToImport::Exact(exact_name) => items_locator::with_for_exact_name(
                sema,
                self.module_with_candidate.krate(),
                exact_name.clone(),
            ),
            // FIXME: ideally, we should avoid using `fst` for seacrhing trait imports for assoc items:
            // instead, we need to look up all trait impls for a certain struct and search through them only
            // see https://github.com/rust-analyzer/rust-analyzer/pull/7293#issuecomment-761585032
            // and https://rust-lang.zulipchat.com/#narrow/stream/185405-t-compiler.2Fwg-rls-2.2E0/topic/Blanket.20trait.20impls.20lookup
            // for the details
            NameToImport::Fuzzy(fuzzy_name) => {
                let (assoc_item_search, limit) = if self.import_candidate.is_trait_candidate() {
                    (AssocItemSearch::AssocItemsOnly, None)
                } else {
                    (AssocItemSearch::Include, Some(DEFAULT_QUERY_SEARCH_LIMIT))
                };

                items_locator::with_similar_name(
                    sema,
                    self.module_with_candidate.krate(),
                    fuzzy_name.clone(),
                    assoc_item_search,
                    limit,
                )
            }
        };

        let scope_definitions = self.scope_definitions();
        self.applicable_defs(sema.db, prefixed, items_with_candidate_name)
            .into_iter()
            .filter(|import| import.import_path.len() > 1)
            .filter(|import| !scope_definitions.contains(&ScopeDef::from(import.item_to_import)))
            .sorted_by_key(|import| import.import_path.clone())
            .collect()
    }

    fn scope_definitions(&self) -> FxHashSet<ScopeDef> {
        let mut scope_definitions = FxHashSet::default();
        self.scope.process_all_names(&mut |_, scope_def| {
            scope_definitions.insert(scope_def);
        });
        scope_definitions
    }

    fn name_to_import(&self) -> &NameToImport {
        match &self.import_candidate {
            ImportCandidate::Path(candidate) => &candidate.name,
            ImportCandidate::TraitAssocItem(candidate)
            | ImportCandidate::TraitMethod(candidate) => &candidate.name,
        }
    }

    fn applicable_defs(
        &self,
        db: &RootDatabase,
        prefixed: Option<PrefixKind>,
        items_with_candidate_name: FxHashSet<ItemInNs>,
    ) -> FxHashSet<LocatedImport> {
        let _p = profile::span("import_assets::applicable_defs");
        let current_crate = self.module_with_candidate.krate();

        let mod_path = |item| {
            get_mod_path(db, item_for_path_search(db, item)?, &self.module_with_candidate, prefixed)
        };

        match &self.import_candidate {
            ImportCandidate::Path(path_candidate) => {
                path_applicable_imports(db, path_candidate, mod_path, items_with_candidate_name)
            }
            ImportCandidate::TraitAssocItem(trait_candidate) => trait_applicable_items(
                db,
                current_crate,
                trait_candidate,
                true,
                mod_path,
                items_with_candidate_name,
            ),
            ImportCandidate::TraitMethod(trait_candidate) => trait_applicable_items(
                db,
                current_crate,
                trait_candidate,
                false,
                mod_path,
                items_with_candidate_name,
            ),
        }
    }
}

fn path_applicable_imports(
    db: &RootDatabase,
    path_candidate: &PathImportCandidate,
    mod_path: impl Fn(ItemInNs) -> Option<ModPath> + Copy,
    items_with_candidate_name: FxHashSet<ItemInNs>,
) -> FxHashSet<LocatedImport> {
    let _p = profile::span("import_assets::path_applicable_imports");

    let (unresolved_first_segment, unresolved_qualifier) = match &path_candidate.qualifier {
        Qualifier::Absent => {
            return items_with_candidate_name
                .into_iter()
                .filter_map(|item| {
                    Some(LocatedImport::new(mod_path(item)?, item, item, mod_path(item)))
                })
                .collect();
        }
        Qualifier::FirstSegmentUnresolved(first_segment, qualifier) => {
            (first_segment.to_string(), qualifier.to_string())
        }
    };

    items_with_candidate_name
        .into_iter()
        .filter_map(|item| {
            import_for_item(db, mod_path, &unresolved_first_segment, &unresolved_qualifier, item)
        })
        .collect()
}

fn import_for_item(
    db: &RootDatabase,
    mod_path: impl Fn(ItemInNs) -> Option<ModPath>,
    unresolved_first_segment: &str,
    unresolved_qualifier: &str,
    original_item: ItemInNs,
) -> Option<LocatedImport> {
    let _p = profile::span("import_assets::import_for_item");

    let original_item_candidate = item_for_path_search(db, original_item)?;
    let import_path_candidate = mod_path(original_item_candidate)?;
    let import_path_string = import_path_candidate.to_string();

    let expected_import_end = if item_as_assoc(db, original_item).is_some() {
        unresolved_qualifier.to_string()
    } else {
        format!("{}::{}", unresolved_qualifier, item_name(db, original_item)?)
    };
    if !import_path_string.contains(unresolved_first_segment)
        || !import_path_string.ends_with(&expected_import_end)
    {
        return None;
    }

    let segment_import =
        find_import_for_segment(db, original_item_candidate, &unresolved_first_segment)?;
    let trait_item_to_import = item_as_assoc(db, original_item)
        .and_then(|assoc| assoc.containing_trait(db))
        .map(|trait_| ItemInNs::from(ModuleDef::from(trait_)));
    Some(match (segment_import == original_item_candidate, trait_item_to_import) {
        (true, Some(_)) => {
            // FIXME we should be able to import both the trait and the segment,
            // but it's unclear what to do with overlapping edits (merge imports?)
            // especially in case of lazy completion edit resolutions.
            return None;
        }
        (false, Some(trait_to_import)) => LocatedImport::new(
            mod_path(trait_to_import)?,
            trait_to_import,
            original_item,
            mod_path(original_item),
        ),
        (true, None) => LocatedImport::new(
            import_path_candidate,
            original_item_candidate,
            original_item,
            mod_path(original_item),
        ),
        (false, None) => LocatedImport::new(
            mod_path(segment_import)?,
            segment_import,
            original_item,
            mod_path(original_item),
        ),
    })
}

fn item_for_path_search(db: &RootDatabase, item: ItemInNs) -> Option<ItemInNs> {
    Some(match item {
        ItemInNs::Types(_) | ItemInNs::Values(_) => match item_as_assoc(db, item) {
            Some(assoc_item) => match assoc_item.container(db) {
                AssocItemContainer::Trait(trait_) => ItemInNs::from(ModuleDef::from(trait_)),
                AssocItemContainer::Impl(impl_) => {
                    ItemInNs::from(ModuleDef::from(impl_.target_ty(db).as_adt()?))
                }
            },
            None => item,
        },
        ItemInNs::Macros(_) => item,
    })
}

fn find_import_for_segment(
    db: &RootDatabase,
    original_item: ItemInNs,
    unresolved_first_segment: &str,
) -> Option<ItemInNs> {
    let segment_is_name = item_name(db, original_item)
        .map(|name| name.to_string() == unresolved_first_segment)
        .unwrap_or(false);

    Some(if segment_is_name {
        original_item
    } else {
        let matching_module =
            module_with_segment_name(db, &unresolved_first_segment, original_item)?;
        ItemInNs::from(ModuleDef::from(matching_module))
    })
}

fn module_with_segment_name(
    db: &RootDatabase,
    segment_name: &str,
    candidate: ItemInNs,
) -> Option<Module> {
    let mut current_module = match candidate {
        ItemInNs::Types(module_def_id) => ModuleDef::from(module_def_id).module(db),
        ItemInNs::Values(module_def_id) => ModuleDef::from(module_def_id).module(db),
        ItemInNs::Macros(macro_def_id) => MacroDef::from(macro_def_id).module(db),
    };
    while let Some(module) = current_module {
        if let Some(module_name) = module.name(db) {
            if module_name.to_string() == segment_name {
                return Some(module);
            }
        }
        current_module = module.parent(db);
    }
    None
}

fn trait_applicable_items(
    db: &RootDatabase,
    current_crate: Crate,
    trait_candidate: &TraitImportCandidate,
    trait_assoc_item: bool,
    mod_path: impl Fn(ItemInNs) -> Option<ModPath>,
    items_with_candidate_name: FxHashSet<ItemInNs>,
) -> FxHashSet<LocatedImport> {
    let _p = profile::span("import_assets::trait_applicable_items");
    let mut required_assoc_items = FxHashSet::default();

    let trait_candidates = items_with_candidate_name
        .into_iter()
        .filter_map(|input| item_as_assoc(db, input))
        .filter_map(|assoc| {
            let assoc_item_trait = assoc.containing_trait(db)?;
            required_assoc_items.insert(assoc);
            Some(assoc_item_trait.into())
        })
        .collect();

    let mut located_imports = FxHashSet::default();

    if trait_assoc_item {
        trait_candidate.receiver_ty.iterate_path_candidates(
            db,
            current_crate,
            &trait_candidates,
            None,
            |_, assoc| {
                if required_assoc_items.contains(&assoc) {
                    if let AssocItem::Function(f) = assoc {
                        if f.self_param(db).is_some() {
                            return None;
                        }
                    }

                    let item = ItemInNs::from(ModuleDef::from(assoc.containing_trait(db)?));
                    let original_item = assoc_to_item(assoc);
                    located_imports.insert(LocatedImport::new(
                        mod_path(item)?,
                        item,
                        original_item,
                        mod_path(original_item),
                    ));
                }
                None::<()>
            },
        )
    } else {
        trait_candidate.receiver_ty.iterate_method_candidates(
            db,
            current_crate,
            &trait_candidates,
            None,
            |_, function| {
                let assoc = function.as_assoc_item(db)?;
                if required_assoc_items.contains(&assoc) {
                    let item = ItemInNs::from(ModuleDef::from(assoc.containing_trait(db)?));
                    let original_item = assoc_to_item(assoc);
                    located_imports.insert(LocatedImport::new(
                        mod_path(item)?,
                        item,
                        original_item,
                        mod_path(original_item),
                    ));
                }
                None::<()>
            },
        )
    };

    located_imports
}

fn assoc_to_item(assoc: AssocItem) -> ItemInNs {
    match assoc {
        AssocItem::Function(f) => ItemInNs::from(ModuleDef::from(f)),
        AssocItem::Const(c) => ItemInNs::from(ModuleDef::from(c)),
        AssocItem::TypeAlias(t) => ItemInNs::from(ModuleDef::from(t)),
    }
}

fn get_mod_path(
    db: &RootDatabase,
    item_to_search: ItemInNs,
    module_with_candidate: &Module,
    prefixed: Option<PrefixKind>,
) -> Option<ModPath> {
    if let Some(prefix_kind) = prefixed {
        module_with_candidate.find_use_path_prefixed(db, item_to_search, prefix_kind)
    } else {
        module_with_candidate.find_use_path(db, item_to_search)
    }
}

impl ImportCandidate {
    fn for_method_call(
        sema: &Semantics<RootDatabase>,
        method_call: &ast::MethodCallExpr,
    ) -> Option<Self> {
        match sema.resolve_method_call(method_call) {
            Some(_) => None,
            None => Some(Self::TraitMethod(TraitImportCandidate {
                receiver_ty: sema.type_of_expr(&method_call.receiver()?)?,
                name: NameToImport::Exact(method_call.name_ref()?.to_string()),
            })),
        }
    }

    fn for_regular_path(sema: &Semantics<RootDatabase>, path: &ast::Path) -> Option<Self> {
        if sema.resolve_path(path).is_some() {
            return None;
        }
        path_import_candidate(
            sema,
            path.qualifier(),
            NameToImport::Exact(path.segment()?.name_ref()?.to_string()),
        )
    }

    fn for_fuzzy_path(
        qualifier: Option<ast::Path>,
        fuzzy_name: String,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        path_import_candidate(sema, qualifier, NameToImport::Fuzzy(fuzzy_name))
    }

    fn is_trait_candidate(&self) -> bool {
        matches!(self, ImportCandidate::TraitAssocItem(_) | ImportCandidate::TraitMethod(_))
    }
}

fn path_import_candidate(
    sema: &Semantics<RootDatabase>,
    qualifier: Option<ast::Path>,
    name: NameToImport,
) -> Option<ImportCandidate> {
    Some(match qualifier {
        Some(qualifier) => match sema.resolve_path(&qualifier) {
            None => {
                let qualifier_start =
                    qualifier.syntax().descendants().find_map(ast::NameRef::cast)?;
                let qualifier_start_path =
                    qualifier_start.syntax().ancestors().find_map(ast::Path::cast)?;
                if sema.resolve_path(&qualifier_start_path).is_none() {
                    ImportCandidate::Path(PathImportCandidate {
                        qualifier: Qualifier::FirstSegmentUnresolved(
                            qualifier_start,
                            ModPath::from_src_unhygienic(qualifier)?,
                        ),
                        name,
                    })
                } else {
                    return None;
                }
            }
            Some(PathResolution::Def(ModuleDef::Adt(assoc_item_path))) => {
                ImportCandidate::TraitAssocItem(TraitImportCandidate {
                    receiver_ty: assoc_item_path.ty(sema.db),
                    name,
                })
            }
            Some(_) => return None,
        },
        None => ImportCandidate::Path(PathImportCandidate { qualifier: Qualifier::Absent, name }),
    })
}

fn item_as_assoc(db: &RootDatabase, item: ItemInNs) -> Option<AssocItem> {
    item.as_module_def_id()
        .and_then(|module_def_id| ModuleDef::from(module_def_id).as_assoc_item(db))
}