From fd37550b633d5081d9563f0f7da011515e8d5916 Mon Sep 17 00:00:00 2001
From: BlackMATov <blackmatov@gmail.com>
Date: Sat, 31 Aug 2019 09:45:55 +0700
Subject: [PATCH] functions for component ensuring

---
 headers/ecs.hpp/ecs.hpp | 43 +++++++++++++++++++++++++++++++++++++++++
 untests/ecs_tests.cpp   | 25 ++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/headers/ecs.hpp/ecs.hpp b/headers/ecs.hpp/ecs.hpp
index 3d0c728..eee23a1 100644
--- a/headers/ecs.hpp/ecs.hpp
+++ b/headers/ecs.hpp/ecs.hpp
@@ -669,6 +669,12 @@ namespace ecs_hpp
                 return components_.get(id);
             }
 
+            template < typename... Args >
+            T& ensure(entity_id id, Args&&... args) {
+                components_.emplace(id, std::forward<Args>(args)...);
+                return components_.get(id);
+            }
+
             bool exists(entity_id id) const noexcept {
                 return components_.has(id);
             }
@@ -740,6 +746,12 @@ namespace ecs_hpp
                 return empty_value_;
             }
 
+            template < typename... Args >
+            T& ensure(entity_id id, Args&&...) {
+                components_.insert(id);
+                return empty_value_;
+            }
+
             bool exists(entity_id id) const noexcept {
                 return components_.has(id);
             }
@@ -839,6 +851,9 @@ namespace ecs_hpp
         template < typename T, typename... Args >
         T& assign_component(Args&&... args);
 
+        template < typename T, typename... Args >
+        T& ensure_component(Args&&... args);
+
         template < typename T >
         bool remove_component();
 
@@ -992,6 +1007,9 @@ namespace ecs_hpp
         template < typename... Args >
         T& assign(Args&&... args);
 
+        template < typename... Args >
+        T& ensure(Args&&... args);
+
         T& get();
         const T& get() const;
 
@@ -1262,6 +1280,9 @@ namespace ecs_hpp
         template < typename T, typename... Args >
         T& assign_component(const uentity& ent, Args&&... args);
 
+        template < typename T, typename... Args >
+        T& ensure_component(const uentity& ent, Args&&... args);
+
         template < typename T >
         bool remove_component(const uentity& ent);
 
@@ -1490,6 +1511,13 @@ namespace ecs_hpp
             std::forward<Args>(args)...);
     }
 
+    template < typename T, typename... Args >
+    T& entity::ensure_component(Args&&... args) {
+        return (*owner_).ensure_component<T>(
+            id_,
+            std::forward<Args>(args)...);
+    }
+
     template < typename T >
     bool entity::remove_component() {
         return (*owner_).remove_component<T>(id_);
@@ -1695,6 +1723,13 @@ namespace ecs_hpp
             std::forward<Args>(args)...);
     }
 
+    template < typename T >
+    template < typename... Args >
+    T& component<T>::ensure(Args&&... args) {
+        return owner_.ensure_component<T>(
+            std::forward<Args>(args)...);
+    }
+
     template < typename T >
     T& component<T>::get() {
         return owner_.get_component<T>();
@@ -2174,6 +2209,14 @@ namespace ecs_hpp
             std::forward<Args>(args)...);
     }
 
+    template < typename T, typename... Args >
+    T& registry::ensure_component(const uentity& ent, Args&&... args) {
+        assert(valid_entity(ent));
+        return get_or_create_storage_<T>().ensure(
+            ent,
+            std::forward<Args>(args)...);
+    }
+
     template < typename T >
     bool registry::remove_component(const uentity& ent) {
         assert(valid_entity(ent));
diff --git a/untests/ecs_tests.cpp b/untests/ecs_tests.cpp
index 4f9e68b..f30866f 100644
--- a/untests/ecs_tests.cpp
+++ b/untests/ecs_tests.cpp
@@ -813,6 +813,31 @@ TEST_CASE("registry") {
             e1.destroy();
         }
     }
+    SECTION("component_ensuring") {
+        {
+            ecs::registry w;
+            ecs::entity e1 = w.create_entity();
+
+            e1.ensure_component<position_c>(1,2);
+            e1.ensure_component<movable_c>();
+
+            REQUIRE(e1.get_component<position_c>() == position_c(1,2));
+            REQUIRE(e1.exists_component<movable_c>());
+
+            e1.ensure_component<position_c>(10,20).x = 15;
+            e1.ensure_component<movable_c>();
+
+            REQUIRE(e1.get_component<position_c>() == position_c(15,2));
+            REQUIRE(e1.exists_component<movable_c>());
+
+            ecs::component<velocity_c> c1 = w.wrap_component<velocity_c>(e1);
+            REQUIRE_FALSE(c1.exists());
+            c1.ensure(2, 1).y = 10;
+            REQUIRE(c1.get() == velocity_c(2, 10));
+            c1.ensure(10, 20).x = 20;
+            REQUIRE(c1.get() == velocity_c(20, 10));
+        }
+    }
     SECTION("component_accessing") {
         {
             ecs::registry w;