memory.hpp

/*
    SPDX-License-Identifier: MIT
    Copyright © 2023-2024 Amebis
*/
 
#pragma once
 
#include "compat.hpp"
#include <memory>
 
#if defined(__GNUC__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunknown-pragmas"
#endif
 
namespace stdex
{
    template <class T>
    struct no_delete {
        constexpr no_delete() noexcept = default;
 
        template <class T2, std::enable_if_t<std::is_convertible_v<T2*, T*>, int> = 0>
        no_delete(const no_delete<T2>&) noexcept {}
 
        void operator()(T* p) const noexcept { _Unreferenced_(p); }
    };
 
    template <class T>
    struct no_delete<T[]> {
        constexpr no_delete() noexcept = default;
 
        template <class T2, std::enable_if_t<std::is_convertible_v<T2(*)[], T(*)[]>, int> = 0>
        no_delete(const no_delete<T2[]>&) noexcept {}
 
        template <class T2, std::enable_if_t<std::is_convertible_v<T2(*)[], T(*)[]>, int> = 0>
        void operator()(T2* p) const noexcept { p; }
    };
 
    template <class T>
    std::shared_ptr<T> make_shared_no_delete(_In_ T* p)
    {
        return std::shared_ptr<T>(p, no_delete<T>{});
    }
 
    // sanitizing_allocator::destroy() member generates p parameter not used warning for primitive datatypes T.
    #pragma warning(push)
    #pragma warning(disable: 4100)
 
    template <class T>
    class sanitizing_allocator : public std::allocator<T>
    {
    public:
        template <class T2>
        struct rebind
        {
            typedef sanitizing_allocator<T2> other; 
        };
 
        sanitizing_allocator() noexcept : std::allocator<T>()
        {}
 
        sanitizing_allocator(_In_ const sanitizing_allocator<T> &other) : std::allocator<T>(other)
        {}
 
        template <class T2>
        sanitizing_allocator(_In_ const sanitizing_allocator<T2> &other) noexcept : std::allocator<T>(other)
        {}
 
        void deallocate(_In_ T* const p, _In_ const std::size_t n)
        {
#ifdef _WIN32
            SecureZeroMemory(p, sizeof(T) * n);
#else
            memset(p, 0, sizeof(T) * n);
#endif
            std::allocator<T>::deallocate(p, n);
        }
    };
 
    #pragma warning(pop)
 
    template <size_t N>
    class sanitizing_blob
    {
    public:
        sanitizing_blob()
        {
            memset(m_data, 0, N);
        }
 
        ~sanitizing_blob()
        {
#ifdef _WIN32
            SecureZeroMemory(m_data, N);
#else
            memset(m_data, 0, N);
#endif
        }
 
    public:
        unsigned char m_data[N]; 
    };
 
    template <typename T, typename D>
    class ref_unique_ptr
    {
    public:
        ref_unique_ptr(_Inout_ std::unique_ptr<T, D> &owner) :
            m_own(owner),
            m_ptr(owner.release())
        {}
 
        ref_unique_ptr(_Inout_ ref_unique_ptr<T, D> &&other) :
            m_own(other.m_own),
            m_ptr(other.m_ptr)
        {
            other.m_ptr = nullptr;
        }
 
        ~ref_unique_ptr()
        {
            if (m_ptr != nullptr)
                m_own.reset(m_ptr);
        }
 
        operator T**()
        {
            return &m_ptr;
        }
 
        operator T*&()
        {
            return m_ptr;
        }
 
    protected:
        std::unique_ptr<T, D> &m_own; 
        T *m_ptr;                     
    };
 
    template<class T, class D>
    ref_unique_ptr<T, D> get_ptr(_Inout_ std::unique_ptr<T, D> &owner) noexcept
    {
        return ref_unique_ptr<T, D>(owner);
    }
 
    template<typename T, typename D>
    class ref_unique_ptr<T[], D>
    {
    public:
        ref_unique_ptr(_Inout_ std::unique_ptr<T[], D> &owner) noexcept :
            m_own(owner),
            m_ptr(owner.release())
        {}
 
        ref_unique_ptr(_Inout_ ref_unique_ptr<T[], D> &&other) :
            m_own(other.m_own),
            m_ptr(other.m_ptr)
        {
            other.m_ptr = nullptr;
        }
 
        virtual ~ref_unique_ptr()
        {
            if (m_ptr != nullptr)
                m_own.reset(m_ptr);
        }
 
        operator T**() noexcept
        {
            return &m_ptr;
        }
 
        operator T*&()
        {
            return m_ptr;
        }
 
    protected:
        std::unique_ptr<T[], D> &m_own; 
        T *m_ptr;                       
    };
 
    template<class T, class D>
    ref_unique_ptr<T[], D> get_ptr(_Inout_ std::unique_ptr<T[], D>& owner) noexcept
    {
        return ref_unique_ptr<T[], D>(owner);
    }
}
 
#if defined(__GNUC__)
#pragma GCC diagnostic pop
#endif