From 636d4880c4b9403aa45d837c62f4b812d0520cf0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felici=C3=A1n=20N=C3=A9meth?= <felician.nemeth@gmail.com>
Date: Sun, 27 Mar 2022 10:45:57 +0200
Subject: [PATCH 1/3] internal: Allow OnTypeFormatting to send SnippetTextEdit

But continue to send TextEdit only.
---
 crates/ide/src/typing.rs              | 22 ++++++++++++++++------
 crates/rust-analyzer/src/handlers.rs  |  6 +++---
 crates/rust-analyzer/src/lsp_ext.rs   | 17 +++++++++++++++--
 crates/rust-analyzer/src/main_loop.rs |  2 +-
 docs/dev/lsp-extensions.md            |  6 +++---
 5 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/crates/ide/src/typing.rs b/crates/ide/src/typing.rs
index a75e6be8b8a..d423d8a0546 100644
--- a/crates/ide/src/typing.rs
+++ b/crates/ide/src/typing.rs
@@ -34,6 +34,11 @@
 // Don't forget to add new trigger characters to `server_capabilities` in `caps.rs`.
 pub(crate) const TRIGGER_CHARS: &str = ".=>{";
 
+struct ExtendedTextEdit {
+    edit: TextEdit,
+    is_snippet: bool,
+}
+
 // Feature: On Typing Assists
 //
 // Some features trigger on typing certain characters:
@@ -68,22 +73,27 @@ pub(crate) fn on_char_typed(
         return None;
     }
     let edit = on_char_typed_inner(file, position.offset, char_typed)?;
-    Some(SourceChange::from_text_edit(position.file_id, edit))
+    let mut sc = SourceChange::from_text_edit(position.file_id, edit.edit);
+    sc.is_snippet = edit.is_snippet;
+    Some(sc)
 }
 
 fn on_char_typed_inner(
     file: &Parse<SourceFile>,
     offset: TextSize,
     char_typed: char,
-) -> Option<TextEdit> {
+) -> Option<ExtendedTextEdit> {
+    fn conv(text_edit: Option<TextEdit>) -> Option<ExtendedTextEdit> {
+        Some(ExtendedTextEdit { edit: text_edit?, is_snippet: false })
+    }
     if !stdx::always!(TRIGGER_CHARS.contains(char_typed)) {
         return None;
     }
     match char_typed {
-        '.' => on_dot_typed(&file.tree(), offset),
-        '=' => on_eq_typed(&file.tree(), offset),
-        '>' => on_arrow_typed(&file.tree(), offset),
-        '{' => on_opening_brace_typed(file, offset),
+        '.' => conv(on_dot_typed(&file.tree(), offset)),
+        '=' => conv(on_eq_typed(&file.tree(), offset)),
+        '>' => conv(on_arrow_typed(&file.tree(), offset)),
+        '{' => conv(on_opening_brace_typed(file, offset)),
         _ => unreachable!(),
     }
 }
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index 261c02816d5..a16f0d904ce 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -276,7 +276,7 @@ pub(crate) fn handle_on_enter(
 pub(crate) fn handle_on_type_formatting(
     snap: GlobalStateSnapshot,
     params: lsp_types::DocumentOnTypeFormattingParams,
-) -> Result<Option<Vec<lsp_types::TextEdit>>> {
+) -> Result<Option<Vec<lsp_ext::SnippetTextEdit>>> {
     let _p = profile::span("handle_on_type_formatting");
     let mut position = from_proto::file_position(&snap, params.text_document_position)?;
     let line_index = snap.file_line_index(position.file_id)?;
@@ -306,9 +306,9 @@ pub(crate) fn handle_on_type_formatting(
     };
 
     // This should be a single-file edit
-    let (_, edit) = edit.source_file_edits.into_iter().next().unwrap();
+    let (_, text_edit) = edit.source_file_edits.into_iter().next().unwrap();
 
-    let change = to_proto::text_edit_vec(&line_index, edit);
+    let change = to_proto::snippet_text_edit_vec(&line_index, edit.is_snippet, text_edit);
     Ok(Some(change))
 }
 
diff --git a/crates/rust-analyzer/src/lsp_ext.rs b/crates/rust-analyzer/src/lsp_ext.rs
index c1b230bd9df..5f0e108624b 100644
--- a/crates/rust-analyzer/src/lsp_ext.rs
+++ b/crates/rust-analyzer/src/lsp_ext.rs
@@ -4,8 +4,8 @@
 
 use lsp_types::request::Request;
 use lsp_types::{
-    notification::Notification, CodeActionKind, PartialResultParams, Position, Range,
-    TextDocumentIdentifier, WorkDoneProgressParams,
+    notification::Notification, CodeActionKind, DocumentOnTypeFormattingParams,
+    PartialResultParams, Position, Range, TextDocumentIdentifier, WorkDoneProgressParams,
 };
 use serde::{Deserialize, Serialize};
 
@@ -512,6 +512,19 @@ pub enum WorkspaceSymbolSearchKind {
     AllSymbols,
 }
 
+/// The document on type formatting request is sent from the client to
+/// the server to format parts of the document during typing.  This is
+/// almost same as lsp_types::request::OnTypeFormatting, but the
+/// result has SnippetTextEdit in it instead of TextEdit.
+#[derive(Debug)]
+pub enum OnTypeFormatting {}
+
+impl Request for OnTypeFormatting {
+    type Params = DocumentOnTypeFormattingParams;
+    type Result = Option<Vec<SnippetTextEdit>>;
+    const METHOD: &'static str = "textDocument/onTypeFormatting";
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct CompletionResolveData {
     pub position: lsp_types::TextDocumentPositionParams,
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index b5ac55e60d6..3c879687433 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -605,7 +605,7 @@ fn on_request(&mut self, request_received: Instant, req: Request) -> Result<()>
             .on::<lsp_ext::OpenCargoToml>(handlers::handle_open_cargo_toml)
             .on::<lsp_ext::MoveItem>(handlers::handle_move_item)
             .on::<lsp_ext::WorkspaceSymbol>(handlers::handle_workspace_symbol)
-            .on::<lsp_types::request::OnTypeFormatting>(handlers::handle_on_type_formatting)
+            .on::<lsp_ext::OnTypeFormatting>(handlers::handle_on_type_formatting)
             .on::<lsp_types::request::DocumentSymbolRequest>(handlers::handle_document_symbol)
             .on::<lsp_types::request::GotoDefinition>(handlers::handle_goto_definition)
             .on::<lsp_types::request::GotoDeclaration>(handlers::handle_goto_declaration)
diff --git a/docs/dev/lsp-extensions.md b/docs/dev/lsp-extensions.md
index 875561608dc..983d8e1ac09 100644
--- a/docs/dev/lsp-extensions.md
+++ b/docs/dev/lsp-extensions.md
@@ -1,5 +1,5 @@
 <!---
-lsp_ext.rs hash: 44e8238e4fbd4128
+lsp_ext.rs hash: 2a188defec26cc7c
 
 If you need to change the above hash to make the test pass, please check if you
 need to adjust this doc as well and ping this issue:
@@ -47,7 +47,7 @@ If a language client does not know about `rust-analyzer`'s configuration options
 
 **Experimental Client Capability:** `{ "snippetTextEdit": boolean }`
 
-If this capability is set, `WorkspaceEdit`s returned from `codeAction` requests might contain `SnippetTextEdit`s instead of usual `TextEdit`s:
+If this capability is set, `WorkspaceEdit`s returned from `codeAction` requests and `TextEdit`s returned from `textDocument/onTypeFormatting` requests might contain `SnippetTextEdit`s instead of usual `TextEdit`s:
 
 ```typescript
 interface SnippetTextEdit extends TextEdit {
@@ -63,7 +63,7 @@ export interface TextDocumentEdit {
 }
 ```
 
-When applying such code action, the editor should insert snippet, with tab stops and placeholder.
+When applying such code action or text edit, the editor should insert snippet, with tab stops and placeholder.
 At the moment, rust-analyzer guarantees that only a single edit will have `InsertTextFormat.Snippet`.
 
 ### Example

From 3bb02f2329623f1bb83512135746ce77ecb72b0b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felici=C3=A1n=20N=C3=A9meth?= <felician.nemeth@gmail.com>
Date: Sun, 27 Mar 2022 10:49:00 +0200
Subject: [PATCH 2/3] feat: Add on-typing handler for left angle

Only advertise this feature in the server capabilities when the client
supports SnippetTextEdit.

Close #11398.

Co-authored-by: unexge <unexge@gmail.com>
---
 crates/ide/src/typing.rs           | 212 ++++++++++++++++++++++++++++-
 crates/rust-analyzer/src/caps.rs   |  10 +-
 crates/rust-analyzer/src/config.rs |   4 +
 3 files changed, 220 insertions(+), 6 deletions(-)

diff --git a/crates/ide/src/typing.rs b/crates/ide/src/typing.rs
index d423d8a0546..be1a6da7ea8 100644
--- a/crates/ide/src/typing.rs
+++ b/crates/ide/src/typing.rs
@@ -20,9 +20,9 @@
     RootDatabase,
 };
 use syntax::{
-    algo::find_node_at_offset,
+    algo::{ancestors_at_offset, find_node_at_offset},
     ast::{self, edit::IndentLevel, AstToken},
-    AstNode, Parse, SourceFile, SyntaxKind, TextRange, TextSize,
+    AstNode, Parse, SourceFile, SyntaxKind, TextRange, TextSize, T,
 };
 
 use text_edit::{Indel, TextEdit};
@@ -32,7 +32,7 @@
 pub(crate) use on_enter::on_enter;
 
 // Don't forget to add new trigger characters to `server_capabilities` in `caps.rs`.
-pub(crate) const TRIGGER_CHARS: &str = ".=>{";
+pub(crate) const TRIGGER_CHARS: &str = ".=<>{";
 
 struct ExtendedTextEdit {
     edit: TextEdit,
@@ -92,7 +92,8 @@ fn conv(text_edit: Option<TextEdit>) -> Option<ExtendedTextEdit> {
     match char_typed {
         '.' => conv(on_dot_typed(&file.tree(), offset)),
         '=' => conv(on_eq_typed(&file.tree(), offset)),
-        '>' => conv(on_arrow_typed(&file.tree(), offset)),
+        '<' => on_left_angle_typed(&file.tree(), offset),
+        '>' => conv(on_right_angle_typed(&file.tree(), offset)),
         '{' => conv(on_opening_brace_typed(file, offset)),
         _ => unreachable!(),
     }
@@ -312,8 +313,40 @@ fn on_dot_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
     Some(TextEdit::replace(TextRange::new(offset - current_indent_len, offset), target_indent))
 }
 
+/// Add closing `>` for generic arguments/parameters.
+fn on_left_angle_typed(file: &SourceFile, offset: TextSize) -> Option<ExtendedTextEdit> {
+    let file_text = file.syntax().text();
+    if !stdx::always!(file_text.char_at(offset) == Some('<')) {
+        return None;
+    }
+    let range = TextRange::at(offset, TextSize::of('<'));
+
+    if let Some(t) = file.syntax().token_at_offset(offset).left_biased() {
+        if T![impl] == t.kind() {
+            return Some(ExtendedTextEdit {
+                edit: TextEdit::replace(range, "<$0>".to_string()),
+                is_snippet: true,
+            });
+        }
+    }
+
+    if ancestors_at_offset(file.syntax(), offset)
+        .find(|n| {
+            ast::GenericParamList::can_cast(n.kind()) || ast::GenericArgList::can_cast(n.kind())
+        })
+        .is_some()
+    {
+        return Some(ExtendedTextEdit {
+            edit: TextEdit::replace(range, "<$0>".to_string()),
+            is_snippet: true,
+        });
+    }
+
+    None
+}
+
 /// Adds a space after an arrow when `fn foo() { ... }` is turned into `fn foo() -> { ... }`
-fn on_arrow_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
+fn on_right_angle_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
     let file_text = file.syntax().text();
     if !stdx::always!(file_text.char_at(offset) == Some('>')) {
         return None;
@@ -335,6 +368,12 @@ mod tests {
 
     use super::*;
 
+    impl ExtendedTextEdit {
+        fn apply(&self, text: &mut String) {
+            self.edit.apply(text);
+        }
+    }
+
     fn do_type_char(char_typed: char, before: &str) -> Option<String> {
         let (offset, mut before) = extract_offset(before);
         let edit = TextEdit::insert(offset, char_typed.to_string());
@@ -879,6 +918,169 @@ fn adds_closing_brace_for_use_tree() {
         );
     }
 
+    #[test]
+    fn adds_closing_angle_bracket_for_generic_args() {
+        type_char(
+            '<',
+            r#"
+fn foo() {
+    bar::$0
+}
+            "#,
+            r#"
+fn foo() {
+    bar::<$0>
+}
+            "#,
+        );
+
+        type_char(
+            '<',
+            r#"
+fn foo(bar: &[u64]) {
+    bar.iter().collect::$0();
+}
+            "#,
+            r#"
+fn foo(bar: &[u64]) {
+    bar.iter().collect::<$0>();
+}
+            "#,
+        );
+    }
+
+    #[test]
+    fn adds_closing_angle_bracket_for_generic_params() {
+        type_char(
+            '<',
+            r#"
+fn foo$0() {}
+            "#,
+            r#"
+fn foo<$0>() {}
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+fn foo$0
+            "#,
+            r#"
+fn foo<$0>
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+struct Foo$0 {}
+            "#,
+            r#"
+struct Foo<$0> {}
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+struct Foo$0();
+            "#,
+            r#"
+struct Foo<$0>();
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+struct Foo$0
+            "#,
+            r#"
+struct Foo<$0>
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+enum Foo$0
+            "#,
+            r#"
+enum Foo<$0>
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+trait Foo$0
+            "#,
+            r#"
+trait Foo<$0>
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+type Foo$0 = Bar;
+            "#,
+            r#"
+type Foo<$0> = Bar;
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+impl$0 Foo {}
+            "#,
+            r#"
+impl<$0> Foo {}
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+impl<T> Foo$0 {}
+            "#,
+            r#"
+impl<T> Foo<$0> {}
+            "#,
+        );
+        type_char(
+            '<',
+            r#"
+impl Foo$0 {}
+            "#,
+            r#"
+impl Foo<$0> {}
+            "#,
+        );
+    }
+
+    #[test]
+    fn dont_add_closing_angle_bracket_for_comparison() {
+        type_char_noop(
+            '<',
+            r#"
+fn main() {
+    42$0
+}
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+fn main() {
+    42 $0
+}
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+fn main() {
+    let foo = 42;
+    foo $0
+}
+            "#,
+        );
+    }
+
     #[test]
     fn regression_629() {
         type_char_noop(
diff --git a/crates/rust-analyzer/src/caps.rs b/crates/rust-analyzer/src/caps.rs
index a653ec289b3..58b1f29df54 100644
--- a/crates/rust-analyzer/src/caps.rs
+++ b/crates/rust-analyzer/src/caps.rs
@@ -56,7 +56,7 @@ pub fn server_capabilities(config: &Config) -> ServerCapabilities {
         },
         document_on_type_formatting_provider: Some(DocumentOnTypeFormattingOptions {
             first_trigger_character: "=".to_string(),
-            more_trigger_character: Some(vec![".".to_string(), ">".to_string(), "{".to_string()]),
+            more_trigger_character: Some(more_trigger_character(&config)),
         }),
         selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
         folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
@@ -189,3 +189,11 @@ fn code_action_capabilities(client_caps: &ClientCapabilities) -> CodeActionProvi
             })
         })
 }
+
+fn more_trigger_character(config: &Config) -> Vec<String> {
+    let mut res = vec![".".to_string(), ">".to_string(), "{".to_string()];
+    if config.snippet_cap() {
+        res.push("<".to_string());
+    }
+    res
+}
diff --git a/crates/rust-analyzer/src/config.rs b/crates/rust-analyzer/src/config.rs
index d7ae4c72f5c..c53f7e8c592 100644
--- a/crates/rust-analyzer/src/config.rs
+++ b/crates/rust-analyzer/src/config.rs
@@ -1070,6 +1070,10 @@ pub fn completion(&self) -> CompletionConfig {
         }
     }
 
+    pub fn snippet_cap(&self) -> bool {
+        self.experimental("snippetTextEdit")
+    }
+
     pub fn assist(&self) -> AssistConfig {
         AssistConfig {
             snippet_cap: SnippetCap::new(self.experimental("snippetTextEdit")),

From f7c963c0f238a6ebb07bbb290d31f772b448c7aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felici=C3=A1n=20N=C3=A9meth?= <felician.nemeth@gmail.com>
Date: Tue, 5 Apr 2022 18:35:03 +0200
Subject: [PATCH 3/3] onTypeFormatting: don't insert > if another > is there

---
 crates/ide/src/typing.rs | 106 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 101 insertions(+), 5 deletions(-)

diff --git a/crates/ide/src/typing.rs b/crates/ide/src/typing.rs
index be1a6da7ea8..6af62d0ab23 100644
--- a/crates/ide/src/typing.rs
+++ b/crates/ide/src/typing.rs
@@ -83,19 +83,20 @@ fn on_char_typed_inner(
     offset: TextSize,
     char_typed: char,
 ) -> Option<ExtendedTextEdit> {
-    fn conv(text_edit: Option<TextEdit>) -> Option<ExtendedTextEdit> {
-        Some(ExtendedTextEdit { edit: text_edit?, is_snippet: false })
-    }
     if !stdx::always!(TRIGGER_CHARS.contains(char_typed)) {
         return None;
     }
-    match char_typed {
+    return match char_typed {
         '.' => conv(on_dot_typed(&file.tree(), offset)),
         '=' => conv(on_eq_typed(&file.tree(), offset)),
         '<' => on_left_angle_typed(&file.tree(), offset),
         '>' => conv(on_right_angle_typed(&file.tree(), offset)),
         '{' => conv(on_opening_brace_typed(file, offset)),
         _ => unreachable!(),
+    };
+
+    fn conv(text_edit: Option<TextEdit>) -> Option<ExtendedTextEdit> {
+        Some(ExtendedTextEdit { edit: text_edit?, is_snippet: false })
     }
 }
 
@@ -319,8 +320,17 @@ fn on_left_angle_typed(file: &SourceFile, offset: TextSize) -> Option<ExtendedTe
     if !stdx::always!(file_text.char_at(offset) == Some('<')) {
         return None;
     }
-    let range = TextRange::at(offset, TextSize::of('<'));
 
+    // Find the next non-whitespace char in the line.
+    let mut next_offset = offset + TextSize::of('<');
+    while file_text.char_at(next_offset) == Some(' ') {
+        next_offset += TextSize::of(' ')
+    }
+    if file_text.char_at(next_offset) == Some('>') {
+        return None;
+    }
+
+    let range = TextRange::at(offset, TextSize::of('<'));
     if let Some(t) = file.syntax().token_at_offset(offset).left_biased() {
         if T![impl] == t.kind() {
             return Some(ExtendedTextEdit {
@@ -1081,6 +1091,92 @@ fn main() {
         );
     }
 
+    #[test]
+    fn dont_add_closing_angle_bracket_if_it_is_already_there() {
+        type_char_noop(
+            '<',
+            r#"
+fn foo() {
+    bar::$0>
+}
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+fn foo(bar: &[u64]) {
+    bar.iter().collect::$0   >();
+}
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+fn foo$0>() {}
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+fn foo$0>
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+struct Foo$0> {}
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+struct Foo$0>();
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+struct Foo$0>
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+enum Foo$0>
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+trait Foo$0>
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+type Foo$0> = Bar;
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+impl$0> Foo {}
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+impl<T> Foo$0> {}
+            "#,
+        );
+        type_char_noop(
+            '<',
+            r#"
+impl Foo$0> {}
+            "#,
+        );
+    }
+
     #[test]
     fn regression_629() {
         type_char_noop(