feat: working on the undo stack! Huge progress, not tested

Tests are coming, but there are some infrastructure issues with the
tests

internal/action/misc.go

@@ -60,3 +60,23 @@ func (a EnterVisualBlockMode) Execute(m Model) tea.Cmd {
 	m.SetMode(core.VisualBlockMode)
 	return nil
 }
+
+// TODO: Implement count?
+type Undo struct{}
+
+func (a Undo) Execute(m Model) tea.Cmd {
+	win := m.ActiveWindow()
+	buf := m.ActiveBuffer()
+	buf.Undo(win)
+	return nil
+}
+
+// TODO: Implement count?
+type Redo struct{}
+
+func (a Redo) Execute(m Model) tea.Cmd {
+	win := m.ActiveWindow()
+	buf := m.ActiveBuffer()
+	buf.Redo(win)
+	return nil
+}

internal/command/handlers.go

@@ -933,3 +933,29 @@ func cmdListColorschemes(m action.Model, args []string, force bool) tea.Cmd {
 
 	return nil
 }
+
+func cmdUndoList(m action.Model, args []string, force bool) tea.Cmd {
+	_, _ = args, force
+
+	lines := m.ActiveBuffer().UndoStack.List()
+
+	// For now, display an error when empty
+	if len(lines) == 0 {
+		m.SetCommandOutput(&core.CommandOutput{
+			Lines:   []string{"Undo stack is empty"},
+			Inline:  true,
+			IsError: true,
+		})
+		return nil
+	}
+
+	m.SetMode(core.CommandOutputMode)
+	m.SetCommandOutput(&core.CommandOutput{
+		Title:   ":undo",
+		Lines:   lines,
+		Inline:  false,
+		IsError: false,
+	})
+
+	return nil
+}

internal/command/registry.go

@@ -237,4 +237,11 @@ func (r *Registry) registerDefaults() {
 		ShortForm: "colorschemes",
 		Handler:   cmdListColorschemes,
 	})
+
+	// Undo stack commands
+	r.Register(Command{
+		Name:      "undo",
+		ShortForm: "u",
+		Handler:   cmdUndoList,
+	})
 }

internal/core/buffer.go

@@ -29,7 +29,7 @@ type Buffer struct {
 	ReadOnly bool
 
 	// Options BufferOptions
-	// UndoTree TODO: This will be big
+	UndoStack *UndoStack
 }
 
 // ==================================================
@@ -49,6 +49,10 @@ func (b *Buffer) Line(idx int) string {
 // index is out of bounds. This function sets the modified flag.
 func (b *Buffer) SetLine(idx int, content string) {
 	if idx >= 0 && idx < len(b.Lines) {
+		// Record set line in undo stack
+		if b.UndoStack != nil {
+			b.UndoStack.RecordSetLine(idx, b.Lines[idx], content)
+		}
 		b.Lines[idx] = content
 	}
 	b.Modified = true
@@ -64,6 +68,12 @@ func (b *Buffer) InsertLine(idx int, content string) {
 	if idx > len(b.Lines) {
 		idx = len(b.Lines)
 	}
+
+	// Record insert line in undo stack
+	if b.UndoStack != nil {
+		b.UndoStack.RecordInsertLine(idx, content)
+	}
+
 	b.Lines = append(b.Lines[:idx], append([]string{content}, b.Lines[idx:]...)...)
 	b.Modified = true
 }
@@ -72,6 +82,10 @@ func (b *Buffer) InsertLine(idx int, content string) {
 // of bounds. This function sets the modified flag.
 func (b *Buffer) DeleteLine(idx int) {
 	if idx >= 0 && idx < len(b.Lines) {
+		// Record delete line in undo stack
+		if b.UndoStack != nil {
+			b.UndoStack.RecordDeleteLine(idx, b.Lines[idx])
+		}
 		b.Lines = append(b.Lines[:idx], b.Lines[idx+1:]...)
 	}
 	b.Modified = true
@@ -82,6 +96,99 @@ func (b *Buffer) LineCount() int {
 	return len(b.Lines)
 }
 
+// ==================================================
+// Undo Stack
+// ==================================================
+func (b *Buffer) Undo(w *Window) bool {
+	if b.UndoStack == nil {
+		return false
+	}
+
+	block := b.UndoStack.Undo()
+	if block == nil {
+		return false
+	}
+
+	// Apply changes in REVERSE order
+	for i := len(block.Changes) - 1; i >= 0; i-- {
+		change := block.Changes[i]
+
+		// Temporarily disable recording while we undo
+		wasRecording := b.UndoStack.recording
+		b.UndoStack.recording = false
+
+		switch change.Type {
+		case SetLineChange:
+			// Restore old data
+			if change.Line >= 0 && change.Line < len(b.Lines) {
+				b.Lines[change.Line] = change.OldData
+			}
+		case InsertLineChange:
+			// Remove the inserted line
+			if change.Line >= 0 && change.Line < len(b.Lines) {
+				b.Lines = append(b.Lines[:change.Line], b.Lines[change.Line+1:]...)
+			}
+		case DeleteLineChange:
+			// Re-insert the deleted line
+			if change.Line <= len(b.Lines) {
+				b.Lines = append(b.Lines[:change.Line], append([]string{change.OldData}, b.Lines[change.Line:]...)...)
+			}
+		}
+
+		b.UndoStack.recording = wasRecording
+	}
+
+	// Restore cursor position
+	w.SetCursorLine(block.OldCursor.Line)
+	w.SetCursorCol(block.OldCursor.Col)
+
+	return true
+}
+
+func (b *Buffer) Redo(w *Window) bool {
+	if b.UndoStack == nil {
+		return false
+	}
+
+	block := b.UndoStack.Redo()
+	if block == nil {
+		return false
+	}
+
+	// Apply changes in FORWARD order
+	for _, change := range block.Changes {
+		// Temporarily disable recording while we redo
+		wasRecording := b.UndoStack.recording
+		b.UndoStack.recording = false
+
+		switch change.Type {
+		case SetLineChange:
+			// Apply new data
+			if change.Line >= 0 && change.Line < len(b.Lines) {
+				b.Lines[change.Line] = change.NewData
+			}
+		case InsertLineChange:
+			// Re-insert the line
+			if change.Line <= len(b.Lines) {
+				b.Lines = append(b.Lines[:change.Line], append([]string{change.NewData}, b.Lines[change.Line:]...)...)
+			}
+		case DeleteLineChange:
+			// Re-delete the line
+			if change.Line >= 0 && change.Line < len(b.Lines) {
+				b.Lines = append(b.Lines[:change.Line], b.Lines[change.Line+1:]...)
+			}
+		}
+
+		b.UndoStack.recording = wasRecording
+	}
+
+	// Restore cursor position
+	w.SetCursorLine(block.NewCursor.Line)
+	w.SetCursorCol(block.NewCursor.Col)
+
+	return true
+}
+
 // ==================================================
 // Setters
 // ==================================================

internal/core/buffer_builder.go

@@ -21,6 +21,7 @@ func NewBufferBuilder() *BufferBuilder {
 			Loaded:    false,
 			Listed:    false,
 			ReadOnly:  false,
+			UndoStack: NewUndoStack(), // Empty undo stack
 		},
 	}
 }

internal/core/undo.go

@@ -0,0 +1,180 @@
+package core
+
+import (
+	"fmt"
+	"slices"
+)
+
+type ChangeType string
+
+const (
+	SetLineChange    ChangeType = "SetLine"
+	InsertLineChange ChangeType = "InsertLine"
+	DeleteLineChange ChangeType = "DeleteLine"
+)
+
+type Change struct {
+	Type    ChangeType
+	Line    int
+	OldData string
+	NewData string
+}
+
+type ChangeBlock struct {
+	Changes   []Change
+	OldCursor Position // Before OP
+	NewCursor Position // After OP
+}
+
+type UndoStack struct {
+	undoStack []ChangeBlock
+	redoStack []ChangeBlock
+	current   []Change
+	recording bool
+	oldCursor Position
+}
+
+func NewUndoStack() *UndoStack {
+	return &UndoStack{
+		undoStack: []ChangeBlock{},
+		redoStack: []ChangeBlock{},
+		current:   []Change{},
+		recording: false,
+		oldCursor: Position{},
+	}
+}
+
+func (u *UndoStack) BeginBlock(cursor Position) {
+	u.current = []Change{}
+	u.recording = true
+	u.oldCursor = cursor
+}
+
+func (u *UndoStack) EndBlock(cursor Position) {
+	// If not recording or nothing changed, we can exit safely
+	if !u.recording || len(u.current) == 0 {
+		return
+	}
+
+	block := ChangeBlock{
+		Changes:   u.current,
+		OldCursor: u.oldCursor,
+		NewCursor: cursor,
+	}
+
+	u.undoStack = append(u.undoStack, block)
+	u.redoStack = []ChangeBlock{} // Reset old changes, can no longer redo
+
+	u.recording = false
+	u.current = []Change{}
+}
+
+func (u *UndoStack) RecordSetLine(line int, oldData, newData string) {
+	if !u.recording {
+		return
+	}
+	change := Change{
+		Type:    SetLineChange,
+		Line:    line,
+		OldData: oldData,
+		NewData: newData,
+	}
+	u.current = append(u.current, change)
+}
+
+func (u *UndoStack) RecordInsertLine(line int, newData string) {
+	if !u.recording {
+		return
+	}
+	change := Change{
+		Type:    InsertLineChange,
+		Line:    line,
+		NewData: newData,
+	}
+	u.current = append(u.current, change)
+}
+
+func (u *UndoStack) RecordDeleteLine(line int, oldData string) {
+	if !u.recording {
+		return
+	}
+	change := Change{
+		Type:    DeleteLineChange,
+		Line:    line,
+		OldData: oldData,
+	}
+	u.current = append(u.current, change)
+}
+
+func (u *UndoStack) Undo() *ChangeBlock {
+	if len(u.undoStack) == 0 {
+		return nil
+	}
+
+	// Pop from undo stack
+	size := len(u.undoStack)
+	block := u.undoStack[size-1]
+	u.undoStack = u.undoStack[:size-1]
+
+	// Push to redo stack
+	u.redoStack = append(u.redoStack, block)
+
+	return &block
+}
+
+func (u *UndoStack) Redo() *ChangeBlock {
+	if len(u.redoStack) == 0 {
+		return nil
+	}
+
+	// Pop from redo stack
+	size := len(u.redoStack)
+	block := u.redoStack[size-1]
+	u.redoStack = u.redoStack[:size-1]
+
+	// Push to undo stack
+	u.undoStack = append(u.undoStack, block)
+
+	return &block
+}
+
+func (u *UndoStack) CanUndo() bool {
+	return len(u.undoStack) > 0
+}
+
+func (u *UndoStack) CanRedo() bool {
+	return len(u.redoStack) > 0
+}
+
+func (u *UndoStack) Recording() bool {
+	return u.recording
+}
+
+func (u *UndoStack) List() []string {
+	var lines []string
+
+	stack := slices.Clone(u.undoStack)
+	slices.Reverse(stack)
+
+	for _, b := range stack {
+		lines = append(lines, fmt.Sprintf(
+			"block (%d:%d) -> (%d:%d)",
+			b.OldCursor.Line,
+			b.OldCursor.Col,
+			b.NewCursor.Line,
+			b.NewCursor.Col,
+		))
+
+		for _, c := range b.Changes {
+			lines = append(lines, fmt.Sprintf(
+				"\t%q #%d (%s) -> (%s)",
+				c.Type,
+				c.Line,
+				c.OldData,
+				c.NewData,
+			))
+		}
+	}
+
+	return lines
+}

internal/editor/helpers_test.go

@@ -30,6 +30,8 @@ func sendKeys(tm *teatest.TestModel, keys ...string) {
 			tm.Send(tea.KeyMsg{Type: tea.KeyCtrlU})
 		case "ctrl+v":
 			tm.Send(tea.KeyMsg{Type: tea.KeyCtrlV})
+		case "ctrl+r":
+			tm.Send(tea.KeyMsg{Type: tea.KeyCtrlR})
 		case "ctrl+w":
 			tm.Send(tea.KeyMsg{Type: tea.KeyCtrlW})
 		default:

internal/input/handler.go

@@ -59,6 +59,12 @@ func (h *Handler) Handle(m action.Model, key string) tea.Cmd {
 	if key == "esc" {
 		h.Reset()
 		if m.Mode() == core.InsertMode {
+			// Before exiting insert mode, end the block in the undo stack
+			win := m.ActiveWindow()
+			buf := m.ActiveBuffer()
+			if buf.UndoStack != nil {
+				buf.UndoStack.EndBlock(win.Cursor)
+			}
 			m.ExitInsertMode()
 		} else {
 			m.SetMode(core.NormalMode)
@@ -180,7 +186,7 @@ func (h *Handler) handleInitial(m action.Model, kind string, binding any, key st
 		if res, ok := mot.(action.Resolvable); ok {
 			mot = res.Resolve(m)
 		}
-		cmd := mot.Execute(m)
+		cmd := h.executeMotion(m, mot)
 		h.Reset()
 		return cmd
 
@@ -194,7 +200,7 @@ func (h *Handler) handleInitial(m action.Model, kind string, binding any, key st
 			if m.Mode() == core.VisualLineMode {
 				mtype = core.Linewise
 			}
-			cmd := op.Operate(m, start, end, mtype)
+			cmd := h.executeOperator(m, op, start, end, mtype)
 			// Only reset to normal mode if operator didn't enter insert mode
 			if m.Mode() != core.InsertMode {
 				m.SetMode(core.NormalMode)
@@ -213,7 +219,7 @@ func (h *Handler) handleInitial(m action.Model, kind string, binding any, key st
 		if r, ok := act.(action.Repeatable); ok {
 			act = r.WithCount(count)
 		}
-		cmd := act.Execute(m)
+		cmd := h.executeAction(m, act)
 		h.Reset()
 		return cmd
 	}
@@ -232,7 +238,7 @@ func (h *Handler) handleAfterOperator(m action.Model, kind string, binding any,
 	if kind == "operator" && key == h.operatorKey {
 		// Only call DoublePress if the operator supports it
 		if dp, ok := h.operator.(action.DoublePresser); ok {
-			cmd := dp.DoublePress(m, count)
+			cmd := h.executeDoublePress(m, dp, count)
 			h.Reset()
 			return cmd
 		}
@@ -258,9 +264,9 @@ func (h *Handler) handleAfterOperator(m action.Model, kind string, binding any,
 		}
 		// Get range and motion type
 		start := win.Cursor
-		mot.Execute(m)
+		h.executeMotion(m, mot)
 		end := win.Cursor
-		cmd := h.operator.Operate(m, start, end, mot.Type())
+		cmd := h.executeOperator(m, h.operator, start, end, mot.Type())
 		h.Reset()
 		return cmd
 	}
@@ -333,15 +339,15 @@ func (h *Handler) handleCharMotion(m action.Model, key string) tea.Cmd {
 	if h.operator != nil {
 		win := m.ActiveWindow()
 		start := win.Cursor
-		mot.Execute(m)
+		h.executeMotion(m, mot)
 		end := win.Cursor
-		cmd := h.operator.Operate(m, start, end, mot.Type())
+		cmd := h.executeOperator(m, h.operator, start, end, mot.Type())
 		h.Reset()
 		return cmd
 	}
 
 	// Otherwise just execute the motion
-	cmd := mot.Execute(m)
+	cmd := h.executeMotion(m, mot)
 	h.Reset()
 	return cmd
 }
@@ -366,7 +372,7 @@ func (h *Handler) handleTextObject(m action.Model, kind string, binding any, key
 
 	// If we have an operator pending (e.g., "diw")
 	if h.operator != nil {
-		cmd := h.operator.Operate(m, start, end, mtype)
+		cmd := h.executeOperator(m, h.operator, start, end, mtype)
 		h.Reset()
 		return cmd
 	}
@@ -468,7 +474,18 @@ func (h *Handler) Pending() string {
 
 // Handler.handleInsertKey: Processes a keypress in insert mode, recording it
 // for count replay and executing it as an action or character insertion.
+//
+// This function does not make use of the execute abstractions, to prevent each
+// key inserted from creating a new block in the undo stack.
 func (h *Handler) handleInsertKey(m action.Model, key string) tea.Cmd {
+	buf := m.ActiveBuffer()
+	win := m.ActiveWindow()
+
+	// Start undo block on first insert key
+	if buf.UndoStack != nil && !buf.UndoStack.Recording() {
+		buf.UndoStack.BeginBlock(win.Cursor)
+	}
+
 	// Record the key for count replay (e.g. 5i...)
 	m.SetInsertKeys(append(m.InsertKeys(), key))
 
@@ -486,7 +503,8 @@ func (h *Handler) handleInsertKey(m action.Model, key string) tea.Cmd {
 }
 
 // Handler.handleCommandKey: Processes a keypress in command mode, executing
-// it as an action or inserting it into the command line.
+// it as an action or inserting it into the command line. This does not record
+// anything into the undo stack.
 func (h *Handler) handleCommandKey(m action.Model, key string) tea.Cmd {
 	kind, binding := h.commandKeymap.Lookup(key)
 	switch kind {
@@ -511,3 +529,59 @@ func normalizeVisualSelection(m action.Model) (core.Position, core.Position) {
 	}
 	return c, a
 }
+
+func (h *Handler) executeAction(m action.Model, act action.Action) tea.Cmd {
+	win := m.ActiveWindow()
+	buf := m.ActiveBuffer()
+
+	if buf.UndoStack != nil {
+		buf.UndoStack.BeginBlock(win.Cursor)
+	}
+
+	cmd := act.Execute(m)
+
+	if buf.UndoStack != nil {
+		buf.UndoStack.EndBlock(win.Cursor)
+	}
+
+	return cmd
+}
+
+func (h *Handler) executeMotion(m action.Model, mot action.Motion) tea.Cmd {
+	// These do not change the buffer, so no need to record anything
+	return mot.Execute(m)
+}
+
+func (h *Handler) executeOperator(m action.Model, op action.Operator, start, end core.Position, mtype core.MotionType) tea.Cmd {
+	buf := m.ActiveBuffer()
+	win := m.ActiveWindow()
+
+	if buf.UndoStack != nil {
+		buf.UndoStack.BeginBlock(win.Cursor)
+	}
+
+	cmd := op.Operate(m, start, end, mtype)
+
+	if buf.UndoStack != nil {
+		buf.UndoStack.EndBlock(win.Cursor)
+	}
+
+	return cmd
+}
+
+func (h *Handler) executeDoublePress(m action.Model, dp action.DoublePresser, count int) tea.Cmd {
+	buf := m.ActiveBuffer()
+	win := m.ActiveWindow()
+
+	if buf.UndoStack != nil {
+		buf.UndoStack.BeginBlock(win.Cursor)
+	}
+
+	cmd := dp.DoublePress(m, count)
+
+	if buf.UndoStack != nil {
+		buf.UndoStack.EndBlock(win.Cursor)
+	}
+
+	return cmd
+}

internal/input/keymap.go

@@ -69,6 +69,8 @@ func NewNormalKeymap() *Keymap {
 			"S":      action.SubstituteLine{Count: 1},
 			"p":      action.Paste{Count: 1},
 			"P":      action.PasteBefore{Count: 1},
+			"u":      action.Undo{},
+			"ctrl+r": action.Redo{},
 		},
 		charMotions: map[string]action.Motion{
 			"f": action.FindChar{Forward: true, Inclusive: true, Repeated: false},