feat: LSP auto-import will try to add to existing use statements (#6354)

compiler/noirc_frontend/src/ast/statement.rs

@@ -327,6 +327,7 @@ pub enum PathKind {
 pub struct UseTree {
     pub prefix: Path,
     pub kind: UseTreeKind,
+    pub span: Span,
 }
 
 impl Display for UseTree {

compiler/noirc_frontend/src/parser/parser/use_tree.rs

@@ -53,13 +53,17 @@ impl<'a> Parser<'a> {
                     Self::parse_use_tree_in_list,
                 );
 
-                UseTree { prefix, kind: UseTreeKind::List(use_trees) }
+                UseTree {
+                    prefix,
+                    kind: UseTreeKind::List(use_trees),
+                    span: self.span_since(start_span),
+                }
             } else {
                 self.expected_token(Token::LeftBrace);
-                self.parse_path_use_tree_end(prefix, nested)
+                self.parse_path_use_tree_end(prefix, nested, start_span)
             }
         } else {
-            self.parse_path_use_tree_end(prefix, nested)
+            self.parse_path_use_tree_end(prefix, nested, start_span)
         }
     }
 
@@ -71,6 +75,7 @@ impl<'a> Parser<'a> {
             return Some(UseTree {
                 prefix: Path { segments: Vec::new(), kind: PathKind::Plain, span: start_span },
                 kind: UseTreeKind::Path(Ident::new("self".to_string(), start_span), None),
+                span: start_span,
             });
         }
 
@@ -89,25 +94,46 @@ impl<'a> Parser<'a> {
         }
     }
 
-    pub(super) fn parse_path_use_tree_end(&mut self, mut prefix: Path, nested: bool) -> UseTree {
+    pub(super) fn parse_path_use_tree_end(
+        &mut self,
+        mut prefix: Path,
+        nested: bool,
+        start_span: Span,
+    ) -> UseTree {
         if prefix.segments.is_empty() {
             if nested {
                 self.expected_identifier();
             } else {
                 self.expected_label(ParsingRuleLabel::UseSegment);
             }
-            UseTree { prefix, kind: UseTreeKind::Path(Ident::default(), None) }
+            UseTree {
+                prefix,
+                kind: UseTreeKind::Path(Ident::default(), None),
+                span: self.span_since(start_span),
+            }
         } else {
             let ident = prefix.segments.pop().unwrap().ident;
             if self.eat_keyword(Keyword::As) {
                 if let Some(alias) = self.eat_ident() {
-                    UseTree { prefix, kind: UseTreeKind::Path(ident, Some(alias)) }
+                    UseTree {
+                        prefix,
+                        kind: UseTreeKind::Path(ident, Some(alias)),
+                        span: self.span_since(start_span),
+                    }
                 } else {
                     self.expected_identifier();
-                    UseTree { prefix, kind: UseTreeKind::Path(ident, None) }
+                    UseTree {
+                        prefix,
+                        kind: UseTreeKind::Path(ident, None),
+                        span: self.span_since(start_span),
+                    }
                 }
             } else {
-                UseTree { prefix, kind: UseTreeKind::Path(ident, None) }
+                UseTree {
+                    prefix,
+                    kind: UseTreeKind::Path(ident, None),
+                    span: self.span_since(start_span),
+                }
             }
         }
     }

tooling/lsp/src/lib.rs

@@ -67,6 +67,7 @@ mod modules;
 mod notifications;
 mod requests;
 mod solver;
+mod tests;
 mod trait_impl_method_stub_generator;
 mod types;
 mod utils;

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

@@ -106,11 +106,12 @@ fn use_tree_without_unused_import(
                 let mut prefix = use_tree.prefix.clone();
                 prefix.segments.extend(new_use_tree.prefix.segments);
 
-                Some(UseTree { prefix, kind: new_use_tree.kind })
+                Some(UseTree { prefix, kind: new_use_tree.kind, span: use_tree.span })
             } else {
                 Some(UseTree {
                     prefix: use_tree.prefix.clone(),
                     kind: UseTreeKind::List(new_use_trees),
+                    span: use_tree.span,
                 })
             };
 

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

@@ -1,11 +1,11 @@
 #![cfg(test)]
 
-use crate::{notifications::on_did_open_text_document, test_utils};
+use crate::{notifications::on_did_open_text_document, test_utils, tests::apply_text_edit};
 
 use lsp_types::{
     CodeActionContext, CodeActionOrCommand, CodeActionParams, CodeActionResponse,
     DidOpenTextDocumentParams, PartialResultParams, Position, Range, TextDocumentIdentifier,
-    TextDocumentItem, TextEdit, WorkDoneProgressParams,
+    TextDocumentItem, WorkDoneProgressParams,
 };
 
 use super::on_code_action_request;
@@ -78,16 +78,3 @@ pub(crate) async fn assert_code_action(title: &str, src: &str, expected: &str) {
         assert_eq!(result, expected);
     }
 }
-
-fn apply_text_edit(src: &str, text_edit: &TextEdit) -> String {
-    let mut lines: Vec<_> = src.lines().collect();
-    assert_eq!(text_edit.range.start.line, text_edit.range.end.line);
-
-    let mut line = lines[text_edit.range.start.line as usize].to_string();
-    line.replace_range(
-        text_edit.range.start.character as usize..text_edit.range.end.character as usize,
-        &text_edit.new_text,
-    );
-    lines[text_edit.range.start.line as usize] = &line;
-    lines.join("\n")
-}

tooling/lsp/src/requests/completion.rs

@@ -88,6 +88,42 @@ pub(crate) fn on_completion_request(
     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,
@@ -115,6 +151,11 @@ struct NodeFinder<'a> {
     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>,
     self_type: Option<Type>,
     in_comptime: bool,
 }
@@ -159,6 +200,7 @@ impl<'a> NodeFinder<'a> {
             suggested_module_def_ids: HashSet::new(),
             nesting: 0,
             auto_import_line: 0,
+            use_segment_positions: HashMap::new(),
             self_type: None,
             in_comptime: false,
         }
@@ -1035,17 +1077,135 @@ impl<'a> NodeFinder<'a> {
         }
     }
 
+    /// 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
     }
 }
 
 impl<'a> Visitor for NodeFinder<'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.gather_use_tree_segments(use_tree, String::new());
         }
 
         self.includes_span(item.span)