//========= Copyright Valve Corporation, All rights reserved. ============//
//
// Purpose: Tools for correctly implementing & handling reference counted
//			objects
//
//=============================================================================

#ifndef REFCOUNT_H
#define REFCOUNT_H

#include "tier0/threadtools.h"

#if defined( _WIN32 )
#pragma once
#endif

template <typename T>
inline void SafeAssign(T** ppInoutDst, T* pInoutSrc )
{
	Assert( ppInoutDst );

	// Do addref before release
	if ( pInoutSrc )
		( pInoutSrc )->AddRef();

	// Do addref before release
	if ( *ppInoutDst )
		( *ppInoutDst )->Release();

	// Do the assignment
	( *ppInoutDst ) = pInoutSrc;
}

template <typename T>
inline void SafeAddRef( T* pObj )
{
	if ( pObj )
		pObj->AddRef();
}

template <typename T>
inline void SafeRelease( T** ppInoutPtr )
{
	Assert( ppInoutPtr  );
	if ( *ppInoutPtr )
		( *ppInoutPtr )->Release();

	( *ppInoutPtr ) = NULL;
}

//-----------------------------------------------------------------------------
// Purpose:	Implement a standard reference counted interface. Use of this
//			is optional insofar as all the concrete tools only require
//			at compile time that the function signatures match.
//-----------------------------------------------------------------------------

class IRefCounted
{
public:
	virtual int AddRef() = 0;
	virtual int Release() = 0;
};


//-----------------------------------------------------------------------------
// Purpose:	Release a pointer and mark it NULL
//-----------------------------------------------------------------------------

template <class REFCOUNTED_ITEM_PTR>
inline int SafeRelease( REFCOUNTED_ITEM_PTR &pRef )
{
	// Use funny syntax so that this works on "auto pointers"
	REFCOUNTED_ITEM_PTR *ppRef = &pRef;
	if ( *ppRef )
	{
		int result = (*ppRef)->Release();
		*ppRef = NULL;
		return result;
	}
	return 0;
}

//-----------------------------------------------------------------------------
// Purpose:	Maintain a reference across a scope
//-----------------------------------------------------------------------------

template <class T = IRefCounted>
class CAutoRef
{
public:
	CAutoRef( T *pRef )
	  : m_pRef( pRef )
	{
		if ( m_pRef )
			m_pRef->AddRef();
	}

   ~CAutoRef()
   {
      if (m_pRef)
         m_pRef->Release();
   }

private:
   T *m_pRef;
};

//-----------------------------------------------------------------------------
// Purpose:	Do a an inline AddRef then return the pointer, useful when
//			returning an object from a function
//-----------------------------------------------------------------------------

#define RetAddRef( p ) ( (p)->AddRef(), (p) )
#define InlineAddRef( p ) ( (p)->AddRef(), (p) )


//-----------------------------------------------------------------------------
// Purpose:	A class to both hold a pointer to an object and its reference.
//			Base exists to support other cleanup models
//-----------------------------------------------------------------------------

template <class T>
class CBaseAutoPtr
{
public:
	CBaseAutoPtr()                                         	: m_pObject(0) {}
	CBaseAutoPtr(T *pFrom)                        			: m_pObject(pFrom) {}

	operator const void *() const          					{ return m_pObject; }
	operator void *()                      					{ return m_pObject; }

	operator const T *() const							    { return m_pObject; }
	operator const T *()          							{ return m_pObject; }
	operator T *()											{ return m_pObject; }

	int			operator=( int i )							{ AssertMsg( i == 0, "Only NULL allowed on integer assign" ); m_pObject = 0; return 0; }
	T *			operator=( T *p )							{ m_pObject = p; return p; }

    bool        operator !() const							{ return ( !m_pObject ); }
    bool        operator!=( int i ) const					{ AssertMsg( i == 0, "Only NULL allowed on integer compare" ); return (m_pObject != NULL); }
	bool		operator==( const void *p ) const			{ return ( m_pObject == p ); }
	bool		operator!=( const void *p ) const			{ return ( m_pObject != p ); }
	bool		operator==( T *p ) const					{ return operator==( (void *)p ); }
	bool		operator!=( T *p ) const					{ return operator!=( (void *)p ); }
	bool		operator==( const CBaseAutoPtr<T> &p ) const { return operator==( (const void *)p ); }
	bool		operator!=( const CBaseAutoPtr<T> &p ) const { return operator!=( (const void *)p ); }

	T *  		operator->()								{ return m_pObject; }
	T &  		operator *()								{ return *m_pObject; }
	T ** 		operator &()								{ return &m_pObject; }

	const T *   operator->() const							{ return m_pObject; }
	const T &   operator *() const							{ return *m_pObject; }
	T * const * operator &() const							{ return &m_pObject; }

protected:
	CBaseAutoPtr( const CBaseAutoPtr<T> &from )				: m_pObject( from.m_pObject ) {}
	void operator=( const CBaseAutoPtr<T> &from ) 			{ m_pObject = from.m_pObject; }

	T *m_pObject;
};

//---------------------------------------------------------

template <class T>
class CRefPtr : public CBaseAutoPtr<T>
{
	typedef CBaseAutoPtr<T> BaseClass;
public:
	CRefPtr()												{}
	CRefPtr( T *pInit )										: BaseClass( pInit ) {}
	CRefPtr( const CRefPtr<T> &from )						: BaseClass( from ) {}
	~CRefPtr()												{ if ( BaseClass::m_pObject ) BaseClass::m_pObject->Release(); }

	void operator=( const CRefPtr<T> &from )				{ BaseClass::operator=( from ); }

	int operator=( int i )									{ return BaseClass::operator=( i ); }
	T *operator=( T *p )									{ return BaseClass::operator=( p ); }

	operator bool() const									{ return !BaseClass::operator!(); }
	operator bool()											{ return !BaseClass::operator!(); }

	void SafeRelease()										{ if ( BaseClass::m_pObject ) BaseClass::m_pObject->Release(); BaseClass::m_pObject = 0; }
	void AssignAddRef( T *pFrom )							{ SafeRelease(); if (pFrom) pFrom->AddRef(); BaseClass::m_pObject = pFrom; }
	void AddRefAssignTo( T *&pTo )							{ ::SafeRelease( pTo ); if ( BaseClass::m_pObject ) BaseClass::m_pObject->AddRef(); pTo = BaseClass::m_pObject; }
};


//-----------------------------------------------------------------------------
// Purpose:	Traits classes defining reference count threading model
//-----------------------------------------------------------------------------

class CRefMT
{
public:
	static int Increment( int *p) { return ThreadInterlockedIncrement( (long *)p ); }
	static int Decrement( int *p) { return ThreadInterlockedDecrement( (long *)p ); }
};

class CRefST
{
public:
	static int Increment( int *p) { return ++(*p); }
	static int Decrement( int *p) { return --(*p); }
};

//-----------------------------------------------------------------------------
// Purpose:	Actual reference counting implementation. Pulled out to reduce
//			code bloat.
//-----------------------------------------------------------------------------

template <const bool bSelfDelete, typename CRefThreading = CRefMT>
class NO_VTABLE CRefCountServiceBase
{
protected:
	CRefCountServiceBase()
	  : m_iRefs( 1 )
	{
	}

	virtual ~CRefCountServiceBase()
	{
	}

	virtual bool OnFinalRelease()
	{
		return true;
	}

	int GetRefCount() const
	{
		return m_iRefs;
	}

	int DoAddRef()
	{
		return CRefThreading::Increment( &m_iRefs );
	}

	int DoRelease()
	{
		int result = CRefThreading::Decrement( &m_iRefs );
		if ( result )
			return result;
		if ( OnFinalRelease() && bSelfDelete )
			delete this;
		return 0;
	}

private:
	int m_iRefs;
};

class CRefCountServiceNull
{
protected:
	static int DoAddRef() { return 1; }
	static int DoRelease() { return 1; }
};

template <typename CRefThreading = CRefMT>
class NO_VTABLE CRefCountServiceDestruct
{
protected:
	CRefCountServiceDestruct()
		: m_iRefs( 1 )
	{
	}

	virtual ~CRefCountServiceDestruct()
	{
	}

	int GetRefCount() const
	{
		return m_iRefs;
	}

	int DoAddRef()
	{
		return CRefThreading::Increment( &m_iRefs );
	}

	int DoRelease()
	{
		int result = CRefThreading::Decrement( &m_iRefs );
		if ( result )
			return result;
		this->~CRefCountServiceDestruct();
		return 0;
	}

private:
	int m_iRefs;
};


typedef CRefCountServiceBase<true, CRefST>	CRefCountServiceST;
typedef CRefCountServiceBase<false, CRefST>	CRefCountServiceNoDeleteST;

typedef CRefCountServiceBase<true, CRefMT>	CRefCountServiceMT;
typedef CRefCountServiceBase<false, CRefMT> CRefCountServiceNoDeleteMT;

// Default to threadsafe
typedef CRefCountServiceNoDeleteMT			CRefCountServiceNoDelete;
typedef CRefCountServiceMT					CRefCountService;

//-----------------------------------------------------------------------------
// Purpose:	Base classes to implement reference counting
//-----------------------------------------------------------------------------

template < class REFCOUNT_SERVICE = CRefCountService > 
class NO_VTABLE CRefCounted : public REFCOUNT_SERVICE
{
public:
	virtual ~CRefCounted()	{}
	int AddRef() 			{ return REFCOUNT_SERVICE::DoAddRef(); }
	int Release()			{ return REFCOUNT_SERVICE::DoRelease(); }
};

//-------------------------------------

template < class BASE1, class REFCOUNT_SERVICE = CRefCountService > 
class NO_VTABLE CRefCounted1 : public BASE1,
							   public REFCOUNT_SERVICE
{
public:
	virtual ~CRefCounted1()	{}
	int AddRef() 			{ return REFCOUNT_SERVICE::DoAddRef(); }
	int Release()			{ return REFCOUNT_SERVICE::DoRelease(); }
};

//-------------------------------------

template < class BASE1, class BASE2, class REFCOUNT_SERVICE = CRefCountService > 
class NO_VTABLE CRefCounted2 : public BASE1, public BASE2,
							   public REFCOUNT_SERVICE
{
public:
	virtual ~CRefCounted2()	{}
	int AddRef() 			{ return REFCOUNT_SERVICE::DoAddRef(); }
	int Release()			{ return REFCOUNT_SERVICE::DoRelease(); }
};

//-------------------------------------

template < class BASE1, class BASE2, class BASE3, class REFCOUNT_SERVICE = CRefCountService > 
class NO_VTABLE CRefCounted3 : public BASE1, public BASE2, public BASE3,
							   public REFCOUNT_SERVICE
{
	virtual ~CRefCounted3()	{}
	int AddRef() 			{ return REFCOUNT_SERVICE::DoAddRef(); }
	int Release()			{ return REFCOUNT_SERVICE::DoRelease(); }
};

//-------------------------------------

template < class BASE1, class BASE2, class BASE3, class BASE4, class REFCOUNT_SERVICE = CRefCountService > 
class NO_VTABLE CRefCounted4 : public BASE1, public BASE2, public BASE3, public BASE4,
							   public REFCOUNT_SERVICE
{
public:
	virtual ~CRefCounted4()	{}
	int AddRef() 			{ return REFCOUNT_SERVICE::DoAddRef(); }
	int Release()			{ return REFCOUNT_SERVICE::DoRelease(); }
};

//-------------------------------------

template < class BASE1, class BASE2, class BASE3, class BASE4, class BASE5, class REFCOUNT_SERVICE = CRefCountService > 
class NO_VTABLE CRefCounted5 : public BASE1, public BASE2, public BASE3, public BASE4, public BASE5,
							   public REFCOUNT_SERVICE
{
public:
	virtual ~CRefCounted5()	{}
	int AddRef() 			{ return REFCOUNT_SERVICE::DoAddRef(); }
	int Release()			{ return REFCOUNT_SERVICE::DoRelease(); }
};

//-----------------------------------------------------------------------------
// Purpose:	Class to throw around a reference counted item to debug
//			referencing problems
//-----------------------------------------------------------------------------

template <class BASE_REFCOUNTED, int FINAL_REFS, const char *pszName>
class CRefDebug : public BASE_REFCOUNTED
{
public:
#ifdef _DEBUG
	CRefDebug()
	{
		AssertMsg( this->GetRefCount() == 1, "Expected initial ref count of 1" );
		DevMsg( "%s:create 0x%x\n", ( pszName ) ? pszName : "", this );
	}

	virtual ~CRefDebug()
	{
		AssertDevMsg( this->GetRefCount() == FINAL_REFS, "Object still referenced on destroy?" );
		DevMsg( "%s:destroy 0x%x\n", ( pszName ) ? pszName : "", this );
	}

	int AddRef()
	{
		DevMsg( "%s:(0x%x)->AddRef() --> %d\n", ( pszName ) ? pszName : "", this, this->GetRefCount() + 1 );
		return BASE_REFCOUNTED::AddRef();
	}

	int Release()
	{
		DevMsg( "%s:(0x%x)->Release() --> %d\n", ( pszName ) ? pszName : "", this, this->GetRefCount() - 1 );
		Assert( this->GetRefCount() > 0 );
		return BASE_REFCOUNTED::Release();
	}
#endif
};

//-----------------------------------------------------------------------------

#endif // REFCOUNT_H