From 0a3be94815675f8a12a5e04824d749c7db8041f0 Mon Sep 17 00:00:00 2001
From: Rochet2 <rochet2@post.com>
Date: Fri, 2 Feb 2018 00:27:14 +0200
Subject: [PATCH] Implement convenience functions for DataMap

See https://github.com/azerothcore/azerothcore-wotlk/pull/748#issuecomment-362411082 for more.
---
 src/common/Utilities/DataMap.h | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/common/Utilities/DataMap.h b/src/common/Utilities/DataMap.h
index 75e8b7555..79edfa6ec 100644
--- a/src/common/Utilities/DataMap.h
+++ b/src/common/Utilities/DataMap.h
@@ -12,20 +12,50 @@
 class DataMap
 {
 public:
+    /**
+     * Base class that you should inherit in your script.
+     * Inheriting classes can be stored to DataMap
+     */
     class Base
     {
     public:
         virtual ~Base() = default;
     };
 
-    template<class T> T* GetCustomData(std::string const & k) const {
+    /**
+     * Returns a pointer to object of requested type stored with given key or nullptr
+     */
+    template<class T> T* Get(std::string const & k) const {
         static_assert(std::is_base_of<Base, T>::value, "T must derive from Base");
         auto it = Container.find(k);
         if (it != Container.end())
             return dynamic_cast<T*>(it->second.get());
         return nullptr;
     }
-    void SetCustomData(std::string const & k, Base* v) { Container[k] = std::unique_ptr<Base>(v); }
+    
+    /**
+     * Returns a pointer to object of requested type stored with given key
+     * or default constructs one and returns that one
+     */
+    template<class T, typename std::enable_if<std::is_default_constructible<T>{}, int>::type = 0>
+    T* GetDefault(std::string const & k) {
+        static_assert(std::is_base_of<Base, T>::value, "T must derive from Base");
+        if (T* v = Get<T>(k))
+            return v;
+        T* v = new T();
+        cont.emplace(k, std::unique_ptr<T>(v));
+        return v;
+    }
+    
+    /**
+     * Stores a new object that inherits the Base class with the given key
+     */
+    void Set(std::string const & k, Base* v) { Container[k] = std::unique_ptr<Base>(v); }
+
+    /**
+     * Removes objects with given key and returns true if one was removed, false otherwise
+     */
+    bool Erase(std::string const & k) { cont.erase(k) != 0; }
 
 private:
     std::unordered_map<std::string, std::unique_ptr<Base>> Container;