Apply additional edits for completion when the buffer is remote

Cargo.lock

@@ -2635,6 +2635,7 @@ dependencies = [
  "rand 0.8.3",
  "rpc",
  "serde",
+ "serde_json",
  "similar",
  "smallvec",
  "smol",

crates/editor/src/editor.rs

@@ -1683,9 +1683,9 @@ impl Editor {
             });
         }
 
-        self.buffer.update(cx, |buffer, cx| {
+        Some(self.buffer.update(cx, |buffer, cx| {
             buffer.apply_additional_edits_for_completion(completion.clone(), cx)
-        })
+        }))
     }
 
     pub fn has_completions(&self) -> bool {

crates/editor/src/multi_buffer.rs

@@ -313,9 +313,9 @@ impl MultiBuffer {
                 .map(|range| range.start.to_offset(&snapshot)..range.end.to_offset(&snapshot));
             return buffer.update(cx, |buffer, cx| {
                 if autoindent {
-                    buffer.edit_with_autoindent(ranges, new_text, cx)
+                    buffer.edit_with_autoindent(ranges, new_text, cx);
                 } else {
-                    buffer.edit(ranges, new_text, cx)
+                    buffer.edit(ranges, new_text, cx);
                 }
             });
         }
@@ -922,14 +922,18 @@ impl MultiBuffer {
         &self,
         completion: Completion<Anchor>,
         cx: &mut ModelContext<Self>,
-    ) -> Option<Task<Result<()>>> {
-        let buffer = self
+    ) -> Task<Result<()>> {
+        let buffer = if let Some(buffer_state) = self
             .buffers
             .borrow()
-            .get(&completion.old_range.start.buffer_id)?
-            .buffer
-            .clone();
-        buffer.update(cx, |buffer, cx| {
+            .get(&completion.old_range.start.buffer_id)
+        {
+            buffer_state.buffer.clone()
+        } else {
+            return Task::ready(Ok(()));
+        };
+
+        let apply_edits = buffer.update(cx, |buffer, cx| {
             buffer.apply_additional_edits_for_completion(
                 Completion {
                     old_range: completion.old_range.start.text_anchor
@@ -937,8 +941,13 @@ impl MultiBuffer {
                     new_text: completion.new_text,
                     lsp_completion: completion.lsp_completion,
                 },
+                true,
                 cx,
             )
+        });
+        cx.foreground().spawn(async move {
+            apply_edits.await?;
+            Ok(())
         })
     }
 

crates/language/Cargo.toml

@@ -36,6 +36,7 @@ parking_lot = "0.11.1"
 postage = { version = "0.4.1", features = ["futures-traits"] }
 rand = { version = "0.8.3", optional = true }
 serde = { version = "1", features = ["derive"] }
+serde_json = { version = "1", features = ["preserve_order"] }
 similar = "1.3"
 smallvec = { version = "1.6", features = ["union"] }
 smol = "1.2"

crates/language/src/buffer.rs

@@ -206,6 +206,13 @@ pub trait File {
         cx: &mut MutableAppContext,
     ) -> Task<Result<Vec<Completion<Anchor>>>>;
 
+    fn apply_additional_edits_for_completion(
+        &self,
+        buffer_id: u64,
+        completion: Completion<Anchor>,
+        cx: &mut MutableAppContext,
+    ) -> Task<Result<Vec<clock::Local>>>;
+
     fn buffer_updated(&self, buffer_id: u64, operation: Operation, cx: &mut MutableAppContext);
 
     fn buffer_removed(&self, buffer_id: u64, cx: &mut MutableAppContext);
@@ -284,6 +291,15 @@ impl File for FakeFile {
         Task::ready(Ok(Default::default()))
     }
 
+    fn apply_additional_edits_for_completion(
+        &self,
+        _: u64,
+        _: Completion<Anchor>,
+        _: &mut MutableAppContext,
+    ) -> Task<Result<Vec<clock::Local>>> {
+        Task::ready(Ok(Default::default()))
+    }
+
     fn buffer_updated(&self, _: u64, _: Operation, _: &mut MutableAppContext) {}
 
     fn buffer_removed(&self, _: u64, _: &mut MutableAppContext) {}
@@ -595,7 +611,8 @@ impl Buffer {
                 if let Some(edits) = edits {
                     this.update(&mut cx, |this, cx| {
                         if this.version == version {
-                            this.apply_lsp_edits(edits, cx)
+                            this.apply_lsp_edits(edits, cx)?;
+                            Ok(())
                         } else {
                             Err(anyhow!("buffer edited since starting to format"))
                         }
@@ -1295,7 +1312,9 @@ impl Buffer {
                 let range = offset..(offset + len);
                 match tag {
                     ChangeTag::Equal => offset += len,
-                    ChangeTag::Delete => self.edit(Some(range), "", cx),
+                    ChangeTag::Delete => {
+                        self.edit(Some(range), "", cx);
+                    }
                     ChangeTag::Insert => {
                         self.edit(Some(offset..offset), &diff.new_text[range], cx);
                         offset += len;
@@ -1409,7 +1428,12 @@ impl Buffer {
             .blocking_send(Some(snapshot));
     }
 
-    pub fn edit<I, S, T>(&mut self, ranges_iter: I, new_text: T, cx: &mut ModelContext<Self>)
+    pub fn edit<I, S, T>(
+        &mut self,
+        ranges_iter: I,
+        new_text: T,
+        cx: &mut ModelContext<Self>,
+    ) -> Option<clock::Local>
     where
         I: IntoIterator<Item = Range<S>>,
         S: ToOffset,
@@ -1423,7 +1447,8 @@ impl Buffer {
         ranges_iter: I,
         new_text: T,
         cx: &mut ModelContext<Self>,
-    ) where
+    ) -> Option<clock::Local>
+    where
         I: IntoIterator<Item = Range<S>>,
         S: ToOffset,
         T: Into<String>,
@@ -1437,7 +1462,8 @@ impl Buffer {
         new_text: T,
         autoindent: bool,
         cx: &mut ModelContext<Self>,
-    ) where
+    ) -> Option<clock::Local>
+    where
         I: IntoIterator<Item = Range<S>>,
         S: ToOffset,
         T: Into<String>,
@@ -1461,7 +1487,7 @@ impl Buffer {
             }
         }
         if ranges.is_empty() {
-            return;
+            return None;
         }
 
         self.start_transaction();
@@ -1488,6 +1514,7 @@ impl Buffer {
         let new_text_len = new_text.len();
 
         let edit = self.text.edit(ranges.iter().cloned(), new_text);
+        let edit_id = edit.timestamp.local();
 
         if let Some((before_edit, edited)) = autoindent_request {
             let mut inserted = None;
@@ -1517,13 +1544,14 @@ impl Buffer {
 
         self.end_transaction(cx);
         self.send_operation(Operation::Buffer(text::Operation::Edit(edit)), cx);
+        Some(edit_id)
     }
 
     fn apply_lsp_edits(
         &mut self,
         edits: Vec<lsp::TextEdit>,
         cx: &mut ModelContext<Self>,
-    ) -> Result<()> {
+    ) -> Result<Vec<clock::Local>> {
         for edit in &edits {
             let range = range_from_lsp(edit.range);
             if self.clip_point_utf16(range.start, Bias::Left) != range.start
@@ -1535,11 +1563,14 @@ impl Buffer {
             }
         }
 
-        for edit in edits.into_iter().rev() {
-            self.edit([range_from_lsp(edit.range)], edit.new_text, cx);
-        }
-
-        Ok(())
+        self.start_transaction();
+        let edit_ids = edits
+            .into_iter()
+            .rev()
+            .filter_map(|edit| self.edit([range_from_lsp(edit.range)], edit.new_text, cx))
+            .collect();
+        self.end_transaction(cx);
+        Ok(edit_ids)
     }
 
     fn did_edit(
@@ -1835,21 +1866,59 @@ impl Buffer {
     pub fn apply_additional_edits_for_completion(
         &mut self,
         completion: Completion<Anchor>,
+        push_to_history: bool,
         cx: &mut ModelContext<Self>,
-    ) -> Option<Task<Result<()>>> {
-        self.file.as_ref()?.as_local()?;
-        let server = self.language_server.as_ref()?.server.clone();
-        Some(cx.spawn(|this, mut cx| async move {
-            let resolved_completion = server
-                .request::<lsp::request::ResolveCompletionItem>(completion.lsp_completion)
-                .await?;
-            if let Some(additional_edits) = resolved_completion.additional_text_edits {
-                this.update(&mut cx, |this, cx| {
-                    this.apply_lsp_edits(additional_edits, cx)
-                })?;
-            }
-            Ok::<_, anyhow::Error>(())
-        }))
+    ) -> Task<Result<Vec<clock::Local>>> {
+        let file = if let Some(file) = self.file.as_ref() {
+            file
+        } else {
+            return Task::ready(Ok(Default::default()));
+        };
+
+        if file.is_local() {
+            let server = if let Some(lang) = self.language_server.as_ref() {
+                lang.server.clone()
+            } else {
+                return Task::ready(Ok(Default::default()));
+            };
+
+            cx.spawn(|this, mut cx| async move {
+                let resolved_completion = server
+                    .request::<lsp::request::ResolveCompletionItem>(completion.lsp_completion)
+                    .await?;
+                if let Some(additional_edits) = resolved_completion.additional_text_edits {
+                    this.update(&mut cx, |this, cx| {
+                        this.avoid_grouping_next_transaction();
+                        this.start_transaction();
+                        let edit_ids = this.apply_lsp_edits(additional_edits, cx);
+                        if let Some(transaction_id) = this.end_transaction(cx) {
+                            if !push_to_history {
+                                this.text.forget_transaction(transaction_id);
+                            }
+                        }
+                        edit_ids
+                    })
+                } else {
+                    Ok(Default::default())
+                }
+            })
+        } else {
+            let apply_edits = file.apply_additional_edits_for_completion(
+                self.remote_id(),
+                completion,
+                cx.as_mut(),
+            );
+            cx.spawn(|this, mut cx| async move {
+                let edit_ids = apply_edits.await?;
+                if push_to_history {
+                    this.update(&mut cx, |this, _| {
+                        this.text
+                            .push_transaction(edit_ids.iter().copied(), Instant::now());
+                    });
+                }
+                Ok(edit_ids)
+            })
+        }
     }
 
     pub fn completion_triggers(&self) -> &[String] {

crates/language/src/proto.rs

@@ -1,4 +1,4 @@
-use crate::{diagnostic_set::DiagnosticEntry, Diagnostic, Operation};
+use crate::{diagnostic_set::DiagnosticEntry, Completion, Diagnostic, Operation};
 use anyhow::{anyhow, Result};
 use clock::ReplicaId;
 use collections::HashSet;
@@ -377,3 +377,28 @@ pub fn deserialize_anchor(anchor: proto::Anchor) -> Option<Anchor> {
         },
     })
 }
+
+pub fn serialize_completion(completion: &Completion<Anchor>) -> proto::Completion {
+    proto::Completion {
+        old_start: Some(serialize_anchor(&completion.old_range.start)),
+        old_end: Some(serialize_anchor(&completion.old_range.end)),
+        new_text: completion.new_text.clone(),
+        lsp_completion: serde_json::to_vec(&completion.lsp_completion).unwrap(),
+    }
+}
+
+pub fn deserialize_completion(completion: proto::Completion) -> Result<Completion<Anchor>> {
+    let old_start = completion
+        .old_start
+        .and_then(deserialize_anchor)
+        .ok_or_else(|| anyhow!("invalid old start"))?;
+    let old_end = completion
+        .old_end
+        .and_then(deserialize_anchor)
+        .ok_or_else(|| anyhow!("invalid old end"))?;
+    Ok(Completion {
+        old_range: old_start..old_end,
+        new_text: completion.new_text,
+        lsp_completion: serde_json::from_slice(&completion.lsp_completion)?,
+    })
+}

crates/project/src/project.rs

@@ -335,6 +335,11 @@ impl Project {
                 client.subscribe_to_entity(remote_id, cx, Self::handle_buffer_saved),
                 client.subscribe_to_entity(remote_id, cx, Self::handle_format_buffer),
                 client.subscribe_to_entity(remote_id, cx, Self::handle_get_completions),
+                client.subscribe_to_entity(
+                    remote_id,
+                    cx,
+                    Self::handle_apply_additional_edits_for_completion,
+                ),
                 client.subscribe_to_entity(remote_id, cx, Self::handle_get_definition),
             ]);
         }
@@ -1712,17 +1717,63 @@ impl Project {
                         receipt,
                         proto::GetCompletionsResponse {
                             completions: completions
+                                .iter()
+                                .map(language::proto::serialize_completion)
+                                .collect(),
+                        },
+                    )
+                    .await
+                }
+                Err(error) => {
+                    rpc.respond_with_error(
+                        receipt,
+                        proto::Error {
+                            message: error.to_string(),
+                        },
+                    )
+                    .await
+                }
+            }
+        })
+        .detach_and_log_err(cx);
+        Ok(())
+    }
+
+    fn handle_apply_additional_edits_for_completion(
+        &mut self,
+        envelope: TypedEnvelope<proto::ApplyCompletionAdditionalEdits>,
+        rpc: Arc<Client>,
+        cx: &mut ModelContext<Self>,
+    ) -> Result<()> {
+        let receipt = envelope.receipt();
+        let sender_id = envelope.original_sender_id()?;
+        let buffer = self
+            .shared_buffers
+            .get(&sender_id)
+            .and_then(|shared_buffers| shared_buffers.get(&envelope.payload.buffer_id).cloned())
+            .ok_or_else(|| anyhow!("unknown buffer id {}", envelope.payload.buffer_id))?;
+        let completion = language::proto::deserialize_completion(
+            envelope
+                .payload
+                .completion
+                .ok_or_else(|| anyhow!("invalid position"))?,
+        )?;
+        cx.spawn(|_, mut cx| async move {
+            match buffer
+                .update(&mut cx, |buffer, cx| {
+                    buffer.apply_additional_edits_for_completion(completion, false, cx)
+                })
+                .await
+            {
+                Ok(edit_ids) => {
+                    rpc.respond(
+                        receipt,
+                        proto::ApplyCompletionAdditionalEditsResponse {
+                            additional_edits: edit_ids
                                 .into_iter()
-                                .map(|completion| proto::Completion {
-                                    old_start: Some(language::proto::serialize_anchor(
-                                        &completion.old_range.start,
-                                    )),
-                                    old_end: Some(language::proto::serialize_anchor(
-                                        &completion.old_range.end,
-                                    )),
-                                    new_text: completion.new_text,
-                                    lsp_completion: serde_json::to_vec(&completion.lsp_completion)
-                                        .unwrap(),
+                                .map(|edit_id| proto::AdditionalEdit {
+                                    replica_id: edit_id.replica_id as u32,
+                                    local_timestamp: edit_id.value,
                                 })
                                 .collect(),
                         },

crates/project/src/worktree.rs

@@ -1448,25 +1448,47 @@ impl language::File for File {
             response
                 .completions
                 .into_iter()
-                .map(|completion| {
-                    let old_start = completion
-                        .old_start
-                        .and_then(language::proto::deserialize_anchor)
-                        .ok_or_else(|| anyhow!("invalid old start"))?;
-                    let old_end = completion
-                        .old_end
-                        .and_then(language::proto::deserialize_anchor)
-                        .ok_or_else(|| anyhow!("invalid old end"))?;
-                    Ok(Completion {
-                        old_range: old_start..old_end,
-                        new_text: completion.new_text,
-                        lsp_completion: serde_json::from_slice(&completion.lsp_completion)?,
-                    })
-                })
+                .map(language::proto::deserialize_completion)
                 .collect()
         })
     }
 
+    fn apply_additional_edits_for_completion(
+        &self,
+        buffer_id: u64,
+        completion: Completion<Anchor>,
+        cx: &mut MutableAppContext,
+    ) -> Task<Result<Vec<clock::Local>>> {
+        let worktree = self.worktree.read(cx);
+        let worktree = if let Some(worktree) = worktree.as_remote() {
+            worktree
+        } else {
+            return Task::ready(Err(anyhow!(
+                "remote additional edits application requested on a local worktree"
+            )));
+        };
+        let rpc = worktree.client.clone();
+        let project_id = worktree.project_id;
+        cx.foreground().spawn(async move {
+            let response = rpc
+                .request(proto::ApplyCompletionAdditionalEdits {
+                    project_id,
+                    buffer_id,
+                    completion: Some(language::proto::serialize_completion(&completion)),
+                })
+                .await?;
+
+            Ok(response
+                .additional_edits
+                .into_iter()
+                .map(|edit| clock::Local {
+                    replica_id: edit.replica_id as ReplicaId,
+                    value: edit.local_timestamp,
+                })
+                .collect())
+        })
+    }
+
     fn buffer_updated(&self, buffer_id: u64, operation: Operation, cx: &mut MutableAppContext) {
         self.worktree.update(cx, |worktree, cx| {
             worktree.send_buffer_update(buffer_id, operation, cx);

crates/rpc/proto/zed.proto

@@ -42,22 +42,24 @@ message Envelope {
         FormatBuffer format_buffer = 34;
         GetCompletions get_completions = 35;
         GetCompletionsResponse get_completions_response = 36;
+        ApplyCompletionAdditionalEdits apply_completion_additional_edits = 37;
+        ApplyCompletionAdditionalEditsResponse apply_completion_additional_edits_response = 38;
 
-        GetChannels get_channels = 37;
-        GetChannelsResponse get_channels_response = 38;
-        JoinChannel join_channel = 39;
-        JoinChannelResponse join_channel_response = 40;
-        LeaveChannel leave_channel = 41;
-        SendChannelMessage send_channel_message = 42;
-        SendChannelMessageResponse send_channel_message_response = 43;
-        ChannelMessageSent channel_message_sent = 44;
-        GetChannelMessages get_channel_messages = 45;
-        GetChannelMessagesResponse get_channel_messages_response = 46;
+        GetChannels get_channels = 39;
+        GetChannelsResponse get_channels_response = 40;
+        JoinChannel join_channel = 41;
+        JoinChannelResponse join_channel_response = 42;
+        LeaveChannel leave_channel = 43;
+        SendChannelMessage send_channel_message = 44;
+        SendChannelMessageResponse send_channel_message_response = 45;
+        ChannelMessageSent channel_message_sent = 46;
+        GetChannelMessages get_channel_messages = 47;
+        GetChannelMessagesResponse get_channel_messages_response = 48;
 
-        UpdateContacts update_contacts = 47;
+        UpdateContacts update_contacts = 49;
 
-        GetUsers get_users = 48;
-        GetUsersResponse get_users_response = 49;
+        GetUsers get_users = 50;
+        GetUsersResponse get_users_response = 51;
     }
 }
 
@@ -215,6 +217,21 @@ message GetCompletionsResponse {
     repeated Completion completions = 1;
 }
 
+message ApplyCompletionAdditionalEdits {
+    uint64 project_id = 1;
+    uint64 buffer_id = 2;
+    Completion completion = 3;
+}
+
+message ApplyCompletionAdditionalEditsResponse {
+    repeated AdditionalEdit additional_edits = 1;
+}
+
+message AdditionalEdit {
+    uint32 replica_id = 1;
+    uint32 local_timestamp = 2;
+}
+
 message Completion {
     Anchor old_start = 1;
     Anchor old_end = 2;

crates/rpc/src/proto.rs

@@ -122,6 +122,8 @@ macro_rules! entity_messages {
 messages!(
     Ack,
     AddProjectCollaborator,
+    ApplyCompletionAdditionalEdits,
+    ApplyCompletionAdditionalEditsResponse,
     BufferReloaded,
     BufferSaved,
     ChannelMessageSent,
@@ -169,6 +171,10 @@ messages!(
 );
 
 request_messages!(
+    (
+        ApplyCompletionAdditionalEdits,
+        ApplyCompletionAdditionalEditsResponse
+    ),
     (FormatBuffer, Ack),
     (GetChannelMessages, GetChannelMessagesResponse),
     (GetChannels, GetChannelsResponse),
@@ -191,6 +197,7 @@ request_messages!(
 entity_messages!(
     project_id,
     AddProjectCollaborator,
+    ApplyCompletionAdditionalEdits,
     BufferReloaded,
     BufferSaved,
     CloseBuffer,

crates/server/src/rpc.rs

@@ -84,6 +84,7 @@ impl Server {
             .add_handler(Server::save_buffer)
             .add_handler(Server::format_buffer)
             .add_handler(Server::get_completions)
+            .add_handler(Server::apply_additional_edits_for_completion)
             .add_handler(Server::get_channels)
             .add_handler(Server::get_users)
             .add_handler(Server::join_channel)
@@ -747,6 +748,30 @@ impl Server {
         Ok(())
     }
 
+    async fn apply_additional_edits_for_completion(
+        self: Arc<Server>,
+        request: TypedEnvelope<proto::ApplyCompletionAdditionalEdits>,
+    ) -> tide::Result<()> {
+        let host;
+        {
+            let state = self.state();
+            let project = state
+                .read_project(request.payload.project_id, request.sender_id)
+                .ok_or_else(|| anyhow!(NO_SUCH_PROJECT))?;
+            host = project.host_connection_id;
+        }
+
+        let sender = request.sender_id;
+        let receipt = request.receipt();
+        let response = self
+            .peer
+            .forward_request(sender, host, request.payload.clone())
+            .await?;
+        self.peer.respond(receipt, response).await?;
+
+        Ok(())
+    }
+
     async fn update_buffer(
         self: Arc<Server>,
         request: TypedEnvelope<proto::UpdateBuffer>,

crates/text/src/text.rs

@@ -233,6 +233,20 @@ impl History {
         }
     }
 
+    fn push_transaction(&mut self, edit_ids: impl IntoIterator<Item = clock::Local>, now: Instant) {
+        assert_eq!(self.transaction_depth, 0);
+        let mut edit_ids = edit_ids.into_iter().peekable();
+
+        if let Some(first_edit_id) = edit_ids.peek() {
+            let version = self.ops[first_edit_id].version.clone();
+            self.start_transaction(version, now);
+            for edit_id in edit_ids {
+                self.push_undo(edit_id);
+            }
+            self.end_transaction(now);
+        }
+    }
+
     fn push_undo(&mut self, edit_id: clock::Local) {
         assert_ne!(self.transaction_depth, 0);
         let last_transaction = self.undo_stack.last_mut().unwrap();
@@ -260,6 +274,17 @@ impl History {
         }
     }
 
+    fn forget(&mut self, transaction_id: TransactionId) {
+        assert_eq!(self.transaction_depth, 0);
+        if let Some(transaction_ix) = self.undo_stack.iter().rposition(|t| t.id == transaction_id) {
+            self.undo_stack.remove(transaction_ix);
+        } else if let Some(transaction_ix) =
+            self.redo_stack.iter().rposition(|t| t.id == transaction_id)
+        {
+            self.undo_stack.remove(transaction_ix);
+        }
+    }
+
     fn pop_redo(&mut self) -> Option<&Transaction> {
         assert_eq!(self.transaction_depth, 0);
         if let Some(transaction) = self.redo_stack.pop() {
@@ -377,14 +402,14 @@ pub struct InsertionTimestamp {
 }
 
 impl InsertionTimestamp {
-    fn local(&self) -> clock::Local {
+    pub fn local(&self) -> clock::Local {
         clock::Local {
             replica_id: self.replica_id,
             value: self.local,
         }
     }
 
-    fn lamport(&self) -> clock::Lamport {
+    pub fn lamport(&self) -> clock::Lamport {
         clock::Lamport {
             replica_id: self.replica_id,
             value: self.lamport,
@@ -1188,6 +1213,7 @@ impl Buffer {
 
     pub fn undo(&mut self) -> Option<(TransactionId, Operation)> {
         if let Some(transaction) = self.history.pop_undo().cloned() {
+            dbg!(&transaction);
             let transaction_id = transaction.id;
             let op = self.undo_or_redo(transaction).unwrap();
             Some((transaction_id, op))
@@ -1205,6 +1231,10 @@ impl Buffer {
         }
     }
 
+    pub fn forget_transaction(&mut self, transaction_id: TransactionId) {
+        self.history.forget(transaction_id);
+    }
+
     pub fn redo(&mut self) -> Option<(TransactionId, Operation)> {
         if let Some(transaction) = self.history.pop_redo().cloned() {
             let transaction_id = transaction.id;
@@ -1245,6 +1275,14 @@ impl Buffer {
         })
     }
 
+    pub fn push_transaction(
+        &mut self,
+        edit_ids: impl IntoIterator<Item = clock::Local>,
+        now: Instant,
+    ) {
+        self.history.push_transaction(edit_ids, now);
+    }
+
     pub fn subscribe(&mut self) -> Subscription {
         self.subscriptions.subscribe()
     }