From e0ddac4e75cc9af97aa124e6884ec79bb6325efc Mon Sep 17 00:00:00 2001
From: TakeV <takev@disroot.org>
Date: Mon, 27 May 2024 00:59:47 -0400
Subject: [PATCH] Handle issue where entity-map:update-entity! would try
 setting an invalid entity

---
 modules/ces/entity/entity-map.scm | 11 +++++++----
 test/tests/entity-map-tests.scm   | 28 +++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/modules/ces/entity/entity-map.scm b/modules/ces/entity/entity-map.scm
index 8498573..e89ced0 100644
--- a/modules/ces/entity/entity-map.scm
+++ b/modules/ces/entity/entity-map.scm
@@ -115,8 +115,9 @@ stored."
 
 (define (delete-entity! entity-map entity)
   "Removes the entity from the entity map"
-  (when (and (entity-map? entity-map)
-             (get-entity entity-map entity))
+  (when (and
+         (entity-map? entity-map)
+         (get-entity entity-map entity))
     (hashtable-delete! (entity-map-entities entity-map)
                        (entity-key entity))))
 
@@ -126,9 +127,11 @@ with the result of proc.
 proc must be a proceedure with one argument: the entity to update.
 proc should be free of side effects, as it may be called multiple times.
 If the entity does not exist, no change is applied. "
-  (let ((old-entity (get-entity entity-map entity-id)))
+  (let ((old-entity (get-entity entity-map entity-id #:default #f)))
     (when old-entity
-      (set-entity! entity-map entity-id (proc old-entity)))))
+      (let ((updated-entity (proc old-entity)))
+        (when (game-entity? updated-entity)
+          (set-entity! entity-map entity-id updated-entity))))))
 
 (define (update-entity entity-map entity-id proc)
   "Pure update-entity call, does not modify the underlying entity.
diff --git a/test/tests/entity-map-tests.scm b/test/tests/entity-map-tests.scm
index e221ab8..fdbe63c 100644
--- a/test/tests/entity-map-tests.scm
+++ b/test/tests/entity-map-tests.scm
@@ -87,8 +87,34 @@
   (let* ((funky-entity (entity-generator))
          (funky-entity-id (game-entity-id funky-entity))
          (correct-id (+ 100 funky-entity-id)))
+
     (set-entity! emap-copy correct-id funky-entity)
-    (test-eqv correct-id (game-entity-id (get-entity emap-copy correct-id)))))
+    (test-eqv correct-id (game-entity-id (get-entity emap-copy correct-id)))
+
+    ;; Test deletion
+    (delete-entity! emap-copy correct-id)
+    (test-eqv 'notfound (get-entity emap-copy correct-id #:default 'notfound))))
+
+(let* ((emap-original list-constructed)
+       (emap-copy (copy-entity-map emap-original))
+       (new-components (create-ihashtable '((a . b))))
+       (add-component-proc (lambda (e)
+                             (set-entity-components e new-components)))
+       (bad-update-proc (lambda (e)
+                          'this-should-fail))
+       (get-component (lambda (e-id c-key)
+                        (ihashtable-get
+                         (get-entity-components
+                          (get-entity emap-copy e-id))
+                         c-key))))
+
+  ;; Components updated
+  (update-entity! emap-copy next-id add-component-proc)
+  (test-eqv 'b (get-component next-id 'a))
+
+  ;; No update on bad input
+  (update-entity! emap-copy next-id bad-update-proc)
+  (test-eqv 'b (get-component next-id 'a)))
 
 (test-end)