From b61a8ba8ee1eced6311924a1eba1aa14c712a924 Mon Sep 17 00:00:00 2001
From: Aram Becker <becker.aram@gmail.com>
Date: Wed, 11 May 2022 23:24:15 +0200
Subject: [PATCH] fix: correctly update blocksort code lenses

---
 src/providers/BlockSortActionProvider.ts     | 25 ++++++++++++++------
 src/providers/BlockSortFormattingProvider.ts |  2 +-
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/providers/BlockSortActionProvider.ts b/src/providers/BlockSortActionProvider.ts
index f4a921e..4e04a25 100644
--- a/src/providers/BlockSortActionProvider.ts
+++ b/src/providers/BlockSortActionProvider.ts
@@ -198,25 +198,36 @@ export default class BlockSortActionProvider
       return markers;
     }
 
-    for (const change of changes) {
+    let lineOffset = 0;
+    for (const change of changes.sort((a, b) => a.range.start.compareTo(b.range.start))) {
       if (token.isCancellationRequested) return markers;
       const { range, text } = change;
-      const lineCountBefore = range.end.line - range.start.line + 1;
+      const { start, end } = range;
+      const lineCountBefore = end.line - start.line + 1;
       const lines = text.split("\n");
 
       // If line count changes, update all markers after the change
-      if (lineCountBefore !== lines.length) {
-        const lineCountChange = lines.length - lineCountBefore;
+      const lineCountChange = lines.length - lineCountBefore;
+      if (lineCountChange) {
         for (let i = 0; i < markers.length; i++)
-          if (markers[i].isAfter(range.end)) markers[i] = markers[i].translate(lineCountChange, 0);
+          if (markers[i].isAfter(end)) markers[i] = markers[i].translate(lineCountChange, 0);
       }
 
       // Delete and recreate all markers inside the changed range
-      for (let i = markers.length - 1; i >= 0; i--) if (range.contains(markers[i])) markers.splice(i, 1);
+      for (let i = markers.length - 1; i >= 0; i--) {
+        if (range.contains(markers[i]) || start.line == markers[i].line || end.line == markers[i].line)
+          markers.splice(i, 1);
+      }
 
       markers.push(
-        ...BlockSortFormattingProvider.getBlockSortMarkers(document, range, token).map(({ range: { start } }) => start)
+        ...BlockSortFormattingProvider.getBlockSortMarkers(
+          document,
+          range.with(start.translate(lineOffset), end.translate(lineCountChange)),
+          token
+        ).map(({ range: { start } }) => start)
       );
+
+      lineOffset += lineCountChange;
     }
 
     markers.sort((a, b) => a.compareTo(b));
diff --git a/src/providers/BlockSortFormattingProvider.ts b/src/providers/BlockSortFormattingProvider.ts
index 10b7a5a..098a11f 100644
--- a/src/providers/BlockSortFormattingProvider.ts
+++ b/src/providers/BlockSortFormattingProvider.ts
@@ -29,7 +29,7 @@ export default class BlockSortFormattingProvider
     const comments = commentMarkers[document.languageId] ?? commentMarkers.default;
     const markerPrefixes = comments.map((comment) => `${comment.start} ${BlockSortFormattingProvider.blockSortMarker}`);
 
-    for (let i = range?.start.line ?? 0; i < (range?.end.line ?? document.lineCount); i++) {
+    for (let i = range?.start.line ?? 0; i <= (range?.end.line ?? document.lineCount - 1); i++) {
       if (token?.isCancellationRequested) return markers;
 
       const line = document.lineAt(i);