From 5afe852f758bc049fdc839c5e48c8fc888a5609d Mon Sep 17 00:00:00 2001
From: NukeBird <nukebird.dev@gmail.com>
Date: Fri, 21 Feb 2025 18:26:24 +0300
Subject: [PATCH] Implement ensure method

---
 tests/zecsy.cpp | 13 +++++++++----
 zecsy.hpp       | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/tests/zecsy.cpp b/tests/zecsy.cpp
index 11f187e..e69c697 100644
--- a/tests/zecsy.cpp
+++ b/tests/zecsy.cpp
@@ -67,13 +67,12 @@ TEST_CASE("Attach a simple component to an entity and verify it is correctly "
     world w;
 
     auto e1 = w.make_entity();
-    w.set(e1, ChoosenOne{});
 
+    w.set(e1, ChoosenOne{});
     REQUIRE(w.has<ChoosenOne>(e1));
 
     auto e2 = w.make_entity();
     w.set(e2, ChoosenOne{});
-
     REQUIRE(w.has<ChoosenOne>(e2));
 }
 
@@ -94,8 +93,15 @@ TEST_CASE("Retrieve a component from an entity and verify its data matches "
     REQUIRE(w.get<Comp>(e).v == 0);
 
     w.get<Comp>(e).v = 77;
-
     REQUIRE(w.get<Comp>(e).v == 77);
+
+    w.ensure<Comp>(e).v = 4;
+    REQUIRE(w.ensure<Comp>(e).v == 4);
+    REQUIRE(w.get<Comp>(e).v == 4);
+
+    w.remove<Comp>(e);
+    w.ensure<Comp>(e).v = 123;
+    REQUIRE(w.get<Comp>(e).v == 123);
 }
 
 TEST_CASE(
@@ -112,7 +118,6 @@ TEST_CASE(
 
     w.set(e, ChoosenOne{});
     REQUIRE_NOTHROW(w.remove<ChoosenOne>(e));
-
     REQUIRE_FALSE(w.has<ChoosenOne>(e));
 }
 
diff --git a/zecsy.hpp b/zecsy.hpp
index 7c15b13..a6362ae 100644
--- a/zecsy.hpp
+++ b/zecsy.hpp
@@ -70,6 +70,9 @@ namespace zecsy
         void set(entity_id e, const First& comp0, const Second& comp1,
                  const Rest&... rest_comps);
 
+        template<Component T>
+        T& ensure(entity_id e);
+
         template<Component T>
         void remove(entity_id e);
 
@@ -219,6 +222,17 @@ namespace zecsy
         return false;
     }
 
+    template<Component T>
+    T& world::ensure(entity_id e)
+    {
+        if(!has<T>(e))
+        {
+            set<T>(e);
+        }
+
+        return get<T>(e);
+    }
+
     template<Component T>
     inline T& world::get(entity_id e)
     {