From aef16300f7cfa79de3bd4797d0196d13445f4761 Mon Sep 17 00:00:00 2001
From: Christofer Nolander <christofer.nolander@gmail.com>
Date: Sat, 21 May 2022 00:00:47 +0200
Subject: [PATCH] Order auto-imports by relevance

This first attempt prefers items that are closer to the module they are
imported in.
---
 .../ide-assists/src/handlers/auto_import.rs   | 142 +++++++++++++++++-
 1 file changed, 140 insertions(+), 2 deletions(-)

diff --git a/crates/ide-assists/src/handlers/auto_import.rs b/crates/ide-assists/src/handlers/auto_import.rs
index 0a0dafb35e..3316b71de2 100644
--- a/crates/ide-assists/src/handlers/auto_import.rs
+++ b/crates/ide-assists/src/handlers/auto_import.rs
@@ -1,7 +1,10 @@
+use std::cmp::Reverse;
+
+use hir::{db::HirDatabase, Module};
 use ide_db::{
     helpers::mod_path_to_ast,
     imports::{
-        import_assets::{ImportAssets, ImportCandidate},
+        import_assets::{ImportAssets, ImportCandidate, LocatedImport},
         insert_use::{insert_use, ImportScope},
     },
 };
@@ -107,6 +110,10 @@ pub(crate) fn auto_import(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
 
     // we aren't interested in different namespaces
     proposed_imports.dedup_by(|a, b| a.import_path == b.import_path);
+
+    // prioritize more relevant imports
+    proposed_imports.sort_by_key(|import| Reverse(relevance_score(ctx, import)));
+
     for import in proposed_imports {
         acc.add_group(
             &group_label,
@@ -158,11 +165,142 @@ fn group_label(import_candidate: &ImportCandidate) -> GroupLabel {
     GroupLabel(name)
 }
 
+/// Determine how relevant a given import is in the current context. Higher scores are more
+/// relevant.
+fn relevance_score(ctx: &AssistContext, import: &LocatedImport) -> i32 {
+    let mut score = 0;
+
+    let db = ctx.db();
+
+    let item_module = match import.item_to_import {
+        hir::ItemInNs::Types(item) | hir::ItemInNs::Values(item) => item.module(db),
+        hir::ItemInNs::Macros(makro) => Some(makro.module(db)),
+    };
+
+    let current_node = match ctx.covering_element() {
+        NodeOrToken::Node(node) => Some(node),
+        NodeOrToken::Token(token) => token.parent(),
+    };
+
+    if let Some(module) = item_module.as_ref() {
+        if module.krate().is_builtin(db) {
+            // prefer items builtin to the language
+            score += 5;
+        }
+    }
+
+    match item_module.zip(current_node) {
+        // get the distance between the modules (prefer items that are more local)
+        Some((item_module, current_node)) => {
+            let current_module = ctx.sema.scope(&current_node).unwrap().module();
+            score -= module_distance_hueristic(&current_module, &item_module, db) as i32;
+        }
+
+        // could not find relevant modules, so just use the length of the path as an estimate
+        None => return -(2 * import.import_path.len() as i32),
+    }
+
+    score
+}
+
+/// A heuristic that gives a higher score to modules that are more separated.
+fn module_distance_hueristic(current: &Module, item: &Module, db: &dyn HirDatabase) -> usize {
+    let mut current_path = current.path_to_root(db);
+    let mut item_path = item.path_to_root(db);
+
+    current_path.reverse();
+    item_path.reverse();
+
+    let mut i = 0;
+    while i < current_path.len() && i < item_path.len() {
+        if current_path[i] == item_path[i] {
+            i += 1
+        } else {
+            break;
+        }
+    }
+
+    let distinct_distance = current_path.len() + item_path.len();
+    let common_prefix = 2 * i;
+
+    // prefer builtin crates and items within the same crate
+    let crate_boundary_cost =
+        if item.krate().is_builtin(db) || current.krate() == item.krate() { 0 } else { 4 };
+
+    distinct_distance - common_prefix + crate_boundary_cost
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    use crate::tests::{check_assist, check_assist_not_applicable, check_assist_target};
+    use hir::Semantics;
+    use ide_db::{
+        assists::AssistResolveStrategy,
+        base_db::{fixture::WithFixture, FileRange},
+        RootDatabase,
+    };
+
+    use crate::tests::{
+        check_assist, check_assist_not_applicable, check_assist_target, TEST_CONFIG,
+    };
+
+    fn check_auto_import_order(before: &str, order: &[&str]) {
+        let (db, file_id, range_or_offset) = RootDatabase::with_range_or_offset(before);
+        let frange = FileRange { file_id, range: range_or_offset.into() };
+
+        let sema = Semantics::new(&db);
+        let config = TEST_CONFIG;
+        let ctx = AssistContext::new(sema, &config, frange);
+        let mut acc = Assists::new(&ctx, AssistResolveStrategy::All);
+        auto_import(&mut acc, &ctx);
+        let assists = acc.finish();
+
+        let labels = assists.iter().map(|assist| assist.label.to_string()).collect::<Vec<_>>();
+
+        assert_eq!(labels, order);
+    }
+
+    #[test]
+    fn prefer_shorter_paths() {
+        let before = r"
+            //- /main.rs crate:main deps:foo,bar
+            HashMap$0::new();
+
+            //- /lib.rs crate:foo
+            pub mod collections { pub struct HashMap; }
+
+            //- /lib.rs crate:bar
+            pub mod collections { pub mod hash_map { pub struct HashMap; } }
+        ";
+
+        check_auto_import_order(
+            before,
+            &["Import `foo::collections::HashMap`", "Import `bar::collections::hash_map::HashMap`"],
+        )
+    }
+
+    #[test]
+    fn prefer_same_crate() {
+        let before = r"
+            //- /main.rs crate:main deps:foo
+            HashMap$0::new();
+
+            mod collections {
+                pub mod hash_map {
+                    pub struct HashMap;
+                }
+            }
+
+            //- /lib.rs crate:foo
+            pub struct HashMap;
+        ";
+
+        check_auto_import_order(
+            before,
+            &["Import `collections::hash_map::HashMap`", "Import `foo::HashMap`"],
+        )
+    }
 
     #[test]
     fn not_applicable_if_scope_inside_macro() {