index based folding

src/vs/workbench/contrib/notebook/browser/contrib/fold/folding.ts

@@ -45,13 +45,14 @@ export class FoldingController extends Disposable {
 			// folding icon
 
 			const cellViewModel = e.target;
-			const state = viewModel.getFoldingState(cellViewModel);
+			const modelIndex = viewModel.getCellIndex(cellViewModel);
+			const state = viewModel.getFoldingState(modelIndex);
 
 			if (state === CellFoldingState.None) {
 				return;
 			}
 
-			viewModel.setFoldingState(cellViewModel, state === CellFoldingState.Collapsed ? CellFoldingState.Expanded : CellFoldingState.Collapsed);
+			viewModel.setFoldingState(modelIndex, state === CellFoldingState.Collapsed ? CellFoldingState.Expanded : CellFoldingState.Collapsed);
 		}
 
 		return;

src/vs/workbench/contrib/notebook/browser/viewModel/foldingModel.ts

@@ -165,6 +165,7 @@ export enum CellFoldingState {
 
 export interface FoldingRegionDelegate {
 	onDidFoldingRegionChanged: Event<void>;
-	getFoldingStartIndex(cell: CellViewModel): number;
-	getFoldingState(cell: CellViewModel): CellFoldingState;
+	getCellIndex(cell: CellViewModel): number;
+	getFoldingStartIndex(index: number): number;
+	getFoldingState(index: number): CellFoldingState;
 }

src/vs/workbench/contrib/notebook/browser/viewModel/markdownCellViewModel.ts

@@ -39,7 +39,7 @@ export class MarkdownCellViewModel extends BaseCellViewModel implements ICellVie
 	readonly onDidChangeLayout = this._onDidChangeLayout.event;
 
 	get foldingState() {
-		return this.foldingDelegate.getFoldingState(this);
+		return this.foldingDelegate.getFoldingState(this.foldingDelegate.getCellIndex(this));
 	}
 
 	constructor(

src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts

@@ -251,43 +251,30 @@ export class NotebookViewModel extends Disposable implements FoldingRegionDelega
 			this._updateFoldingRanges();
 		}));
 	}
-	getFoldingStartIndex(cell: CellViewModel): number {
-		const modelIndex = this.viewCells.indexOf(cell);
-		if (modelIndex < 0) {
-			return -1;
-		}
 
-		const range = this._foldingModel.regions.findRange(modelIndex + 1);
+	getFoldingStartIndex(index: number): number {
+		const range = this._foldingModel.regions.findRange(index + 1);
 		const startIndex = this._foldingModel.regions.getStartLineNumber(range) - 1;
 		return startIndex;
 	}
 
-	getFoldingState(cell: CellViewModel): CellFoldingState {
-		const modelIndex = this.viewCells.indexOf(cell);
-		if (modelIndex < 0) {
-			return -1;
-		}
-
-		const range = this._foldingModel.regions.findRange(modelIndex + 1);
+	getFoldingState(index: number): CellFoldingState {
+		const range = this._foldingModel.regions.findRange(index + 1);
 		const startIndex = this._foldingModel.regions.getStartLineNumber(range) - 1;
 
-		if (startIndex !== modelIndex) {
+		if (startIndex !== index) {
 			return CellFoldingState.None;
 		}
 
 		return this._foldingModel.regions.isCollapsed(range) ? CellFoldingState.Collapsed : CellFoldingState.Expanded;
 	}
 
-	setFoldingState(cell: CellViewModel, state: CellFoldingState): void {
-		const modelIndex = this.viewCells.indexOf(cell);
-		if (modelIndex < 0) {
-			return;
-		}
+	setFoldingState(index: number, state: CellFoldingState): void {
 
-		const range = this._foldingModel.regions.findRange(modelIndex + 1);
+		const range = this._foldingModel.regions.findRange(index + 1);
 		const startIndex = this._foldingModel.regions.getStartLineNumber(range) - 1;
 
-		if (startIndex !== modelIndex) {
+		if (startIndex !== index) {
 			return;
 		}
 

src/vs/workbench/contrib/notebook/test/notebookFolding.test.ts

@@ -10,7 +10,6 @@ import { withTestNotebook, TestCell } from 'vs/workbench/contrib/notebook/test/t
 import { IBulkEditService } from 'vs/editor/browser/services/bulkEditService';
 import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
 import { FoldingModel, CellFoldingState } from 'vs/workbench/contrib/notebook/browser/viewModel/foldingModel';
-import { CellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel';
 
 suite('Notebook Folding', () => {
 	const instantiationService = new TestInstantiationService();
@@ -96,8 +95,7 @@ suite('Notebook Folding', () => {
 				[['var e = 7;'], 'markdown', CellKind.Markdown, [], {}],
 			],
 			(editor, viewModel) => {
-				const firstCell = viewModel.viewCells[0] as CellViewModel;
-				viewModel.setFoldingState(firstCell, CellFoldingState.Collapsed);
+				viewModel.setFoldingState(0, CellFoldingState.Collapsed);
 				assert.deepEqual(viewModel.getHiddenRanges(), [
 					{ start: 1, length: 6 }
 				]);
@@ -118,8 +116,7 @@ suite('Notebook Folding', () => {
 				[['var e = 7;'], 'markdown', CellKind.Markdown, [], {}],
 			],
 			(editor, viewModel) => {
-				const thirdCell = viewModel.viewCells[2] as CellViewModel;
-				viewModel.setFoldingState(thirdCell, CellFoldingState.Collapsed);
+				viewModel.setFoldingState(2, CellFoldingState.Collapsed);
 				assert.deepEqual(viewModel.getHiddenRanges(), [
 					{ start: 3, length: 2 }
 				]);
@@ -140,8 +137,7 @@ suite('Notebook Folding', () => {
 				[['var e = 7;'], 'markdown', CellKind.Markdown, [], {}],
 			],
 			(editor, viewModel) => {
-				const thirdCell = viewModel.viewCells[2] as CellViewModel;
-				viewModel.setFoldingState(thirdCell, CellFoldingState.Collapsed);
+				viewModel.setFoldingState(2, CellFoldingState.Collapsed);
 				assert.deepEqual(viewModel.getHiddenRanges(), [
 					{ start: 3, length: 4 }
 				]);
@@ -164,19 +160,19 @@ suite('Notebook Folding', () => {
 				[['var e = 7;'], 'markdown', CellKind.Markdown, [], {}],
 			],
 			(editor, viewModel) => {
-				viewModel.setFoldingState(viewModel.viewCells[0] as CellViewModel, CellFoldingState.Collapsed);
+				viewModel.setFoldingState(0, CellFoldingState.Collapsed);
 				assert.deepEqual(viewModel.getHiddenRanges(), [
 					{ start: 1, length: 1 }
 				]);
 
-				viewModel.setFoldingState(viewModel.viewCells[5] as CellViewModel, CellFoldingState.Collapsed);
-				viewModel.setFoldingState(viewModel.viewCells[2] as CellViewModel, CellFoldingState.Collapsed);
+				viewModel.setFoldingState(5, CellFoldingState.Collapsed);
+				viewModel.setFoldingState(2, CellFoldingState.Collapsed);
 				assert.deepEqual(viewModel.getHiddenRanges(), [
 					{ start: 1, length: 1 },
 					{ start: 3, length: 4 }
 				]);
 
-				viewModel.setFoldingState(viewModel.viewCells[2] as CellViewModel, CellFoldingState.Expanded);
+				viewModel.setFoldingState(2, CellFoldingState.Expanded);
 				assert.deepEqual(viewModel.getHiddenRanges(), [
 					{ start: 1, length: 1 },
 					{ start: 6, length: 1 }