/******************************************************************************/
/* ALLOCATORS *****************************************************************/
/******************************************************************************/

#ifndef _H_EBCL_ALLOC
#define _H_EBCL_ALLOC
#include <lw/lib/Config.hh>
#include <lw/lib/Utilities.hh>
#include <atomic>
namespace lw {


/* Helpers for pool allocators. */
struct T_PoolHelper
{
	/* Storage type for an object */
	template<
		size_t ObjectSize ,
		size_t ObjectAlign
	> using T_Object = std::aligned_storage_t< ObjectSize , ObjectAlign >;

	/* List of object storage items. The storage type is assumed to
	 * have a void* field called list.
	 */
	template<
		typename T_Storage ,
		size_t ListSize
	> struct T_List final
	{
		static constexpr size_t BitmapSize
			= ( ListSize / 8 ) + ( ListSize % 8 ? 1 : 0 );

		using T_Self = T_List< T_Storage , ListSize >;

		T_Storage storage[ ListSize ];	// Stored data
		T_Self* next;			// Next list in chain
		size_t free;			// Free objects
		uint8_t bitmap[ BitmapSize ];	// Allocation map

		T_List( ) = delete;
		T_List( T_Self const& ) = delete;
		T_List( T_Self&& ) = delete;

		// Create a new list, replacing the chain's head
		explicit T_List( T_Self*& head ) noexcept;

		// Find an unused item in the list. The list must have free
		// items.
		T_Storage* findUnused( ) noexcept;
		// Find the index of an item. The item must be part of the list.
		size_t indexOf( T_Storage const* item ) noexcept;

		// Destroy a complete chain of lists
		static void destroy( T_Self*& head ) noexcept;
	};

	/* The core of a pool allocator */
	template<
		typename T_Storage ,
		size_t PerList ,
		size_t MaxFreeLists
	> struct T_Allocator final
	{
		static_assert( PerList > 1 , "allocator lists must contain at least 2 items" );
		static_assert( MaxFreeLists > 0 , "allocator must have at least 1 free list" );

		using T_List_ = T_List< T_Storage , PerList >;
		using T_Self = T_Allocator< T_Storage , PerList , MaxFreeLists >;

		// The chains
		T_List_* free_;
		T_List_* partial_;
		T_List_* full_;
		size_t nFree_;

		T_Allocator( ) noexcept;
		~T_Allocator( );

		T_Allocator( T_Self const& ) = delete;
		T_Allocator( T_Self&& ) = delete;

		/* Allocate an object. */
		T_Storage* allocate( ) noexcept;
		/* Free an object */
		void free( T_Storage* item ) noexcept;

		/* Get the amount of lists. Mostly used for testing. */
		size_t countFreeLists( ) const noexcept;
		size_t countPartialLists( ) const noexcept;
		size_t countFullLists( ) const noexcept;

		/* Get the amount of total, free and used objects */
		void getUsage( size_t& total ,
				size_t& free ,
				size_t& used ) const noexcept;

	    private:
		/* Move a list of objects from one of the allocation lists to another. */
		static void moveList(
				T_List_* list ,
				T_List_*& from ,
				T_List_*& to ) noexcept;

		/* Count the amount of lists in an allocation list */
		static size_t countLists(
				T_List_ const* head ) noexcept;
	};
};

/*============================================================================*/

/* The pool allocator implements a thread-unsafe pool of objects, similar to a
 * slab allocator.
 */
template<
	size_t ObjectSize , size_t ObjectAlign ,
	size_t PerList = std::max( size_t( 8 ) , 4096u / ObjectSize ) ,
	size_t MaxFreeLists = 4
>
class T_PoolAllocator
{
    public:
	using T_Self = T_PoolAllocator< ObjectSize , ObjectAlign , PerList , MaxFreeLists >;

    private:
	// Structure that stores a single allocated item
	struct T_Storage_ final
	{
		T_PoolHelper::T_Object< ObjectSize , ObjectAlign > data;	
		void* list;
	};

	using T_Alloc_ = T_PoolHelper::T_Allocator< T_Storage_ , PerList , MaxFreeLists >;
	T_Alloc_ alloc_;
	
    public:
	T_PoolAllocator( ) noexcept = default;
	T_PoolAllocator( T_Self const& ) = delete;
	T_PoolAllocator( T_Self&& ) = delete;

	/* Allocate an object. */
	void* allocate( size_t requested ) noexcept;
	/* Free an object */
	void free( void* item ) noexcept;

	/* Get the amount of lists. Mostly used for testing. */
	size_t countFreeLists( ) const noexcept;
	size_t countPartialLists( ) const noexcept;
	size_t countFullLists( ) const noexcept;

	/* Get the amount of total, free and used objects */
	void getUsage( size_t& total ,
			size_t& free ,
			size_t& used ) const noexcept;
};

/*============================================================================*/

/* Pool allocator with multithreading support. Allocation is performed from
 * a thread-specific pool allocator; when the memory is freed, the memory is
 * either returned directly to the pool if the freeing thread is the same as
 * the allocating thread, or queued for the original thread to free later.
 */
template<
	size_t ObjectSize , size_t ObjectAlign ,
	size_t PerList = std::max( size_t( 8 ) , 4096u / ObjectSize ) ,
	size_t MaxFreeLists = 4
>
class T_ThreadedPoolAllocator
{
    public:
	using T_Self = T_ThreadedPoolAllocator<
		ObjectSize , ObjectAlign , PerList , MaxFreeLists >;

    private:
	// Pointer to either the owning pool or the next object in the
	// pending list
	struct T_Storage_;
	union T_ExtraPointer_
	{
		T_Storage_* next;
		T_Self* pool;
	};

	// Structure that stores a single allocated item
	struct T_Storage_ final
	{
		T_PoolHelper::T_Object< ObjectSize , ObjectAlign > data;	
		void* list;
		T_ExtraPointer_ extra;
	};

	// The head for the "remote" freeing stack
	struct T_FreeHead_ final
	{
		uintptr_t aba{ 0 };
		T_Storage_* head{ nullptr };
	};

	// Atomic stack head, plus padding to prevent false sharing between
	// the stack head and the actual allocator.
	char padding_0_[ 64 ];
	std::atomic< T_FreeHead_ > freeHead_;
	char padding_1_[ 64 - sizeof( std::atomic< T_FreeHead_ > ) ];

	using T_Alloc_ = T_PoolHelper::T_Allocator<
		T_Storage_ , PerList , MaxFreeLists >;
	T_Alloc_ alloc_;
	
    public:
	T_ThreadedPoolAllocator( ) noexcept = default;
	T_ThreadedPoolAllocator( T_Self const& ) = delete;
	T_ThreadedPoolAllocator( T_Self&& ) = delete;

	/* Allocate an object. */
	void* allocate( size_t requested ) noexcept;
	/* Free an object */
	void free( void* item ) noexcept;

	/* Get the amount of lists. Mostly used for testing. */
	size_t countFreeLists( ) const noexcept;
	size_t countPartialLists( ) const noexcept;
	size_t countFullLists( ) const noexcept;

	/* Get the amount of total, free and used objects */
	void getUsage( size_t& total ,
			size_t& free ,
			size_t& used ) const noexcept;

    private:
	// Try to take an item from the stack.
	T_Storage_* takeFromStack( ) noexcept;
	// Add an item to the stack
	void addToStack( T_Storage_* item ) noexcept;
};


}
#endif //_H_EBCL_ALLOC
#include <lw/lib/inline/Alloc.hh>