Implement branching undo system with tests and updates.

- Added branching model for undo/redo, enabling multiple redo paths and branch selection.
- Updated `UndoNode` to include `parent` and refined hierarchical navigation.
- Extended `UndoSystem` with branching logic for redo operations, supporting sibling branch selection.
- Overhauled tests to validate branching behavior and tree invariants.
- Refined editor command logic for undo/redo with repeat counts and branch selection.
- Enabled test-only introspection hooks for undo tree validation.
- Updated CMake to include test definitions (`KTE_TESTS` flag).

CMakeLists.txt

@@ -319,6 +319,9 @@ if (BUILD_TESTS)
             ${SYNTAX_SOURCES}
     )
 
+    # Allow test-only introspection hooks (guarded in headers) without affecting production builds.
+    target_compile_definitions(kte_tests PRIVATE KTE_TESTS=1)
+
     # Allow tests to include project headers like "Buffer.h"
     target_include_directories(kte_tests PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
 

Command.cc

@@ -3068,6 +3068,8 @@ cmd_undo(CommandContext &ctx)
 	if (auto *u = buf->Undo()) {
 		// Ensure pending batch is finalized so it can be undone
 		u->commit();
+		int repeat = ctx.count > 0 ? ctx.count : 1;
+		for (int i = 0; i < repeat; ++i)
 			u->undo();
 		// Keep cursor within buffer bounds
 		ensure_cursor_visible(ctx.editor, *buf);
@@ -3087,7 +3089,14 @@ cmd_redo(CommandContext &ctx)
 	if (auto *u = buf->Undo()) {
 		// Finalize any pending batch before redoing
 		u->commit();
+		// With branching undo, a universal-argument count selects an alternate redo branch:
+		//  - no count (or 1): redo the active branch
+		//  - n>1: redo the (n-1)th sibling branch from this point and make it active
+		if (ctx.count > 1) {
+			u->redo(ctx.count - 1);
+		} else {
 			u->redo();
+		}
 		ensure_cursor_visible(ctx.editor, *buf);
 		ctx.editor.SetStatus("Redone");
 		return true;

UndoNode.h

@@ -16,6 +16,7 @@ struct UndoNode {
 	int row{};
 	int col{};
 	std::string text;
+	UndoNode *parent = nullptr; // previous state; null means pre-first-edit
 	UndoNode *child  = nullptr; // next in current timeline
 	UndoNode *next   = nullptr; // redo branch
 };

UndoNodePool.h

@@ -20,6 +20,7 @@ public:
 		available_.pop();
 		// Node comes zeroed; ensure links are reset
 		node->text.clear();
+		node->parent = nullptr;
 		node->child  = nullptr;
 		node->next   = nullptr;
 		node->row    = node->col = 0;
@@ -34,6 +35,7 @@ public:
 			return;
 		// Clear heavy fields to free memory held by strings
 		node->text.clear();
+		node->parent = nullptr;
 		node->child  = nullptr;
 		node->next   = nullptr;
 		node->row    = node->col = 0;

UndoSystem.cc

@@ -69,6 +69,7 @@ UndoSystem::Begin(UndoType type)
 	tree_.pending->row  = row;
 	tree_.pending->col  = col;
 	tree_.pending->text.clear();
+	tree_.pending->parent = nullptr;
 	tree_.pending->child  = nullptr;
 	tree_.pending->next   = nullptr;
 	pending_mode_         = PendingAppendMode::Append;
@@ -119,36 +120,27 @@ UndoSystem::commit()
 		return;
 	}
 
-	// Linear semantics: if we are not at the tip, discard redo.
-	if (tree_.current && tree_.current->child) {
-		// Prevent dangling `saved` pointer if it sits in the discarded redo chain.
-		if (tree_.saved && is_descendant(tree_.current->child, tree_.saved)) {
-			tree_.saved = nullptr;
-		}
-		free_branch(tree_.current->child);
-		tree_.current->child = nullptr;
-	}
-
 	if (!tree_.root) {
 		tree_.root            = tree_.pending;
+		tree_.pending->parent = nullptr;
 		tree_.current         = tree_.pending;
 	} else if (!tree_.current) {
-		// We are at the "pre-first-edit" state. Attach as the new root child.
+		// We are at the "pre-first-edit" state (undo past the first node).
-		// For v1 linear history, this means starting the chain anew.
+		// In branching history, preserve the existing root chain as an alternate branch.
-		// The existing root represents edits from the past; attach the new node as the new root.
+		tree_.pending->parent = nullptr;
-		// (This situation happens after undoing past the first node.)
+		tree_.pending->next   = tree_.root;
-		if (tree_.saved && is_descendant(tree_.root, tree_.saved)) {
-			// ok
-		}
-		// Discard the old root chain because it is redo from the pre-edit state.
-		if (tree_.saved && is_descendant(tree_.root, tree_.saved)) {
-			tree_.saved = nullptr;
-		}
-		free_node(tree_.root);
 		tree_.root            = tree_.pending;
 		tree_.current         = tree_.pending;
 	} else {
+		// Branching semantics: attach as a new redo branch under current.
+		// Make the new edit the active child by inserting it at the head.
+		tree_.pending->parent = tree_.current;
+		if (!tree_.current->child) {
 			tree_.current->child = tree_.pending;
+		} else {
+			tree_.pending->next  = tree_.current->child;
+			tree_.current->child = tree_.pending;
+		}
 		tree_.current = tree_.pending;
 	}
 
@@ -167,27 +159,44 @@ UndoSystem::undo()
 		return;
 	debug_log("undo");
 	apply(tree_.current, -1);
-	UndoNode *parent = find_parent(tree_.root, tree_.current);
+	tree_.current = tree_.current->parent;
-	tree_.current    = parent;
 	update_dirty_flag();
 }
 
 
 void
-UndoSystem::redo()
+UndoSystem::redo(int branch_index)
 {
 	commit();
-	UndoNode *next = nullptr;
+	UndoNode **head = nullptr;
 	if (!tree_.current) {
-		next = tree_.root;
+		head = &tree_.root;
 	} else {
-		next = tree_.current->child;
+		head = &tree_.current->child;
 	}
-	if (!next)
+	if (!head || !*head)
 		return;
+	if (branch_index < 0)
+		branch_index = 0;
+
+	// Select the Nth sibling from the branch list and make it the active head.
+	UndoNode *prev = nullptr;
+	UndoNode *sel  = *head;
+	for (int i = 0; i < branch_index && sel; ++i) {
+		prev = sel;
+		sel  = sel->next;
+	}
+	if (!sel)
+		return;
+	if (prev) {
+		prev->next = sel->next;
+		sel->next  = *head;
+		*head      = sel;
+	}
+
 	debug_log("redo");
-	apply(next, +1);
+	apply(*head, +1);
-	tree_.current = next;
+	tree_.current = *head;
 	update_dirty_flag();
 }
 

UndoSystem.h

@@ -22,7 +22,10 @@ public:
 
 	void undo();
 
-	void redo();
+	// Redo the current node's active child branch.
+	// If `branch_index` > 0, selects that redo sibling (0-based) and makes it active.
+	// When current is null (pre-first-edit), branches are selected among `tree_.root` siblings.
+	void redo(int branch_index = 0);
 
 	void mark_saved();
 
@@ -32,6 +35,14 @@ public:
 
 	void UpdateBufferReference(Buffer &new_buf);
 
+#if defined(KTE_TESTS)
+	// Test-only introspection hook.
+	const UndoTree &TreeForTests() const
+	{
+		return tree_;
+	}
+#endif
+
 private:
 	enum class PendingAppendMode : std::uint8_t {
 		Append,

tests/test_undo.cc

@@ -1,6 +1,57 @@
 #include "Test.h"
 #include "Buffer.h"
 
+#include "Command.h"
+#include "Editor.h"
+
+#include <cstddef>
+#include <random>
+
+#if defined(KTE_TESTS)
+#include <unordered_set>
+
+static void
+validate_undo_subtree(const UndoNode *node, const UndoNode *expected_parent,
+                      std::unordered_set<const UndoNode *> &seen)
+{
+	ASSERT_TRUE(node != nullptr);
+	ASSERT_TRUE(seen.find(node) == seen.end());
+	seen.insert(node);
+	ASSERT_TRUE(node->parent == expected_parent);
+
+	// Validate each redo branch under this node.
+	for (const UndoNode *ch = node->child; ch != nullptr; ch = ch->next) {
+		validate_undo_subtree(ch, node, seen);
+	}
+}
+
+
+static void
+validate_undo_tree(const UndoSystem &u)
+{
+	const UndoTree &t = u.TreeForTests();
+
+	std::unordered_set<const UndoNode *> seen;
+	for (const UndoNode *root = t.root; root != nullptr; root = root->next) {
+		validate_undo_subtree(root, nullptr, seen);
+	}
+
+	// current/saved must either be null or be reachable from some root.
+	if (t.current)
+		ASSERT_TRUE(seen.find(t.current) != seen.end());
+	if (t.saved)
+		ASSERT_TRUE(seen.find(t.saved) != seen.end());
+
+	// pending is detached (not part of the committed tree).
+	if (t.pending) {
+		ASSERT_TRUE(seen.find(t.pending) == seen.end());
+		ASSERT_TRUE(t.pending->parent == nullptr);
+		ASSERT_TRUE(t.pending->child == nullptr);
+		ASSERT_TRUE(t.pending->next == nullptr);
+	}
+}
+#endif
+
 
 TEST (Undo_InsertRun_Coalesces)
 {
@@ -67,7 +118,7 @@ TEST (Undo_BackspaceRun_Coalesces)
 }
 
 
-TEST (Undo_Linear_RedoDiscardedAfterNewEdit)
+TEST (Undo_Branching_RedoPreservedAfterNewEdit)
 {
 	Buffer b;
 	UndoSystem *u = b.Undo();
@@ -91,7 +142,7 @@ TEST (Undo_Linear_RedoDiscardedAfterNewEdit)
 	u->undo();
 	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
 
-	// New edit after undo should discard redo.
+	// New edit after undo creates a new branch; the old redo should remain as an alternate branch.
 	u->Begin(UndoType::Insert);
 	b.insert_text(0, 1, std::string_view("c"));
 	u->Append('c');
@@ -99,8 +150,16 @@ TEST (Undo_Linear_RedoDiscardedAfterNewEdit)
 	u->commit();
 
 	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+
+	// No further redo from the tip.
 	u->redo();
 	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+
+	// Undo back to the branch point and redo the original branch.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+	u->redo(1);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
 }
 
 
@@ -233,7 +292,7 @@ TEST (Undo_UndoPastFirstEdit_RedoFromPreFirstEdit)
 }
 
 
-TEST (Undo_NewEditFromPreFirstEdit_DiscardsOldHistory)
+TEST (Undo_NewEditFromPreFirstEdit_PreservesOldHistoryAsAlternateRootBranch)
 {
 	Buffer b;
 	UndoSystem *u = b.Undo();
@@ -271,9 +330,15 @@ TEST (Undo_NewEditFromPreFirstEdit_DiscardsOldHistory)
 
 	ASSERT_EQ(std::string(b.Rows()[0]), std::string("x"));
 
-	// Old history should be gone: redo should not resurrect "ab".
+	// From the tip, no further redo.
 	u->redo();
 	ASSERT_EQ(std::string(b.Rows()[0]), std::string("x"));
+
+	// Undo back to pre-first-edit and select the older root branch.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string(""));
+	u->redo(1);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
 }
 
 
@@ -341,3 +406,535 @@ TEST (Undo_DeleteIndent_UndoRestoresCursorAtText)
 	ASSERT_EQ(b.Cury(), (std::size_t) 1);
 	ASSERT_EQ(b.Curx(), (std::size_t) 2);
 }
+
+
+TEST (Undo_StructuralInvariants_BranchingAndRoots)
+{
+	Buffer b;
+	UndoSystem *u = b.Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	// Build history: a -> b
+	b.SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 0, std::string_view("a"));
+	u->Append('a');
+	b.SetCursor(1, 0);
+	u->commit();
+
+	b.SetCursor(1, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("b"));
+	u->Append('b');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+
+	// Undo past first edit; now create a new root-level branch x.
+	u->undo();
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string(""));
+
+	b.SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 0, std::string_view("x"));
+	u->Append('x');
+	b.SetCursor(1, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("x"));
+
+	// Return to the older root branch.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string(""));
+	u->redo(1);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Create a normal branch under 'a'.
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("c"));
+	u->Append('c');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_BranchSelection_ThreeSiblingsAndHeadPersists)
+{
+	Buffer b;
+	UndoSystem *u = b.Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	// Root: a
+	b.SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 0, std::string_view("a"));
+	u->Append('a');
+	b.SetCursor(1, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Branch 1: a->b
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("b"));
+	u->Append('b');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+
+	// Back to branch point.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Branch 2: a->c
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("c"));
+	u->Append('c');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Branch 3: a->d
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("d"));
+	u->Append('d');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ad"));
+
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Under 'a', the sibling list should now contain 3 branches.
+	validate_undo_tree(*u);
+
+	// Select the 3rd sibling (branch_index=2) which should be the oldest ("b"), and make it active.
+	u->redo(2);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Since we selected "b", redo with default should now follow "b" again.
+	u->redo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Select another branch by index and ensure it becomes the new default.
+	u->redo(1);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ad"));
+	u->undo();
+	u->redo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ad"));
+	u->undo();
+
+	// Out-of-range selection should be a no-op.
+	u->redo(99);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_Branching_SwitchBetweenTwoRedoBranches_TextAndCursor)
+{
+	Buffer b;
+	UndoSystem *u = b.Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	// Build A->B.
+	b.SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 0, std::string_view("a"));
+	u->Append('a');
+	b.SetCursor(1, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+	ASSERT_EQ(b.Cury(), (std::size_t) 0);
+	ASSERT_EQ(b.Curx(), (std::size_t) 1);
+
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("b"));
+	u->Append('b');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+	ASSERT_EQ(b.Curx(), (std::size_t) 2);
+
+	// Undo to A.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+	ASSERT_EQ(b.Curx(), (std::size_t) 1);
+
+	// Create sibling branch A->C.
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("c"));
+	u->Append('c');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+	ASSERT_EQ(b.Curx(), (std::size_t) 2);
+
+	// Back to A.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+	ASSERT_EQ(b.Curx(), (std::size_t) 1);
+
+	// Redo into B as the alternate branch (older sibling), and confirm cursor is consistent.
+	u->redo(1);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+	ASSERT_EQ(b.Curx(), (std::size_t) 2);
+
+	// Both branches remain reachable: undo to A, redo defaults to B (head reordered).
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+	u->redo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+
+	// And the other branch C should still be selectable.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+	u->redo(1);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+	ASSERT_EQ(b.Curx(), (std::size_t) 2);
+
+	// After selecting C, default redo from A should now follow C.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+	u->redo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_Randomized_Deterministic_EditUndoRedoBranchSelect)
+{
+	Buffer b;
+	UndoSystem *u = b.Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	std::mt19937 rng(0xC0FFEEu);
+	std::uniform_int_distribution<int> op(0, 99);
+	std::uniform_int_distribution<int> ch(0, 25);
+
+	const int steps      = 300;
+	const int max_len    = 40;
+	const int max_branch = 4;
+
+	for (int i = 0; i < steps; ++i) {
+		ASSERT_TRUE(!b.Rows().empty());
+		ASSERT_EQ(b.Cury(), (std::size_t) 0);
+		ASSERT_EQ(b.Rows().size(), (std::size_t) 1);
+		ASSERT_TRUE(b.Curx() <= b.Rows()[0].size());
+
+		validate_undo_tree(*u);
+
+		int r           = op(rng);
+		std::string cur = std::string(b.Rows()[0]);
+		int len         = static_cast<int>(cur.size());
+
+		if (r < 40 && len < max_len) {
+			// Insert one char at end as a standalone committed node.
+			char c = static_cast<char>('a' + ch(rng));
+			b.SetCursor(static_cast<std::size_t>(len), 0);
+			u->Begin(UndoType::Insert);
+			b.insert_text(0, len, std::string_view(&c, 1));
+			u->Append(c);
+			b.SetCursor(static_cast<std::size_t>(len + 1), 0);
+			u->commit();
+		} else if (r < 60 && len > 0) {
+			// Backspace at end as a standalone committed node.
+			char deleted = cur[static_cast<std::size_t>(len - 1)];
+			b.delete_text(0, len - 1, 1);
+			b.SetCursor(static_cast<std::size_t>(len - 1), 0);
+			u->Begin(UndoType::Delete);
+			u->Append(deleted);
+			u->commit();
+		} else if (r < 80) {
+			// Undo then redo should round-trip to the exact same node/text/cursor when possible.
+			const UndoNode *before_node = u->TreeForTests().current;
+			const std::string before_text(std::string(b.Rows()[0]));
+			const std::size_t before_x = b.Curx();
+
+			if (before_node) {
+				u->undo();
+				u->redo();
+				ASSERT_TRUE(u->TreeForTests().current == before_node);
+				ASSERT_EQ(std::string(b.Rows()[0]), before_text);
+				ASSERT_EQ(b.Curx(), before_x);
+			} else {
+				// Nothing to undo; just exercise redo/branch-select paths.
+				u->redo();
+			}
+		} else if (r < 90) {
+			u->undo();
+		} else {
+			int idx = static_cast<int>(rng() % static_cast<std::uint32_t>(max_branch));
+			if ((rng() % 8u) == 0u)
+				idx = 99; // intentionally out of range sometimes
+			u->redo(idx);
+		}
+	}
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_PendingCoalescedRun_UndoCommitsThenUndoes)
+{
+	Buffer b;
+	UndoSystem *u = b.Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	// Create a coalesced insert run without an explicit commit.
+	b.SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 0, std::string_view("a"));
+	u->Append('a');
+	b.SetCursor(1, 0);
+
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("b"));
+	u->Append('b');
+	b.SetCursor(2, 0);
+
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+
+	// undo() should implicitly commit pending and then undo it as one step.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string(""));
+
+	u->redo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_PendingRunAtBranchPoint_UndoThenBranchSelectionStillWorks)
+{
+	Buffer b;
+	UndoSystem *u = b.Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	// Build a->b.
+	b.SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 0, std::string_view("a"));
+	u->Append('a');
+	b.SetCursor(1, 0);
+	u->commit();
+
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("b"));
+	u->Append('b');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+
+	// Undo to the branch point.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Start a pending insert "c" at the branch point, but don't commit.
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("c"));
+	u->Append('c');
+	b.SetCursor(2, 0);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+
+	// Undo should seal the pending "c" as a new branch, then undo it, leaving us at "a".
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// The active redo should now be "c".
+	u->redo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	// Select the older "b" branch.
+	u->redo(1);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_SavedNodeOnOtherBranch_DirtyClearsWhenReturning)
+{
+	Buffer b;
+	UndoSystem *u = b.Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	// Build a->b and mark saved at the tip.
+	b.SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 0, std::string_view("a"));
+	u->Append('a');
+	b.SetCursor(1, 0);
+	u->commit();
+
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("b"));
+	u->Append('b');
+	b.SetCursor(2, 0);
+	u->commit();
+	u->mark_saved();
+	ASSERT_TRUE(!b.Dirty());
+
+	// Move to a different branch.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("c"));
+	u->Append('c');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ac"));
+	ASSERT_TRUE(b.Dirty());
+
+	// Return to the saved node by selecting the older branch.
+	u->undo();
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("a"));
+	u->redo(1);
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("ab"));
+	ASSERT_TRUE(!b.Dirty());
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_Clear_AfterSaved_ResetsStateSafely)
+{
+	Buffer b;
+	UndoSystem *u = b.Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	b.SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 0, std::string_view("x"));
+	u->Append('x');
+	b.SetCursor(1, 0);
+	u->commit();
+	u->mark_saved();
+	ASSERT_TRUE(!b.Dirty());
+
+	u->Begin(UndoType::Insert);
+	b.insert_text(0, 1, std::string_view("y"));
+	u->Append('y');
+	b.SetCursor(2, 0);
+	u->commit();
+	ASSERT_TRUE(b.Dirty());
+
+	u->clear();
+	ASSERT_TRUE(!b.Dirty());
+	// clear() resets undo history, but does not mutate buffer contents.
+	ASSERT_EQ(std::string(b.Rows()[0]), std::string("xy"));
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_Command_UndoHonorsRepeatCount)
+{
+	InstallDefaultCommands();
+
+	Editor ed;
+	ed.SetDimensions(24, 80);
+
+	Buffer b;
+	ed.AddBuffer(std::move(b));
+	Buffer *buf = ed.CurrentBuffer();
+	ASSERT_TRUE(buf != nullptr);
+	UndoSystem *u = buf->Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	// Create two committed steps using the undo system directly.
+	buf->SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	buf->insert_text(0, 0, std::string_view("a"));
+	u->Append('a');
+	buf->SetCursor(1, 0);
+	u->commit();
+
+	u->Begin(UndoType::Insert);
+	buf->insert_text(0, 1, std::string_view("b"));
+	u->Append('b');
+	buf->SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string("ab"));
+
+	// Undo twice via command repeat count.
+	ed.SetUniversalArg(1, 2);
+	ASSERT_TRUE(Execute(ed, CommandId::Undo));
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string(""));
+
+	validate_undo_tree(*u);
+}
+
+
+TEST (Undo_Command_RedoCountSelectsBranch)
+{
+	InstallDefaultCommands();
+
+	Editor ed;
+	ed.SetDimensions(24, 80);
+
+	Buffer b;
+	ed.AddBuffer(std::move(b));
+	Buffer *buf = ed.CurrentBuffer();
+	ASSERT_TRUE(buf != nullptr);
+	UndoSystem *u = buf->Undo();
+	ASSERT_TRUE(u != nullptr);
+
+	// Build a->b.
+	buf->SetCursor(0, 0);
+	u->Begin(UndoType::Insert);
+	buf->insert_text(0, 0, std::string_view("a"));
+	u->Append('a');
+	buf->SetCursor(1, 0);
+	u->commit();
+
+	u->Begin(UndoType::Insert);
+	buf->insert_text(0, 1, std::string_view("b"));
+	u->Append('b');
+	buf->SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string("ab"));
+
+	// Undo to the branch point and create a sibling branch "c".
+	u->undo();
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string("a"));
+
+	u->Begin(UndoType::Insert);
+	buf->insert_text(0, 1, std::string_view("c"));
+	u->Append('c');
+	buf->SetCursor(2, 0);
+	u->commit();
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string("ac"));
+
+	// Back to branch point.
+	ASSERT_TRUE(Execute(ed, CommandId::Undo));
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string("a"));
+
+	// Command redo with count=2 should select branch_index=1 (the older "b" branch).
+	ed.SetUniversalArg(1, 2);
+	ASSERT_TRUE(Execute(ed, CommandId::Redo));
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string("ab"));
+
+	// After selection, "b" should be the default redo from the branch point.
+	ASSERT_TRUE(Execute(ed, CommandId::Undo));
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string("a"));
+	ASSERT_TRUE(Execute(ed, CommandId::Redo));
+	ASSERT_EQ(std::string(buf->Rows()[0]), std::string("ab"));
+
+	validate_undo_tree(*u);
+}