1 files changed, 91 insertions, 7 deletions

diff --git a/src/Editor.cxx b/src/Editor.cxx
index b60d96643..8a7b15719 100644
--- a/src/Editor.cxx
+++ b/src/Editor.cxx
@@ -926,6 +926,37 @@ void Editor::SetLastXChosen() {
 	lastXChosen = static_cast<int>(pt.x) + xOffset;
 }
 
+void Editor::RememberSelectionForUndo(int index) {
+	EnsureModelState();
+	if (modelState) {
+		modelState->RememberSelectionForUndo(index, sel);
+		needRedoRemembered = true;
+		// Remember selection at end of processing current message
+	}
+}
+
+void Editor::RememberSelectionOntoStack(int index) {
+	EnsureModelState();
+	if (modelState) {
+		// Is undo currently inside a group?
+		if (!pdoc->AfterUndoSequenceStart()) {
+			// Don't remember selections inside a grouped sequence as can only
+			// unto or redo to the start and end of the group.
+			modelState->RememberSelectionOntoStack(index);
+		}
+	}
+}
+
+void Editor::RememberCurrentSelectionForRedoOntoStack() {
+	if (needRedoRemembered && (pdoc->UndoSequenceDepth() == 0)) {
+		EnsureModelState();
+		if (modelState) {
+			modelState->RememberSelectionForRedoOntoStack(pdoc->UndoCurrent(), sel);
+			needRedoRemembered = false;
+		}
+	}
+}
+
 void Editor::ScrollTo(Sci::Line line, bool moveThumb) {
 	const Sci::Line topLineNew = std::clamp<Sci::Line>(line, 0, MaxScrollPos());
 	if (topLineNew != topLine) {
@@ -2085,13 +2116,13 @@ void Editor::InsertCharacter(std::string_view sv, CharacterSource charSource) {
 				}
 			}
 		}
+		ThinRectangularRange();
 	}
 	if (wrapOccurred) {
 		SetScrollBars();
 		SetVerticalScrollPos();
 		Redraw();
 	}
-	ThinRectangularRange();
 	// If in wrap mode rewrap current line so EnsureCaretVisible has accurate information
 	EnsureCaretVisible();
 	// Avoid blinking during rapid typing:
@@ -2367,22 +2398,43 @@ void Editor::SelectAll() {
 	Redraw();
 }
 
+void Editor::RestoreSelection(Sci::Position newPos, UndoRedo history) {
+	if (rememberingSelectionForUndo && modelState) {
+		// Undo wants the element after the current as it just undid it
+		const int index = pdoc->UndoCurrent() + (history == UndoRedo::undo ? 1 : 0);
+		const SelectionSimple *pss = modelState->SelectionFromStack(index, history);
+		if (pss) {
+			sel.selType = pss->selType;
+			if (sel.IsRectangular()) {
+				sel.Rectangular() = pss->rangeRectangular;
+				// Reconstitute ranges from rectangular range
+				SetRectangularRange();
+			} else {
+				sel.SetRanges(pss->ranges);
+			}
+			// Unsure if this is safe with SetMain potentially failing if document doesn't appear the same.
+			// Maybe this can occur if the user changes font or wrap mode?
+			sel.SetMain(pss->mainRange);
+			newPos = -1; // Used selection from stack so don't use position returned from undo/redo.
+		}
+	}
+	if (newPos >= 0)
+		SetEmptySelection(newPos);
+	EnsureCaretVisible();
+}
+
 void Editor::Undo() {
 	if (pdoc->CanUndo()) {
 		InvalidateCaret();
 		const Sci::Position newPos = pdoc->Undo();
-		if (newPos >= 0)
-			SetEmptySelection(newPos);
-		EnsureCaretVisible();
+		RestoreSelection(newPos, UndoRedo::undo);
 	}
 }
 
 void Editor::Redo() {
 	if (pdoc->CanRedo()) {
 		const Sci::Position newPos = pdoc->Redo();
-		if (newPos >= 0)
-			SetEmptySelection(newPos);
-		EnsureCaretVisible();
+		RestoreSelection(newPos, UndoRedo::redo);
 	}
 }
 
@@ -2458,6 +2510,17 @@ void Editor::NotifyErrorOccurred(Document *, void *, Status status) {
 	errorStatus = status;
 }
 
+void Editor::NotifyGroupCompleted(Document *, void *) noexcept {
+	// RememberCurrentSelectionForRedoOntoStack may throw (for memory exhaustion)
+	// but this method may not as it is called in UndoGroup destructor so ignore
+	// exception.
+	try {
+		RememberCurrentSelectionForRedoOntoStack();
+	} catch (...) {
+		// Ignore any exception
+	}
+}
+
 void Editor::NotifyChar(int ch, CharacterSource charSource) {
 	NotificationData scn = {};
 	scn.nmhdr.code = Notification::CharAdded;
@@ -2726,6 +2789,14 @@ void Editor::NotifyModified(Document *, DocModification mh, void *) {
 			view.llc.Invalidate(LineLayout::ValidLevel::checkTextAndStyle);
 		}
 	} else {
+		if (rememberingSelectionForUndo && FlagSet(mh.modificationType, ModificationFlags::User)) {
+			if (FlagSet(mh.modificationType, ModificationFlags::BeforeInsert | ModificationFlags::BeforeDelete)) {
+				RememberSelectionForUndo(pdoc->UndoCurrent());
+			}
+			if (FlagSet(mh.modificationType, ModificationFlags::InsertText | ModificationFlags::DeleteText)) {
+				RememberSelectionOntoStack(pdoc->UndoCurrent());
+			}
+		}
 		// Move selection and brace highlights
 		if (FlagSet(mh.modificationType, ModificationFlags::InsertText)) {
 			sel.MovePositions(true, mh.position, mh.length);
@@ -5460,6 +5531,7 @@ void Editor::SetDocPointer(Document *document) {
 		pdoc = document;
 	}
 	pdoc->AddRef();
+	modelState.reset();
 	pcs = ContractionStateCreate(pdoc->IsLarge());
 
 	// Ensure all positions within document
@@ -8615,6 +8687,13 @@ sptr_t Editor::WndProc(Message iMessage, uptr_t wParam, sptr_t lParam) {
 	case Message::GetChangeHistory:
 		return static_cast<sptr_t>(changeHistoryOption);
 
+	case Message::SetSelectionUndoHistory:
+		rememberingSelectionForUndo = wParam;
+		break;
+
+	case Message::GetSelectionUndoHistory:
+		return rememberingSelectionForUndo;
+
 	case Message::SetExtraAscent:
 		vs.extraAscent = static_cast<int>(wParam);
 		InvalidateStyleRedraw();
@@ -9024,6 +9103,11 @@ sptr_t Editor::WndProc(Message iMessage, uptr_t wParam, sptr_t lParam) {
 	default:
 		return DefWndProc(iMessage, wParam, lParam);
 	}
+
+	// If there was a change that needs its selection saved and it wasn't explicity saved
+	// then do that here.
+	RememberCurrentSelectionForRedoOntoStack();
+
 	//Platform::DebugPrintf("end wnd proc\n");
 	return 0;
 }