diff --git a/src/vs/editor/common/services/modelServiceImpl.ts b/src/vs/editor/common/services/modelServiceImpl.ts
index 8667276868a2fdda1f212cf357bdf9668ccb9ba7..668c697a39d8a5504c87095e6b70c6a1d1a8e6ce 100644
--- a/src/vs/editor/common/services/modelServiceImpl.ts
+++ b/src/vs/editor/common/services/modelServiceImpl.ts
@@ -513,58 +513,6 @@ export class ModelServiceImpl extends Disposable implements IModelService {
 		if (!modelData) {
 			return;
 		}
-		const model = modelData.model;
-		const sharesUndoRedoStack = (this._undoRedoService.getUriComparisonKey(model.uri) !== model.uri.toString());
-		let maintainUndoRedoStack = false;
-		let heapSize = 0;
-		if (sharesUndoRedoStack || (this._shouldRestoreUndoStack() && schemaShouldMaintainUndoRedoElements(resource))) {
-			const elements = this._undoRedoService.getElements(resource);
-			if (elements.past.length > 0 || elements.future.length > 0) {
-				for (const element of elements.past) {
-					if (isEditStackElement(element) && element.matchesResource(resource)) {
-						maintainUndoRedoStack = true;
-						heapSize += element.heapSize(resource);
-						element.setModel(resource); // remove reference from text buffer instance
-					}
-				}
-				for (const element of elements.future) {
-					if (isEditStackElement(element) && element.matchesResource(resource)) {
-						maintainUndoRedoStack = true;
-						heapSize += element.heapSize(resource);
-						element.setModel(resource); // remove reference from text buffer instance
-					}
-				}
-			}
-		}
-
-		if (!maintainUndoRedoStack) {
-			if (!sharesUndoRedoStack) {
-				const initialUndoRedoSnapshot = modelData.model.getInitialUndoRedoSnapshot();
-				if (initialUndoRedoSnapshot !== null) {
-					this._undoRedoService.restoreSnapshot(initialUndoRedoSnapshot);
-				}
-			}
-			modelData.model.dispose();
-			return;
-		}
-
-		const maxMemory = ModelServiceImpl.MAX_MEMORY_FOR_CLOSED_FILES_UNDO_STACK;
-		if (!sharesUndoRedoStack && heapSize > maxMemory) {
-			// the undo stack for this file would never fit in the configured memory, so don't bother with it.
-			const initialUndoRedoSnapshot = modelData.model.getInitialUndoRedoSnapshot();
-			if (initialUndoRedoSnapshot !== null) {
-				this._undoRedoService.restoreSnapshot(initialUndoRedoSnapshot);
-			}
-			modelData.model.dispose();
-			return;
-		}
-
-		this._ensureDisposedModelsHeapSize(maxMemory - heapSize);
-
-		// We only invalidate the elements, but they remain in the undo-redo service.
-		this._undoRedoService.setElementsValidFlag(resource, false, (element) => (isEditStackElement(element) && element.matchesResource(resource)));
-		this._insertDisposedModel(new DisposedModelInfo(resource, modelData.model.getInitialUndoRedoSnapshot(), Date.now(), sharesUndoRedoStack, heapSize, computeModelSha1(model), model.getVersionId(), model.getAlternativeVersionId()));
-
 		modelData.model.dispose();
 	}
 
@@ -599,6 +547,50 @@ export class ModelServiceImpl extends Disposable implements IModelService {
 		const modelId = MODEL_ID(model.uri);
 		const modelData = this._models[modelId];
 
+		const sharesUndoRedoStack = (this._undoRedoService.getUriComparisonKey(model.uri) !== model.uri.toString());
+		let maintainUndoRedoStack = false;
+		let heapSize = 0;
+		if (sharesUndoRedoStack || (this._shouldRestoreUndoStack() && schemaShouldMaintainUndoRedoElements(model.uri))) {
+			const elements = this._undoRedoService.getElements(model.uri);
+			if (elements.past.length > 0 || elements.future.length > 0) {
+				for (const element of elements.past) {
+					if (isEditStackElement(element) && element.matchesResource(model.uri)) {
+						maintainUndoRedoStack = true;
+						heapSize += element.heapSize(model.uri);
+						element.setModel(model.uri); // remove reference from text buffer instance
+					}
+				}
+				for (const element of elements.future) {
+					if (isEditStackElement(element) && element.matchesResource(model.uri)) {
+						maintainUndoRedoStack = true;
+						heapSize += element.heapSize(model.uri);
+						element.setModel(model.uri); // remove reference from text buffer instance
+					}
+				}
+			}
+		}
+
+		const maxMemory = ModelServiceImpl.MAX_MEMORY_FOR_CLOSED_FILES_UNDO_STACK;
+		if (!maintainUndoRedoStack) {
+			if (!sharesUndoRedoStack) {
+				const initialUndoRedoSnapshot = modelData.model.getInitialUndoRedoSnapshot();
+				if (initialUndoRedoSnapshot !== null) {
+					this._undoRedoService.restoreSnapshot(initialUndoRedoSnapshot);
+				}
+			}
+		} else if (!sharesUndoRedoStack && heapSize > maxMemory) {
+			// the undo stack for this file would never fit in the configured memory, so don't bother with it.
+			const initialUndoRedoSnapshot = modelData.model.getInitialUndoRedoSnapshot();
+			if (initialUndoRedoSnapshot !== null) {
+				this._undoRedoService.restoreSnapshot(initialUndoRedoSnapshot);
+			}
+		} else {
+			this._ensureDisposedModelsHeapSize(maxMemory - heapSize);
+			// We only invalidate the elements, but they remain in the undo-redo service.
+			this._undoRedoService.setElementsValidFlag(model.uri, false, (element) => (isEditStackElement(element) && element.matchesResource(model.uri)));
+			this._insertDisposedModel(new DisposedModelInfo(model.uri, modelData.model.getInitialUndoRedoSnapshot(), Date.now(), sharesUndoRedoStack, heapSize, computeModelSha1(model), model.getVersionId(), model.getAlternativeVersionId()));
+		}
+
 		delete this._models[modelId];
 		modelData.dispose();