feat: let the LSP import code action insert into existing use statements

tooling/lsp/src/lib.rs

@@ -70,6 +70,7 @@ mod solver;
 mod tests;
 mod trait_impl_method_stub_generator;
 mod types;
+mod use_segment_positions;
 mod utils;
 mod visibility;
 

tooling/lsp/src/requests/code_action.rs

@@ -22,7 +22,7 @@
     ParsedModule,
 };
 
-use crate::{utils, LspState};
+use crate::{use_segment_positions::UseSegmentPositions, utils, LspState};
 
 use super::{process_request, to_lsp_location};
 
@@ -82,6 +82,7 @@
     nesting: usize,
     /// The line where an auto_import must be inserted
     auto_import_line: usize,
+    use_segment_positions: UseSegmentPositions,
     /// Text edits for the "Remove all unused imports" code action
     unused_imports_text_edits: Vec<TextEdit>,
     code_actions: Vec<CodeAction>,
@@ -121,6 +122,7 @@
             interner,
             nesting: 0,
             auto_import_line: 0,
+            use_segment_positions: UseSegmentPositions::default(),
             unused_imports_text_edits: vec![],
             code_actions: vec![],
         }
@@ -166,7 +168,7 @@
 
         CodeAction {
             title,
             kind: Some(CodeActionKind::QUICKFIX),
             diagnostics: None,
             edit: Some(workspace_edit),
             command: None,
@@ -184,10 +186,11 @@
 
 impl<'a> Visitor for CodeActionFinder<'a> {
     fn visit_item(&mut self, item: &Item) -> bool {
-        if let ItemKind::Import(..) = &item.kind {
+        if let ItemKind::Import(use_tree, _) = &item.kind {
             if let Some(lsp_location) = to_lsp_location(self.files, self.file, item.span) {
                 self.auto_import_line = (lsp_location.range.end.line + 1) as usize;
             }
+            self.use_segment_positions.add(use_tree);
         }
 
         self.includes_span(item.span)

tooling/lsp/src/requests/code_action/import_or_qualify.rs

@@ -1,4 +1,4 @@
-use lsp_types::{Position, Range, TextEdit};
+use lsp_types::TextEdit;
 use noirc_errors::Location;
 use noirc_frontend::{
     ast::{Ident, Path},
@@ -8,6 +8,9 @@ use noirc_frontend::{
 use crate::{
     byte_span_to_range,
     modules::{relative_module_full_path, relative_module_id_path},
+    use_segment_positions::{
+        use_completion_item_additional_text_edits, UseCompletionItemAdditionTextEditsRequest,
+    },
 };
 
 use super::CodeActionFinder;
@@ -82,25 +85,21 @@ impl<'a> CodeActionFinder<'a> {
     }
 
     fn push_import_code_action(&mut self, full_path: &str) {
-        let line = self.auto_import_line as u32;
-        let character = (self.nesting * 4) as u32;
-        let indent = " ".repeat(self.nesting * 4);
-        let mut newlines = "\n";
-
-        // If the line we are inserting into is not an empty line, insert an extra line to make some room
-        if let Some(line_text) = self.lines.get(line as usize) {
-            if !line_text.trim().is_empty() {
-                newlines = "\n\n";
-            }
-        }
-
         let title = format!("Import {}", full_path);
-        let text_edit = TextEdit {
-            range: Range { start: Position { line, character }, end: Position { line, character } },
-            new_text: format!("use {};{}{}", full_path, newlines, indent),
-        };
 
-        let code_action = self.new_quick_fix(title, text_edit);
+        let text_edits = use_completion_item_additional_text_edits(
+            UseCompletionItemAdditionTextEditsRequest {
+                full_path,
+                files: self.files,
+                file: self.file,
+                lines: &self.lines,
+                nesting: self.nesting,
+                auto_import_line: self.auto_import_line,
+            },
+            &self.use_segment_positions,
+        );
+
+        let code_action = self.new_quick_fix_multiple_edits(title, text_edits);
         self.code_actions.push(code_action);
     }
 
@@ -238,4 +237,31 @@ fn main() {
 
         assert_code_action(title, src, expected).await;
     }
+
+    #[test]
+    async fn test_import_code_action_for_struct_inserts_into_existing_use() {
+        let title = "Import foo::bar::SomeTypeInBar";
+
+        let src = r#"use foo::bar::SomeOtherType;
+
+mod foo {
+    mod bar {
+        pub struct SomeTypeInBar {}
+    }
+}
+
+fn foo(x: SomeType>|<InBar) {}"#;
+
+        let expected = r#"use foo::bar::{SomeOtherType, SomeTypeInBar};
+
+mod foo {
+    mod bar {
+        pub struct SomeTypeInBar {}
+    }
+}
+
+fn foo(x: SomeTypeInBar) {}"#;
+
+        assert_code_action(title, src, expected).await;
+    }
 }

tooling/lsp/src/requests/code_action/tests.rs

@@ -1,6 +1,6 @@
 #![cfg(test)]
 
-use crate::{notifications::on_did_open_text_document, test_utils, tests::apply_text_edit};
+use crate::{notifications::on_did_open_text_document, test_utils, tests::apply_text_edits};
 
 use lsp_types::{
     CodeActionContext, CodeActionOrCommand, CodeActionParams, CodeActionResponse,
@@ -70,9 +70,8 @@ pub(crate) async fn assert_code_action(title: &str, src: &str, expected: &str) {
 
     let workspace_edit = action.edit.as_ref().unwrap();
     let text_edits = workspace_edit.changes.as_ref().unwrap().iter().next().unwrap().1;
-    assert_eq!(text_edits.len(), 1);
 
-    let result = apply_text_edit(&src.replace(">|<", ""), &text_edits[0]);
+    let result = apply_text_edits(&src.replace(">|<", ""), text_edits);
     if result != expected {
         println!("Expected:\n```\n{}\n```\n\nGot:\n```\n{}\n```", expected, result);
         assert_eq!(result, expected);

tooling/lsp/src/requests/completion.rs

@@ -38,13 +38,13 @@
 
 use crate::{
     requests::to_lsp_location, trait_impl_method_stub_generator::TraitImplMethodStubGenerator,
-    utils, visibility::is_visible, LspState,
+    use_segment_positions::UseSegmentPositions, utils, visibility::is_visible, LspState,
 };
 
 use super::process_request;
 
 mod auto_import;
 mod builtins;
 mod completion_items;
 mod kinds;
 mod sort_text;
@@ -88,42 +88,6 @@
     future::ready(result)
 }
 
-/// The position of a segment in a `use` statement.
-/// We use this to determine how an auto-import should be inserted.
-#[derive(Debug, Default, Copy, Clone)]
-enum UseSegmentPosition {
-    /// The segment either doesn't exist in the source code or there are multiple segments.
-    /// In this case auto-import will add a new use statement.
-    #[default]
-    NoneOrMultiple,
-    /// The segment is the last one in the `use` statement (or nested use statement):
-    ///
-    /// use foo::bar;
-    ///          ^^^
-    ///
-    /// Auto-import will transform it to this:
-    ///
-    /// use foo::bar::{self, baz};
-    Last { span: Span },
-    /// The segment happens before another simple (ident) segment:
-    ///
-    /// use foo::bar::qux;
-    ///          ^^^
-    ///
-    /// Auto-import will transform it to this:
-    ///
-    /// use foo::bar::{qux, baz};
-    BeforeSegment { segment_span_until_end: Span },
-    /// The segment happens before a list:
-    ///
-    /// use foo::bar::{qux, another};
-    ///
-    /// Auto-import will transform it to this:
-    ///
-    /// use foo::bar::{qux, another, baz};
-    BeforeList { first_entry_span: Span, list_is_empty: bool },
-}
-
 struct NodeFinder<'a> {
     files: &'a FileMap,
     file: FileId,
@@ -151,11 +115,7 @@
     nesting: usize,
     /// The line where an auto_import must be inserted
     auto_import_line: usize,
-    /// Remember where each segment in a `use` statement is located.
-    /// The key is the full segment, so for `use foo::bar::baz` we'll have three
-    /// segments: `foo`, `foo::bar` and `foo::bar::baz`, where the span is just
-    /// for the last identifier (`foo`, `bar` and `baz` in the previous example).
-    use_segment_positions: HashMap<String, UseSegmentPosition>,
+    use_segment_positions: UseSegmentPositions,
     self_type: Option<Type>,
     in_comptime: bool,
 }
@@ -200,7 +160,7 @@
             suggested_module_def_ids: HashSet::new(),
             nesting: 0,
             auto_import_line: 0,
-            use_segment_positions: HashMap::new(),
+            use_segment_positions: UseSegmentPositions::default(),
             self_type: None,
             in_comptime: false,
         }
@@ -280,7 +240,7 @@
         let mut idents: Vec<Ident> = Vec::new();
 
         // Find in which ident we are in, and in which part of it
         // (it could be that we are completting in the middle of an ident)
         for segment in &path.segments {
             let ident = &segment.ident;
 
@@ -1078,121 +1038,6 @@
     }
 
     /// Determine where each segment in a `use` statement is located.
-    fn gather_use_tree_segments(&mut self, use_tree: &UseTree, mut prefix: String) {
-        let kind_string = match use_tree.prefix.kind {
-            PathKind::Crate => Some("crate".to_string()),
-            PathKind::Super => Some("super".to_string()),
-            PathKind::Dep | PathKind::Plain => None,
-        };
-        if let Some(kind_string) = kind_string {
-            if let Some(segment) = use_tree.prefix.segments.first() {
-                self.insert_use_segment_position(
-                    kind_string,
-                    UseSegmentPosition::BeforeSegment {
-                        segment_span_until_end: Span::from(
-                            segment.ident.span().start()..use_tree.span.end() - 1,
-                        ),
-                    },
-                );
-            } else {
-                self.insert_use_segment_position_before_use_tree_kind(use_tree, kind_string);
-            }
-        }
-
-        let prefix_segments_len = use_tree.prefix.segments.len();
-        for (index, segment) in use_tree.prefix.segments.iter().enumerate() {
-            let ident = &segment.ident;
-            if !prefix.is_empty() {
-                prefix.push_str("::");
-            };
-            prefix.push_str(&ident.0.contents);
-
-            if index < prefix_segments_len - 1 {
-                self.insert_use_segment_position(
-                    prefix.clone(),
-                    UseSegmentPosition::BeforeSegment {
-                        segment_span_until_end: Span::from(
-                            use_tree.prefix.segments[index + 1].ident.span().start()
-                                ..use_tree.span.end() - 1,
-                        ),
-                    },
-                );
-            } else {
-                self.insert_use_segment_position_before_use_tree_kind(use_tree, prefix.clone());
-            }
-        }
-
-        match &use_tree.kind {
-            UseTreeKind::Path(ident, alias) => {
-                if !prefix.is_empty() {
-                    prefix.push_str("::");
-                }
-                prefix.push_str(&ident.0.contents);
-
-                if alias.is_none() {
-                    self.insert_use_segment_position(
-                        prefix,
-                        UseSegmentPosition::Last { span: ident.span() },
-                    );
-                } else {
-                    self.insert_use_segment_position(prefix, UseSegmentPosition::NoneOrMultiple);
-                }
-            }
-            UseTreeKind::List(use_trees) => {
-                for use_tree in use_trees {
-                    self.gather_use_tree_segments(use_tree, prefix.clone());
-                }
-            }
-        }
-    }
-
-    fn insert_use_segment_position_before_use_tree_kind(
-        &mut self,
-        use_tree: &UseTree,
-        prefix: String,
-    ) {
-        match &use_tree.kind {
-            UseTreeKind::Path(ident, _alias) => {
-                self.insert_use_segment_position(
-                    prefix,
-                    UseSegmentPosition::BeforeSegment {
-                        segment_span_until_end: Span::from(
-                            ident.span().start()..use_tree.span.end() - 1,
-                        ),
-                    },
-                );
-            }
-            UseTreeKind::List(use_trees) => {
-                if let Some(first_use_tree) = use_trees.first() {
-                    self.insert_use_segment_position(
-                        prefix,
-                        UseSegmentPosition::BeforeList {
-                            first_entry_span: first_use_tree.prefix.span(),
-                            list_is_empty: false,
-                        },
-                    );
-                } else {
-                    self.insert_use_segment_position(
-                        prefix,
-                        UseSegmentPosition::BeforeList {
-                            first_entry_span: Span::from(
-                                use_tree.span.end() - 1..use_tree.span.end() - 1,
-                            ),
-                            list_is_empty: true,
-                        },
-                    );
-                }
-            }
-        }
-    }
-
-    fn insert_use_segment_position(&mut self, segment: String, position: UseSegmentPosition) {
-        if self.use_segment_positions.get(&segment).is_none() {
-            self.use_segment_positions.insert(segment, position);
-        } else {
-            self.use_segment_positions.insert(segment, UseSegmentPosition::NoneOrMultiple);
-        }
-    }
 
     fn includes_span(&self, span: Span) -> bool {
         span.start() as usize <= self.byte_index && self.byte_index <= span.end() as usize
@@ -1205,7 +1050,7 @@
             if let Some(lsp_location) = to_lsp_location(self.files, self.file, item.span) {
                 self.auto_import_line = (lsp_location.range.end.line + 1) as usize;
             }
-            self.gather_use_tree_segments(use_tree, String::new());
+            self.use_segment_positions.add(use_tree);
         }
 
         self.includes_span(item.span)
@@ -1910,8 +1755,8 @@
 ///
 /// For example:
 ///
 /// // "merk" and "ro" match "merkle" and "root" and are in order
 /// name_matches("compute_merkle_root", "merk_ro") == true
 ///
 /// // "ro" matches "root", but "merkle" comes before it, so no match
 /// name_matches("compute_merkle_root", "ro_mer") == false