fix: preserve code lenses on resolve

Author: matoous

helix-term/src/commands/lsp.rs

@@ -1512,7 +1512,7 @@ pub fn show_code_lenses_under_cursor(cx: &mut Context) {
             }
 
             if let Some(doc) = editor.documents.get_mut(&doc_id) {
-                doc.set_code_lenses(lenses.clone());
+                doc.update_code_lenses(lenses.clone());
             }
 
             let mut picker = ui::Menu::new(lenses, (), move |editor, lens, event| {

helix-view/src/document.rs

@@ -2240,16 +2240,16 @@ impl Document {
         &self.code_lenses
     }
 
-    pub fn set_code_lenses(&mut self, lenses: Vec<lsp::CodeLens>) {
+    fn map_code_lenses(&self, lenses: Vec<lsp::CodeLens>) -> Vec<CodeLens> {
         let text = self.text().clone();
         let Some(language_server) = self
             .language_servers_with_feature(LanguageServerFeature::CodeLens)
             .next()
         else {
-            return;
+            return Vec::new();
         };
         let offset_encoding = language_server.offset_encoding();
-        self.code_lenses = lenses
+        lenses
             .into_iter()
             .filter_map(|l| {
                 if let Some(char_idx) = lsp_pos_to_pos(&text, l.range.start, offset_encoding) {
@@ -2262,7 +2262,22 @@ impl Document {
                     None
                 }
             })
-            .collect();
+            .collect()
+    }
+
+    pub fn set_code_lenses(&mut self, lenses: Vec<lsp::CodeLens>) {
+        self.code_lenses = self.map_code_lenses(lenses);
+    }
+
+    pub fn update_code_lenses(&mut self, lenses: Vec<lsp::CodeLens>) {
+        for updated_lens in lenses {
+            if let Some(existing_lens) = self.code_lenses.iter_mut().find(|existing_lens| {
+                existing_lens.lens.range == updated_lens.range
+                    && existing_lens.lens.data == updated_lens.data
+            }) {
+                existing_lens.lens = updated_lens;
+            }
+        }
     }
 
     /// Get the document's auto pairs. If the document has a recognized