[folding] fix restoring collapse state (for #36555)

src/vs/editor/contrib/folding/browser/folding.ts

@@ -17,7 +17,7 @@ import { ICommonCodeEditor, ScrollType } from 'vs/editor/common/editorCommon';
 import { editorAction, ServicesAccessor, EditorAction, CommonEditorRegistry } from 'vs/editor/common/editorCommonExtensions';
 import { ICodeEditor, IEditorMouseEvent, MouseTargetType } from 'vs/editor/browser/editorBrowser';
 import { editorContribution } from 'vs/editor/browser/editorBrowserExtensions';
-import { FoldingModel, FoldingRegion, setCollapseStateAtLevel, setCollapseStateRecursivly, fold, unfold } from 'vs/editor/contrib/folding/common/foldingModel';
+import { FoldingModel, setCollapseStateAtLevel, setCollapseStateRecursivly, fold, unfold, CollapseState } from 'vs/editor/contrib/folding/common/foldingModel';
 import { computeRanges, limitByIndent } from 'vs/editor/contrib/folding/common/indentFoldStrategy';
 import { FoldingDecorationProvider } from './foldingDecorations';
 import { EditorContextKeys } from 'vs/editor/common/editorContextKeys';
@@ -99,27 +99,18 @@ export class FoldingController {
 	/**
 	 * Store view state.
 	 */
-	public saveViewState(): any {
+	public saveViewState(): { collapsedRegions?: CollapseState, lineCount?: number } {
 		let model = this.editor.getModel();
 		if (!model) {
 			return {};
 		}
-		let collapsedIndexes: number[] = [];
-		for (let region of this.foldingModel.regions) {
-			if (region.isCollapsed && region.editorDecorationId) {
-				var range = model.getDecorationRange(region.editorDecorationId);
-				if (range) {
-					collapsedIndexes.push(range.startLineNumber);
-				}
-			}
-		}
-		return { collapsedIndexes, lineCount: model.getLineCount() };
+		return { collapsedRegions: this.foldingModel.getCollapseState(), lineCount: model.getLineCount() };
 	}
 
 	/**
 	 * Restore view state.
 	 */
-	public restoreViewState(state: any): void {
+	public restoreViewState(state: { collapsedRegions?: CollapseState, lineCount?: number }): void {
 		let model = this.editor.getModel();
 		if (!model) {
 			return;
@@ -127,20 +118,17 @@ export class FoldingController {
 		if (!this._isEnabled) {
 			return;
 		}
-		if (!state || !Array.isArray(state.collapsedIndexes) || state.collapsedIndexes.length === 0 || state.lineCount !== model.getLineCount()) {
+		if (!state || !state.collapsedRegions || state.lineCount !== model.getLineCount()) {
 			return;
 		}
+
+		// set the hidden ranges right way, before waiting for the folding model.
+		if (this.hiddenRangeModel.applyCollapseState(state.collapsedRegions)) {
 			this.getFoldingModel().then(foldingModel => {
-			let toToogle: FoldingRegion[] = [];
-			for (let index of state.collapsedIndexes) {
-				let region = foldingModel.getRegionAtLine(index);
-				if (region && !region.isCollapsed) {
-					toToogle.push(region);
-				}
-			}
-			foldingModel.toggleCollapseState(toToogle);
+				foldingModel.applyCollapseState(state.collapsedRegions);
 			});
 		}
+	}
 
 	private onModelChanged(): void {
 		this.localToDispose = dispose(this.localToDispose);

src/vs/editor/contrib/folding/common/foldingModel.ts

@@ -32,6 +32,8 @@ export interface FoldingModelChangeEvent {
 	collapseStateChanged?: FoldingRegion[];
 }
 
+export type CollapseState = ILineRange[];
+
 export class FoldingModel {
 	private _textModel: IModel;
 	private _decorationProvider: IDecorationProvider;
@@ -134,6 +136,44 @@ export class FoldingModel {
 		this._updateEventEmitter.fire({ model: this });
 	}
 
+	/**
+	 * Collapse state, for persistence only
+	 */
+	public getCollapseState(): CollapseState {
+		let collapsedRanges: ILineRange[] = [];
+		for (let region of this._regions) {
+			if (region.isCollapsed && region.editorDecorationId) {
+				let range = this._textModel.getDecorationRange(region.editorDecorationId);
+				if (range) {
+					let startLineNumber = range.startLineNumber;
+					let endLineNumber = range.endLineNumber + region.range.endLineNumber + region.range.startLineNumber;
+					collapsedRanges.push({ startLineNumber, endLineNumber });
+				}
+			}
+		}
+		if (collapsedRanges.length > 0) {
+			return collapsedRanges;
+		}
+		return null;
+	}
+
+	/**
+	 * Apply persisted state, for persistence only
+	 */
+	public applyCollapseState(state: CollapseState) {
+		if (!Array.isArray(state)) {
+			return;
+		}
+		let toToogle: FoldingRegion[] = [];
+		for (let range of state) {
+			let region = this.getRegionAtLine(range.startLineNumber);
+			if (region && !region.isCollapsed) {
+				toToogle.push(region);
+			}
+		}
+		this.toggleCollapseState(toToogle);
+	}
+
 	public dispose() {
 		let editorDecorationIds = [];
 		for (let region of this._regions) {

src/vs/editor/contrib/folding/common/hiddenRangeModel.ts

@@ -5,12 +5,11 @@
 
 import Event, { Emitter } from 'vs/base/common/event';
 import { Range, IRange } from 'vs/editor/common/core/range';
-import { FoldingRegion, FoldingModel, IFoldingRange } from 'vs/editor/contrib/folding/common/foldingModel';
+import { FoldingRegion, FoldingModel, IFoldingRange, CollapseState } from 'vs/editor/contrib/folding/common/foldingModel';
 import { IDisposable } from 'vs/base/common/lifecycle';
 import { Selection } from 'vs/editor/common/core/selection';
 import { findFirst } from 'vs/base/common/arrays';
 
-
 export class HiddenRangeModel {
 	private _foldingModel: FoldingModel;
 	private _hiddenRanges: IRange[] = [];
@@ -53,10 +52,29 @@ export class HiddenRangeModel {
 			}
 		};
 		if (updateHiddenAreas || i < this._hiddenRanges.length) {
+			this.applyHiddenRanges(newHiddenAreas);
+		}
+	}
+
+	public applyCollapseState(state: CollapseState): boolean {
+		if (!Array.isArray(state) || state.length === 0) {
+			return false;
+		}
+		let hiddenRanges = [];
+		for (let r of state) {
+			if (!r.startLineNumber || !r.endLineNumber) {
+				return false;
+			}
+			hiddenRanges.push(new Range(r.startLineNumber, 1, r.endLineNumber, 1));
+		}
+		this.applyHiddenRanges(hiddenRanges);
+		return true;
+	}
+
+	private applyHiddenRanges(newHiddenAreas: IRange[]) {
 		this._hiddenRanges = newHiddenAreas;
 		this._updateEventEmitter.fire(newHiddenAreas);
 	}
-	}
 
 	public hasRanges() {
 		return this._hiddenRanges.length > 0;