Add completion labels

src/wit.rs

@@ -1,4 +1,7 @@
-use zed_extension_api::{register_extension, Extension, LanguageServerId, Worktree};
+use zed_extension_api::{
+    lsp::{Completion, CompletionKind},
+    register_extension, CodeLabel, CodeLabelSpan, Extension, LanguageServerId, Worktree,
+};
 
 struct WitExtension;
 
@@ -25,6 +28,27 @@ impl Extension for WitExtension {
             env: Default::default(),
         })
     }
+
+    fn label_for_completion(
+        &self,
+        _language_server_id: &LanguageServerId,
+        completion: Completion,
+    ) -> Option<CodeLabel> {
+        // https://github.com/Michael-F-Bryan/wit-lsp/blob/main/crates/wit-language-server/src/ops/completion.rs#L110
+        let highlight_name = match completion.kind? {
+            CompletionKind::Struct | CompletionKind::Interface | CompletionKind::Module => {
+                Some("type".to_string())
+            }
+            CompletionKind::Keyword => Some("keyword".to_string()),
+            _ => None,
+        };
+
+        Some(CodeLabel {
+            spans: vec![CodeLabelSpan::literal(&completion.label, highlight_name)],
+            filter_range: (0..completion.label.len()).into(),
+            code: completion.label,
+        })
+    }
 }
 
 register_extension!(WitExtension);