fix: restore selection before saving document

src/DocumentWatcher.ts

@@ -4,6 +4,7 @@ import {
 	Selection,
 	TextDocument,
 	TextDocumentSaveReason,
+	TextEdit,
 	TextEditor,
 	TextEditorOptions,
 	window,
@@ -39,6 +40,10 @@ export default class DocumentWatcher {
 
 		const subscriptions: Disposable[] = []
 
+		let edits: TextEdit[]
+
+		let editedDoc: TextDocument | undefined
+
 		let previousSelections: Selection[] = []
 
 		this.handleTextEditorChange(window.activeTextEditor)
@@ -65,29 +70,51 @@ export default class DocumentWatcher {
 
 		subscriptions.push(
 			workspace.onDidSaveTextDocument(doc => {
-				const activeEditor = window.activeTextEditor
-				if (activeEditor && previousSelections.length) {
-					activeEditor.selections = previousSelections
-				}
-
 				if (path.basename(doc.fileName) === '.editorconfig') {
 					this.log('.editorconfig file saved.')
 				}
 			}),
 		)
 
+		subscriptions.push(
+			workspace.onDidChangeTextDocument(e => {
+				const activeEditor = window.activeTextEditor
+				const activeDoc = activeEditor?.document
+				if (
+					!edits.length ||
+					!activeEditor ||
+					activeDoc !== editedDoc ||
+					activeDoc !== e.document
+				) {
+					return
+				}
+				changesLoop: for (const change of e.contentChanges) {
+					for (const edit of edits) {
+						if (edit.newText === change.text) {
+							activeEditor.selections = [...previousSelections]
+							edits = []
+							editedDoc = undefined
+							break changesLoop
+						}
+					}
+				}
+			}),
+		)
+
 		subscriptions.push(
 			workspace.onWillSaveTextDocument(async e => {
 				const activeEditor = window.activeTextEditor
 				const activeDoc = activeEditor?.document
-				if (activeDoc && activeDoc === e.document && activeEditor) {
+				if (activeEditor && activeDoc === e.document) {
 					previousSelections = [...activeEditor.selections]
 				}
 				const transformations = this.calculatePreSaveTransformations(
 					e.document,
 					e.reason,
 				)
 				e.waitUntil(transformations)
+				edits = await transformations
+				editedDoc = activeDoc
 			}),
 		)