From d2cbf9e1893c91f7d4da3dd9ef0af6821892d358 Mon Sep 17 00:00:00 2001
From: CismonX <admin@cismon.net>
Date: Wed, 31 Mar 2021 20:51:09 +0800
Subject: [PATCH] Re-calculate block information on document save.

---
 src/context_mapping.ts        | 5 ++++-
 src/contexts/folding_range.ts | 9 +++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/context_mapping.ts b/src/context_mapping.ts
index aa4ad40..347a6c7 100644
--- a/src/context_mapping.ts
+++ b/src/context_mapping.ts
@@ -50,7 +50,10 @@ export default class ContextMapping implements vscode.Disposable {
 
     static onDocumentSave(document: vscode.TextDocument) {
         const documentContext = ContextMapping.getDocumentContextIfExist(document);
-        documentContext?.getPreview()?.updateWebview();
+        if (documentContext !== undefined) {
+            documentContext.foldingRange.clear();
+            documentContext.getPreview()?.updateWebview();
+        }
     }
 
     static onDocumentClose(document: vscode.TextDocument) {
diff --git a/src/contexts/folding_range.ts b/src/contexts/folding_range.ts
index 901110f..a9492c1 100644
--- a/src/contexts/folding_range.ts
+++ b/src/contexts/folding_range.ts
@@ -51,6 +51,8 @@ export default class FoldingRangeContext {
 
     private closingSubsection?: number;
 
+    private contentMayChange = true;
+
     /**
      * Get VSCode folding ranges from the context.
      */
@@ -72,6 +74,7 @@ export default class FoldingRangeContext {
      * @param events Events describing the changes in the document.
      */
     update(events: readonly vscode.TextDocumentContentChangeEvent[]) {
+        this.contentMayChange = true;
         if (this.foldingRanges === undefined) return false;
         for (const event of events) {
             const updatedLines = event.text.split(this.document.eol === vscode.EndOfLine.LF ? '\n' : '\r\n').length;
@@ -85,6 +88,11 @@ export default class FoldingRangeContext {
         return false;
     }
 
+    clear() {
+        if (!this.contentMayChange) return;
+        this.foldingRanges = undefined;
+    }
+
     constructor(private readonly document: vscode.TextDocument) {}
 
     /**
@@ -94,6 +102,7 @@ export default class FoldingRangeContext {
      * @param end Ending line number.
      */
     private calculateFoldingRanges() {
+        this.contentMayChange = false;
         this.foldingRanges = [];
         this.clearTemporaries();
         let closingBlocks = <NamedLine[]>[];