@@ -1,5 +1,6 @@
11use std::collections::HashMap;
22use std::sync::Arc;
3+use std::sync::atomic::{AtomicBool, Ordering};
34use std::time::Duration;
45
56use serde_json::json;
@@ -36,6 +37,9 @@ struct Backend {
3637 /// Per-URI debounce handles. The value is a join handle for an in-flight
3738 /// diagnostic publish; replacing the entry cancels the prior handle.
3839 pending_diags: Arc<Mutex<HashMap<Url, JoinHandle<()>>>>,
40+ /// Whether the client accepted the `utf-8` position encoding. When false,
41+ /// positions are UTF-16 code units (the LSP default).
42+ utf8_positions: AtomicBool,
3943}
4044
4145const DIAG_DEBOUNCE: Duration = Duration::from_millis(150);
@@ -55,7 +59,11 @@ impl Backend {
5559 .write()
5660 .await
5761 .insert(uri.clone(), entry.clone());
62+ self.doc_updated(uri, entry).await;
63+ }
5864
65+ /// Follow-up work after a document entry has been stored.
66+ async fn doc_updated(&self, uri: Url, entry: Arc<DocumentEntry>) {
5967 // Refresh the workspace symbol index for this file synchronously —
6068 // symbol extraction is cheap (single pass over text) and keeps
6169 // workspace-wide queries consistent without waiting on the debounce.
@@ -152,9 +160,22 @@ impl LanguageServer for Backend {
152160 self.index.write().await.set_root(root);
153161 }
154162
163+ // Positions are byte offsets throughout this server, so prefer the
164+ // `utf-8` encoding — but only if the client offers it. Otherwise we
165+ // must stay on the UTF-16 default (VS Code refuses to start a server
166+ // that answers with an encoding it didn't offer).
167+ let client_supports_utf8 = params
168+ .capabilities
169+ .general
170+ .as_ref()
171+ .and_then(|g| g.position_encodings.as_ref())
172+ .is_some_and(|encs| encs.contains(&PositionEncodingKind::UTF8));
173+ self.utf8_positions
174+ .store(client_supports_utf8, Ordering::Relaxed);
175+
155176 Ok(InitializeResult {
156177 capabilities: ServerCapabilities {
157− position_encoding: Some(PositionEncodingKind::UTF8),
178+ position_encoding: client_supports_utf8.then_some(PositionEncodingKind::UTF8),
158179 text_document_sync: Some(TextDocumentSyncCapability::Kind(
159180 TextDocumentSyncKind::INCREMENTAL,
160181 )),
@@ -279,15 +300,22 @@ impl LanguageServer for Backend {
279300 let uri = params.text_document.uri.clone();
280301 let version = params.text_document.version;
281302
282− // Start from the current text and apply each change in order.
283− let mut text = match self.doc(&uri).await {
284− Some(d) => d.text.clone(),
285− None => String::new(),
303+ let utf8 = self.utf8_positions.load(Ordering::Relaxed);
304+
305+ // Hold the write lock across read-modify-write: tower-lsp may run
306+ // notification handlers concurrently, and two interleaved incremental
307+ // edits applied to the same base text would lose one of them.
308+ let entry = {
309+ let mut docs = self.documents.write().await;
310+ let mut text = docs.get(&uri).map(|d| d.text.clone()).unwrap_or_default();
311+ for change in ¶ms.content_changes {
312+ apply_change(&mut text, change, utf8);
313+ }
314+ let entry = Arc::new(DocumentEntry::new(text, version));
315+ docs.insert(uri.clone(), entry.clone());
316+ entry
286317 };
287− for change in ¶ms.content_changes {
288− apply_change(&mut text, change);
289− }
290− self.set_doc(uri, text, version).await;
318+ self.doc_updated(uri, entry).await;
291319 }
292320
293321 async fn did_close(&self, params: DidCloseTextDocumentParams) {
@@ -303,6 +331,16 @@ impl LanguageServer for Backend {
303331 {
304332 handle.abort();
305333 }
334+ // The index may hold symbols from unsaved buffer edits; resync it
335+ // with what's actually on disk.
336+ if let Ok(path) = params.text_document.uri.to_file_path() {
337+ let mut idx = self.index.write().await;
338+ if path.exists() {
339+ idx.update_from_disk(&path);
340+ } else {
341+ idx.remove(&path);
342+ }
343+ }
306344 self.client
307345 .publish_diagnostics(params.text_document.uri, vec![], None)
308346 .await;
@@ -613,8 +651,12 @@ impl LanguageServer for Backend {
613651 if formatted == doc.text {
614652 return Ok(None);
615653 }
616− let last_line = doc.text.lines().count().saturating_sub(1) as u32;
617− let last_col = doc.text.lines().last().map_or(0, |l| l.len()) as u32;
654+ // End the edit at the true end of the document. `str::lines` drops a
655+ // trailing newline, which would leave it outside the replaced range
656+ // and append an extra blank line after the (newline-terminated)
657+ // formatted text.
658+ let last_line = doc.text.matches('\n').count() as u32;
659+ let last_col = doc.text.rsplit('\n').next().map_or(0, |l| l.len()) as u32;
618660 Ok(Some(vec![TextEdit {
619661 range: Range::new(Position::new(0, 0), Position::new(last_line, last_col)),
620662 new_text: formatted,
@@ -738,6 +780,7 @@ impl LanguageServer for Backend {
738780 #[derive(Clone)]
739781 struct Def {
740782 line: u32,
783+ col: u32,
741784 name: String,
742785 kind: &'static str,
743786 }
@@ -746,6 +789,9 @@ impl LanguageServer for Backend {
746789 for (i, line) in lines.iter().enumerate() {
747790 let trimmed = line.trim_start();
748791 if let Some(rest) = trimmed.strip_prefix("@let ") {
792+ // `rest` is a suffix of `line`, so the name's column is
793+ // however much of the line precedes it.
794+ let col = (line.len() - rest.trim_start().len()) as u32;
749795 let rest = rest.trim();
750796 if let Some(name) = rest.split_whitespace().next() {
751797 let has_body = lines
@@ -758,18 +804,21 @@ impl LanguageServer for Backend {
758804 {
759805 defs.push(Def {
760806 line: i as u32,
807+ col,
761808 name: name.to_string(),
762809 kind: "fn",
763810 });
764811 } else if value_after_name.starts_with('[') {
765812 defs.push(Def {
766813 line: i as u32,
814+ col,
767815 name: name.to_string(),
768816 kind: "define",
769817 });
770818 } else {
771819 defs.push(Def {
772820 line: i as u32,
821+ col,
773822 name: name.to_string(),
774823 kind: "let",
775824 });
@@ -780,7 +829,7 @@ impl LanguageServer for Backend {
780829
781830 let mut lenses = Vec::with_capacity(defs.len());
782831 for def in &defs {
783− let mut count: usize = 0;
832+ let mut locations: Vec<Location> = Vec::new();
784833 let needle = match def.kind {
785834 "fn" => format!("@{}", def.name),
786835 _ => format!("${}", def.name),
@@ -798,27 +847,36 @@ impl LanguageServer for Backend {
798847 Some(c) => !(c.is_ascii_alphanumeric() || c == b'_' || c == b'-'),
799848 };
800849 if ok {
801− count += 1;
850+ locations.push(Location {
851+ uri: uri.clone(),
852+ range: Range::new(
853+ Position::new(i as u32, pos as u32),
854+ Position::new(i as u32, (pos + needle.len()) as u32),
855+ ),
856+ });
802857 }
803858 from = pos + needle.len();
804859 }
805860 }
806861
807− let title = if count == 1 {
862+ let title = if locations.len() == 1 {
808863 "1 reference".to_string()
809864 } else {
810− format!("{} references", count)
865+ format!("{} references", locations.len())
811866 };
812867 // A clickable lens that opens VS Code's references panel at this
813− // definition. The command name matches what VS Code's built-in
868+ // definition. The arguments mirror what VS Code's built-in
814869 // `editor.action.showReferences` accepts: (uri, position, locations).
870+ // The locations are sent precomputed because a reference query at
871+ // the definition site would resolve the `@let` keyword instead.
815872 let line_pos = Position::new(def.line, 0);
873+ let name_pos = Position::new(def.line, def.col);
816874 lenses.push(CodeLens {
817875 range: Range::new(line_pos, line_pos),
818876 command: Some(Command {
819877 title,
820878 command: "htmlang.showReferences".into(),
821− arguments: Some(vec![json!(uri), json!(line_pos)]),
879+ arguments: Some(vec![json!(uri), json!(name_pos), json!(locations)]),
822880 }),
823881 data: None,
824882 });
@@ -917,6 +975,7 @@ async fn main() {
917975 documents: Arc::new(RwLock::new(HashMap::new())),
918976 index: Arc::new(RwLock::new(WorkspaceIndex::new())),
919977 pending_diags: Arc::new(Mutex::new(HashMap::new())),
978+ utf8_positions: AtomicBool::new(false),
920979 });
921980 Server::new(stdin, stdout, socket).serve(service).await;
922981}