From 2a77e40d12887097f8ec811f88f7ffbca02eaa02 Mon Sep 17 00:00:00 2001
From: griffi-gh <prasol258@gmail.com>
Date: Wed, 16 Oct 2024 10:57:19 +0200
Subject: [PATCH] keep selection on restore

---
 src/src/gundo.cc | 29 ++++++++++++++++++++++++++---
 src/src/gundo.hh | 10 +++++++++-
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/src/gundo.cc b/src/src/gundo.cc
index 8a2e002b..2f237d87 100644
--- a/src/src/gundo.cc
+++ b/src/src/gundo.cc
@@ -3,6 +3,7 @@
 #include "game.hh"
 #include "tms/backend/print.h"
 #include "world.hh"
+#include <cstdint>
 #include <cstdlib>
 
 #ifdef DEBUG
@@ -75,14 +76,25 @@ void undo_stack::checkpoint(const char *reason, void *snapshot /* = nullptr */)
     );
 }
 
-const char* undo_stack::restore(bool keep_cam_pos /* = true */) {
+const char* undo_stack::restore(uint8_t flags) {
+    // Save the current camera position
     float cx, cy, cz;
-    if (keep_cam_pos) {
+    if (flags & UNDO_KEEP_CAM_POS) {
         cx = G->cam->_position.x;
         cy = G->cam->_position.y;
         cz = G->cam->_position.z;
     }
 
+    // Save the current selection
+    uint32_t saved_selection_id = 0;
+    if (flags & UNDO_KEEP_SELECTION) {
+        if (G->selection.e != nullptr) {
+            saved_selection_id = G->selection.e->id;
+        } else {
+            tms_warnf("undo_restore: no selection to keep");
+        }
+    }
+
     if (this->items.size() == 0) {
         tms_fatalf("undo_load: no items to load");
     }
@@ -95,8 +107,19 @@ const char* undo_stack::restore(bool keep_cam_pos /* = true */) {
 
     free(item.data);
 
-    if (keep_cam_pos) {
+    if (flags & UNDO_KEEP_CAM_POS) {
         G->cam->set_position(cx, cy, cz);
     }
+    if (flags & UNDO_KEEP_SELECTION) {
+        G->selection.disable();
+        if (saved_selection_id) {
+            if (entity *e = W->get_entity_by_id(saved_selection_id)) {
+                G->selection.select(e);
+            } else {
+                tms_warnf("undo_restore: selection entity not found");
+            }
+        }
+    }
+
     return item.reason;
 }
diff --git a/src/src/gundo.hh b/src/src/gundo.hh
index 03741181..ebf08748 100644
--- a/src/src/gundo.hh
+++ b/src/src/gundo.hh
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <cstddef>
+#include <cstdint>
 #include <vector>
 
 #define MAX_UNDO_ITEMS 100
@@ -11,6 +12,13 @@ struct undo_item {
     size_t size;
 };
 
+enum {
+    UNDO_KEEP_NONE      = 0b00,
+    UNDO_KEEP_CAM_POS   = 0b01,
+    UNDO_KEEP_SELECTION = 0b10,
+    UNDO_KEEP_ALL       = 0b11,
+};
+
 struct undo_stack {
     private:
         std::vector<undo_item> items;
@@ -39,7 +47,7 @@ struct undo_stack {
         // Returns the reason for the last checkpoint
         //
         // Avoid using this function directly, use P.add_action(ACTION_UNDO_RESTORE, 0) instead
-        const char* restore(bool keep_cam_pos = true);
+        const char* restore(uint8_t flags = UNDO_KEEP_ALL);
 };
 
 extern struct undo_stack undo;