hdk-grid/include/hdk/grid/SharedPtrWrapper.hpp

#pragma once
/// @file SharedPtrWrapper.hpp
/// For complete documentation, see src/SharedPtrWrapper.cpp
#include <functional>
#include <memory>
#include <unordered_map>
namespace hdk::grid {
  template <typename StructType> class SharedPtrWrapper {
    public:
      using SharedPtr = std::shared_ptr<StructType>;
      using WeakPtr = std::weak_ptr<StructType>;
      using SharedCacheMap = std::unordered_map<StructType*, WeakPtr>;
      using DeleteFunc = std::function<void(StructType*)>;
      static void null_deleter(StructType*) { }
      SharedPtrWrapper() = delete;
      SharedPtrWrapper(std::nullptr_t)
        : shared_struct(nullptr) { }
      SharedPtrWrapper(StructType* struct_ptr, DeleteFunc deleteFunc)
        : shared_struct(struct_ptr, deleteFunc) {
        // Cache the shared pointer to ensure consistent reference counting for the same raw pointer.
        get_shared_weak_ptr_cache()[struct_ptr] = shared_struct;
      }
      SharedPtrWrapper(SharedPtrWrapper&& other) noexcept
        : shared_struct(std::move(other.shared_struct)) { }
      /** Existing Shared Pointer */
      SharedPtrWrapper(SharedPtr shared_ptr)

      {
        shared_struct = shared_ptr;
        if ((shared_ptr.get() != nullptr) && !is_cached(shared_ptr.get())) {
          get_shared_weak_ptr_cache()[shared_ptr.get()] = shared_struct;
        }
      }
      SharedPtrWrapper& operator=(SharedPtrWrapper&& other) noexcept {
        if (this != &other) {
          shared_struct = std::move(other.shared_struct);
        }
        return *this;
      }
      SharedPtrWrapper(const SharedPtrWrapper& other)
        : shared_struct(other.shared_struct) { }
      SharedPtrWrapper& operator=(const SharedPtrWrapper& other) {
        if (this != &other) {
          shared_struct = other.shared_struct;
        }
        return *this;
      }
      operator StructType*() const { return shared_struct.get(); }
      SharedPtrWrapper& operator=(std::nullptr_t) {
        shared_struct.reset();
        return *this;
      }
      operator bool() const { return shared_struct != nullptr; }
      StructType* operator->() const { return shared_struct.get(); }
    private: // internal use only.
      SharedPtr shared_struct { nullptr };

    protected: // For use by derived classes.
      static SharedCacheMap& get_shared_weak_ptr_cache() {
        static SharedCacheMap cache_map;
        return cache_map;
      }
      static bool is_cached(StructType* ptr) {
        auto& cache = get_shared_weak_ptr_cache();
        return cache.find(ptr) != cache.end() && !cache[ptr].expired();
      }
      static SharedPtr get_or_view(StructType* ptr) {
        /** if no pointer */
        if (!ptr) {
          /** then return a non-owning null */
          return SharedPtr(nullptr, null_deleter);
        }

        if (is_cached(ptr)) {
          return get_shared_weak_ptr_cache()[ptr].lock();
        } else {
          SharedPtr new_shared(ptr, null_deleter);
          // cache[ptr] = new_shared;
          return new_shared;
        }
      }
      static SharedPtr get_or_cache(StructType* ptr, DeleteFunc deleteFunc) {
        /** If no ptr, no need to do anything else  */
        if (!ptr) {
          return SharedPtr(nullptr, deleteFunc);
        }
        auto& cache = get_shared_weak_ptr_cache();
        if (is_cached(ptr)) {
          return cache[ptr].lock();
        } else {
          SharedPtr new_shared(ptr, deleteFunc);
          cache[ptr] = new_shared;
          return new_shared;
        }
      }
  };
}