//===------------------SharedCluster.h --------------------------*- C++ -*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#ifndef utility_SharedCluster_h_
#define utility_SharedCluster_h_

#include "lldb/Utility/LLDBAssert.h"
#include "lldb/Utility/SharingPtr.h"

#include "llvm/ADT/SmallPtrSet.h"

namespace lldb_private {

namespace imp
{
    template <typename T>
    class shared_ptr_refcount : public lldb_private::imp::shared_count
    {
    public:
        template<class Y> shared_ptr_refcount (Y *in) : shared_count (0), manager(in) {}
        
        shared_ptr_refcount() : shared_count (0) {}
        
        ~shared_ptr_refcount() override
        {
        }
        
        void on_zero_shared() override
        {
            manager->DecrementRefCount();
        }

    private:
        T *manager;
    };

} // namespace imp

template <class T>
class ClusterManager
{
public:
    ClusterManager() : m_objects(), m_external_ref(0), m_mutex() {}

    ~ClusterManager()
    {
        for (typename llvm::SmallPtrSet<T *, 16>::iterator pos = m_objects.begin(), end = m_objects.end(); pos != end;
             ++pos)
        {
            T *object = *pos;
            delete object;
        }

        // Decrement refcount should have been called on this ClusterManager,
        // and it should have locked the mutex, now we will unlock it before
        // we destroy it...
        m_mutex.unlock();
    }

    void
    ManageObject(T *new_object)
    {
        std::lock_guard<std::mutex> guard(m_mutex);
        m_objects.insert(new_object);
    }

    typename lldb_private::SharingPtr<T>
    GetSharedPointer(T *desired_object)
    {
        {
            std::lock_guard<std::mutex> guard(m_mutex);
            m_external_ref++;
            if (0 == m_objects.count(desired_object))
            {
                lldbassert(false && "object not found in shared cluster when expected");
                desired_object = nullptr;
            }
        }
        return typename lldb_private::SharingPtr<T>(desired_object, new imp::shared_ptr_refcount<ClusterManager>(this));
    }

private:
    void
    DecrementRefCount()
    {
        m_mutex.lock();
        m_external_ref--;
        if (m_external_ref == 0)
            delete this;
        else
            m_mutex.unlock();
    }

    friend class imp::shared_ptr_refcount<ClusterManager>;

    llvm::SmallPtrSet<T *, 16> m_objects;
    int m_external_ref;
    std::mutex m_mutex;
};

} // namespace lldb_private

#endif // utility_SharedCluster_h_
