HexEditor: Implement undo and redo actions

6 files changed, 204 insertions, 7 deletions

diff --git a/Userland/Applications/HexEditor/HexDocument.cpp b/Userland/Applications/HexEditor/HexDocument.cpp
index 260fc39d83..7b42026f38 100644
--- a/Userland/Applications/HexEditor/HexDocument.cpp
+++ b/Userland/Applications/HexEditor/HexDocument.cpp
@@ -174,3 +174,63 @@ void HexDocumentFile::ensure_position_in_buffer(size_t position)
         m_buffer_file_pos = position;
     }
 }
+
+HexDocumentUndoCommand::HexDocumentUndoCommand(WeakPtr<HexDocument> document, size_t position)
+    : m_document(move(document))
+    , m_position(position)
+{
+}
+
+void HexDocumentUndoCommand::undo()
+{
+    for (size_t i = 0; i < m_old.size(); i++)
+        m_document->set(m_position + i, m_old[i]);
+}
+
+void HexDocumentUndoCommand::redo()
+{
+    for (size_t i = 0; i < m_new.size(); i++)
+        m_document->set(m_position + i, m_new[i]);
+}
+
+bool HexDocumentUndoCommand::merge_with(GUI::Command const& other)
+{
+    if (!is<HexDocumentUndoCommand>(other) || commit_time_expired())
+        return false;
+
+    auto const& typed_other = static_cast<HexDocumentUndoCommand const&>(other);
+
+    size_t relative_start = typed_other.m_position - m_position;
+    size_t other_length = typed_other.m_old.size();
+    size_t length = m_old.size();
+
+    if (typed_other.m_position < m_position || m_position + length < typed_other.m_position)
+        return false;
+
+    m_old.resize(relative_start + other_length);
+    m_new.resize(relative_start + other_length);
+
+    for (size_t i = 0; i < other_length; i++) {
+        m_new[relative_start + i] = typed_other.m_new[i];
+
+        if (relative_start + i >= length)
+            m_old[relative_start + i] = typed_other.m_old[i];
+    }
+
+    m_timestamp = Time::now_monotonic();
+    return true;
+}
+
+ErrorOr<void> HexDocumentUndoCommand::try_add_changed_byte(u8 old_value, u8 new_value)
+{
+    TRY(m_old.try_append(old_value));
+    TRY(m_new.try_append(new_value));
+    return {};
+}
+
+ErrorOr<void> HexDocumentUndoCommand::try_add_changed_bytes(ByteBuffer old_values, ByteBuffer new_values)
+{
+    TRY(m_old.try_append(move(old_values)));
+    TRY(m_new.try_append(move(new_values)));
+    return {};
+}
diff --git a/Userland/Applications/HexEditor/HexDocument.h b/Userland/Applications/HexEditor/HexDocument.h
index ba3f00e5a0..5a170b4a65 100644
--- a/Userland/Applications/HexEditor/HexDocument.h
+++ b/Userland/Applications/HexEditor/HexDocument.h
@@ -7,10 +7,15 @@
 #pragma once
 
 #include <AK/StringView.h>
+#include <AK/Time.h>
 #include <AK/Types.h>
+#include <AK/WeakPtr.h>
 #include <LibCore/File.h>
+#include <LibGUI/Command.h>
 
-class HexDocument {
+constexpr Time COMMAND_COMMIT_TIME = Time::from_milliseconds(400);
+
+class HexDocument : public Weakable<HexDocument> {
 public:
     enum class Type {
         Memory,
@@ -77,3 +82,26 @@ private:
     Array<u8, 2048> m_buffer;
     size_t m_buffer_file_pos;
 };
+
+class HexDocumentUndoCommand : public GUI::Command {
+public:
+    HexDocumentUndoCommand(WeakPtr<HexDocument> document, size_t position);
+
+    virtual void undo() override;
+    virtual void redo() override;
+    virtual String action_text() const override { return "Update cell"; }
+
+    virtual bool merge_with(GUI::Command const& other) override;
+
+    ErrorOr<void> try_add_changed_byte(u8 old_value, u8 new_value);
+    ErrorOr<void> try_add_changed_bytes(ByteBuffer old_values, ByteBuffer new_values);
+
+private:
+    bool commit_time_expired() const { return Time::now_monotonic() - m_timestamp >= COMMAND_COMMIT_TIME; }
+
+    Time m_timestamp = Time::now_monotonic();
+    WeakPtr<HexDocument> m_document;
+    size_t m_position;
+    ByteBuffer m_old;
+    ByteBuffer m_new;
+};
diff --git a/Userland/Applications/HexEditor/HexEditor.cpp b/Userland/Applications/HexEditor/HexEditor.cpp
index 701a461666..5aff8d3f5d 100644
--- a/Userland/Applications/HexEditor/HexEditor.cpp
+++ b/Userland/Applications/HexEditor/HexEditor.cpp
@@ -82,8 +82,22 @@ void HexEditor::fill_selection(u8 fill_byte)
     if (!has_selection())
         return;
 
-    for (size_t i = m_selection_start; i < m_selection_end; i++)
-        m_document->set(i, fill_byte);
+    ByteBuffer old_values;
+    ByteBuffer new_values;
+
+    size_t length = m_selection_end - m_selection_start;
+
+    new_values.resize(length);
+    old_values.resize(length);
+
+    for (size_t i = 0; i < length; i++) {
+        size_t position = m_selection_start + i;
+        old_values[i] = m_document->get(position).value;
+        new_values[i] = fill_byte;
+        m_document->set(position, fill_byte);
+    }
+
+    did_complete_action(m_selection_start, move(old_values), move(new_values));
 
     update();
     did_change();
@@ -465,6 +479,8 @@ void HexEditor::hex_mode_keydown_event(GUI::KeyEvent& event)
 
         VERIFY(m_position <= m_document->size());
 
+        auto old_value = m_document->get(m_position).value;
+
         // yes, this is terrible... but it works.
         auto value = (event.key() >= KeyCode::Key_0 && event.key() <= KeyCode::Key_9)
             ? event.key() - KeyCode::Key_0
@@ -472,11 +488,14 @@ void HexEditor::hex_mode_keydown_event(GUI::KeyEvent& event)
 
         if (!m_cursor_at_low_nibble) {
             u8 existing_change = m_document->get(m_position).value;
-            existing_change = value << 4 | (existing_change & 0xF); // shift new value left 4 bits, OR with existing last 4 bits
-            m_document->set(m_position, existing_change);
+            u8 new_value = value << 4 | (existing_change & 0xF); // shift new value left 4 bits, OR with existing last 4 bits
+            m_document->set(m_position, new_value);
+            did_complete_action(m_position, old_value, new_value);
             m_cursor_at_low_nibble = true;
         } else {
-            m_document->set(m_position, (m_document->get(m_position).value & 0xF0) | value); // save the first 4 bits, OR the new value in the last 4
+            u8 new_value = (m_document->get(m_position).value & 0xF0) | value; // save the first 4 bits, OR the new value in the last 4
+            m_document->set(m_position, new_value);
+            did_complete_action(m_position, old_value, new_value);
             if (m_position + 1 < m_document->size())
                 m_position++;
             m_cursor_at_low_nibble = false;
@@ -498,7 +517,11 @@ void HexEditor::text_mode_keydown_event(GUI::KeyEvent& event)
     if (event.code_point() == 0) // This is a control key
         return;
 
-    m_document->set(m_position, event.code_point());
+    auto old_value = m_document->get(m_position).value;
+    auto new_value = event.code_point();
+    m_document->set(m_position, new_value);
+    did_complete_action(m_position, old_value, new_value);
+
     if (m_position + 1 < m_document->size())
         m_position++;
     m_cursor_at_low_nibble = false;
@@ -797,3 +820,56 @@ void HexEditor::reset_cursor_blink_state()
     m_cursor_blink_active = true;
     m_blink_timer->restart();
 }
+
+void HexEditor::did_complete_action(size_t position, u8 old_value, u8 new_value)
+{
+    if (old_value == new_value)
+        return;
+
+    auto command = make<HexDocumentUndoCommand>(m_document->make_weak_ptr(), position);
+
+    // We know this won't allocate because the buffers start empty
+    MUST(command->try_add_changed_byte(old_value, new_value));
+    // FIXME: Handle errors
+    MUST(m_undo_stack.try_push(move(command)));
+}
+
+void HexEditor::did_complete_action(size_t position, ByteBuffer&& old_values, ByteBuffer&& new_values)
+{
+    auto command = make<HexDocumentUndoCommand>(m_document->make_weak_ptr(), position);
+
+    // FIXME: Handle errors
+    MUST(command->try_add_changed_bytes(move(old_values), move(new_values)));
+    MUST(m_undo_stack.try_push(move(command)));
+}
+
+bool HexEditor::undo()
+{
+    if (!m_undo_stack.can_undo())
+        return false;
+
+    m_undo_stack.undo();
+    reset_cursor_blink_state();
+    update();
+    update_status();
+    did_change();
+    return true;
+}
+
+bool HexEditor::redo()
+{
+    if (!m_undo_stack.can_redo())
+        return false;
+
+    m_undo_stack.redo();
+    reset_cursor_blink_state();
+    update();
+    update_status();
+    did_change();
+    return true;
+}
+
+GUI::UndoStack& HexEditor::undo_stack()
+{
+    return m_undo_stack;
+}
diff --git a/Userland/Applications/HexEditor/HexEditor.h b/Userland/Applications/HexEditor/HexEditor.h
index 0bda644ddd..1b2cac4fea 100644
--- a/Userland/Applications/HexEditor/HexEditor.h
+++ b/Userland/Applications/HexEditor/HexEditor.h
@@ -19,6 +19,7 @@
 #include <AK/StdLibExtras.h>
 #include <LibCore/Timer.h>
 #include <LibGUI/AbstractScrollableWidget.h>
+#include <LibGUI/UndoStack.h>
 #include <LibGfx/Font/Font.h>
 #include <LibGfx/TextAlignment.h>
 
@@ -40,6 +41,10 @@ public:
     bool save_as(NonnullRefPtr<Core::File>);
     bool save();
 
+    bool undo();
+    bool redo();
+    GUI::UndoStack& undo_stack();
+
     void select_all();
     bool has_selection() const { return m_selection_start < m_selection_end && m_document->size() > 0; }
     size_t selection_size();
@@ -83,6 +88,7 @@ private:
     NonnullRefPtr<Core::Timer> m_blink_timer;
     bool m_cursor_blink_active { false };
     NonnullOwnPtr<HexDocument> m_document;
+    GUI::UndoStack m_undo_stack;
 
     static constexpr int m_address_bar_width = 90;
     static constexpr int m_padding = 5;
@@ -101,6 +107,8 @@ private:
     void set_content_length(size_t); // I might make this public if I add fetching data on demand.
     void update_status();
     void did_change();
+    void did_complete_action(size_t position, u8 old_value, u8 new_value);
+    void did_complete_action(size_t position, ByteBuffer&& old_values, ByteBuffer&& new_values);
 
     void reset_cursor_blink_state();
 };
diff --git a/Userland/Applications/HexEditor/HexEditorWidget.cpp b/Userland/Applications/HexEditor/HexEditorWidget.cpp
index beb50ae7b5..84963272d8 100644
--- a/Userland/Applications/HexEditor/HexEditorWidget.cpp
+++ b/Userland/Applications/HexEditor/HexEditorWidget.cpp
@@ -83,6 +83,11 @@ HexEditorWidget::HexEditorWidget()
         window()->set_modified(is_document_dirty);
     };
 
+    m_editor->undo_stack().on_state_change = [this] {
+        m_undo_action->set_enabled(m_editor->undo_stack().can_undo());
+        m_redo_action->set_enabled(m_editor->undo_stack().can_redo());
+    };
+
     m_search_results->set_activates_on_selection(true);
     m_search_results->on_activation = [this](const GUI::ModelIndex& index) {
         if (!index.is_valid())
@@ -151,6 +156,16 @@ HexEditorWidget::HexEditorWidget()
         dbgln("Wrote document to {}", file->filename());
     });
 
+    m_undo_action = GUI::CommonActions::make_undo_action([&](auto&) {
+        m_editor->undo();
+    });
+    m_undo_action->set_enabled(false);
+
+    m_redo_action = GUI::CommonActions::make_redo_action([&](auto&) {
+        m_editor->redo();
+    });
+    m_redo_action->set_enabled(false);
+
     m_find_action = GUI::Action::create("&Find", { Mod_Ctrl, Key_F }, Gfx::Bitmap::try_load_from_file("/res/icons/16x16/find.png"sv).release_value_but_fixme_should_propagate_errors(), [&](const GUI::Action&) {
         auto old_buffer = m_search_buffer;
         bool find_all = false;
@@ -243,6 +258,9 @@ HexEditorWidget::HexEditorWidget()
     m_toolbar->add_action(*m_open_action);
     m_toolbar->add_action(*m_save_action);
     m_toolbar->add_separator();
+    m_toolbar->add_action(*m_undo_action);
+    m_toolbar->add_action(*m_redo_action);
+    m_toolbar->add_separator();
     m_toolbar->add_action(*m_find_action);
     m_toolbar->add_action(*m_goto_offset_action);
 
@@ -378,6 +396,9 @@ void HexEditorWidget::initialize_menubar(GUI::Window& window)
     }));
 
     auto& edit_menu = window.add_menu("&Edit");
+    edit_menu.add_action(*m_undo_action);
+    edit_menu.add_action(*m_redo_action);
+    edit_menu.add_separator();
     edit_menu.add_action(GUI::CommonActions::make_select_all_action([this](auto&) {
         m_editor->select_all();
         m_editor->update();
diff --git a/Userland/Applications/HexEditor/HexEditorWidget.h b/Userland/Applications/HexEditor/HexEditorWidget.h
index 12bf8d155f..b59e7435a5 100644
--- a/Userland/Applications/HexEditor/HexEditorWidget.h
+++ b/Userland/Applications/HexEditor/HexEditorWidget.h
@@ -53,6 +53,10 @@ private:
     RefPtr<GUI::Action> m_open_action;
     RefPtr<GUI::Action> m_save_action;
     RefPtr<GUI::Action> m_save_as_action;
+
+    RefPtr<GUI::Action> m_undo_action;
+    RefPtr<GUI::Action> m_redo_action;
+
     RefPtr<GUI::Action> m_find_action;
     RefPtr<GUI::Action> m_goto_offset_action;
     RefPtr<GUI::Action> m_layout_toolbar_action;