/*
 * Copyright 2007, Hugo Santos. All Rights Reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *      Hugo Santos, hugosantos@gmail.com
 */

#ifndef _SLAB_H_
#define _SLAB_H_

#include <stdint.h>
#include <stdlib.h>

#include <algorithm> // for swap()
#include <new>
#include <utility> // for pair<>

#include <util/AutoLock.h>
#include <util/DoublyLinkedList.h>
#include <util/OpenHashTable.h>

#include <OS.h>


#define smp_get_current_cpu() 0
#define smp_get_num_cpus() 1


// C interface
extern "C" {
	typedef void *object_cache_t;
	object_cache_t
	object_cache_create(const char *name, size_t object_size, size_t alignment,
		void (*_constructor)(void *, void *), void (*_destructor)(void *, void *),
		void *cookie);
	void *object_cache_alloc(object_cache_t cache);
	void *object_cache_alloc_etc(object_cache_t cache, uint32_t flags);
	void object_cache_free(object_cache_t cache, void *object);
	void object_cache_destroy(object_cache_t cache);
}


// TODO this values should be dynamically tuned per cache.
static const int kMinimumSlabItems = 32;

// base Slab implementation, opaque to the backend used.
class BaseCache {
public:
	typedef void (*Constructor)(void *cookie, void *object);
	typedef void (*Destructor)(void *cookie, void *object);

	BaseCache(const char *_name, size_t objectSize, size_t alignment,
		Constructor _constructor, Destructor _destructor, void *_cookie);

	struct ObjectLink {
		struct ObjectLink *next;
	};

	struct Slab : DoublyLinkedListLinkImpl<Slab> {
		void *pages;
		size_t count, size;
		ObjectLink *free;
	};

	typedef std::pair<Slab *, ObjectLink *> ObjectInfo;

	ObjectLink *AllocateObject();
	bool ReturnObject(const ObjectInfo &object);

	Slab *ConstructSlab(Slab *slab, void *pages, size_t byteCount,
		ObjectLink *(*getLink)(void *parent, void *object), void *parent);
	void DestructSlab(Slab *slab);

	const char *Name() const { return fName; }
	size_t ObjectSize() const { return fObjectSize; }

protected:
	typedef DoublyLinkedList<Slab> SlabList;

	char fName[32];
	size_t fObjectSize, fCacheColorCycle;
	SlabList fPartialSlabs, fFullSlabs;

	Constructor fConstructor;
	Destructor fDestructor;
	void *fCookie;
};


enum {
	CACHE_ALIGN_TO_PAGE_TOTAL	= 1 << 0,
};

struct MallocBackend {
	typedef void *AllocationID;

	static int PageSize() { return 4096; }
	static status_t AllocatePages(BaseCache *cache, AllocationID *id,
		void **pages, size_t byteCount, uint32_t flags);
	static void FreePages(BaseCache *cache, void *pages);
};


struct AreaBackend {
	typedef area_id AllocationID;

	static int PageSize() { return B_PAGE_SIZE; }
	static status_t AllocatePages(BaseCache *cache, area_id *id, void **pages,
		size_t byteCount, uint32_t flags);
	static void FreePages(BaseCache *cache, area_id id);
};


// Slab implementation, glues together the frontend, backend as
// well as the Slab strategy used.
template<typename Strategy>
class Cache : protected BaseCache {
public:
	Cache(const char *_name, size_t objectSize, size_t alignment,
		Constructor _constructor, Destructor _destructor, void *_cookie)
		: BaseCache(_name, Strategy::RequiredSpace(objectSize), alignment,
			_constructor, _destructor, _cookie), fStrategy(this) {}

	void *AllocateObject(uint32_t flags)
	{
		if (fPartialSlabs.IsEmpty()) {
			Slab *newSlab = fStrategy.NewSlab(flags);
			if (newSlab == NULL)
				return NULL;
			fPartialSlabs.Add(newSlab);
		}

		return fStrategy.Object(BaseCache::AllocateObject());
	}

	void ReturnObject(void *object)
	{
		ObjectInfo location = fStrategy.ObjectInformation(object);

		if (BaseCache::ReturnObject(location))
			fStrategy.ReturnSlab(location.first);
	}

private:
	Strategy fStrategy;
};


static inline const void *
LowerBoundary(void *object, size_t byteCount)
{
	const uint8_t *null = (uint8_t *)NULL;
	return null + ((((uint8_t *)object) - null) & ~(byteCount - 1));
}


template<typename Backend>
class BaseCacheStrategy {
protected:
	typedef BaseCache::ObjectLink ObjectLink;

	BaseCacheStrategy(BaseCache *parent)
		: fParent(parent) {}

	size_t SlabSize(size_t tailSpace) const
	{
		size_t pageCount = (kMinimumSlabItems * fParent->ObjectSize()
			+ tailSpace) / Backend::PageSize();
		if (pageCount < 1)
			pageCount = 1;
		return pageCount * Backend::PageSize();
	}

	struct Slab : BaseCache::Slab {
		typename Backend::AllocationID id;
	};

	BaseCache::Slab *_ConstructSlab(Slab *slab, void *pages, size_t tailSpace,
		ObjectLink *(*getLink)(void *parent, void *object), void *parent)
	{
		return fParent->ConstructSlab(slab, pages, SlabSize(tailSpace)
			- tailSpace, getLink, parent);
	}

	void _DestructSlab(BaseCache::Slab *slab)
	{
		fParent->DestructSlab(slab);
		Backend::FreePages(fParent, ((Slab *)slab)->id);
	}

	BaseCache *fParent;
};


// This slab strategy includes the ObjectLink at the end of each object and the
// slab at the end of the allocated pages. It uses aligned allocations to
// provide object to slab mapping with zero storage, thus there is only one
// word of overhead per object. This is optimized for small objects.
template<typename Backend>
class MergedLinkCacheStrategy : public BaseCacheStrategy<Backend> {
public:
	typedef typename BaseCacheStrategy<Backend>::Slab Slab;
	typedef BaseCache::ObjectLink Link;
	typedef BaseCache::ObjectInfo ObjectInfo;

	MergedLinkCacheStrategy(BaseCache *parent)
		: BaseCacheStrategy<Backend>(parent) {}

	static size_t RequiredSpace(size_t objectSize)
	{
		return objectSize + sizeof(Link);
	}

	void *Object(Link *link) const
	{
		return ((uint8_t *)link) - (fParent->ObjectSize() - sizeof(Link));
	}

	ObjectInfo ObjectInformation(void *object) const
	{
		Slab *slab = _SlabInPages(LowerBoundary(object, _SlabSize()));
		return ObjectInfo(slab, _Linkage(object));
	}

	BaseCache::Slab *NewSlab(uint32_t flags)
	{
		typename Backend::AllocationID id;
		void *pages;

		// in order to save a pointer per object or a hash table to
		// map objects to slabs we required this set of pages to be
		// aligned in a (pageCount * PAGE_SIZE) boundary.
		if (Backend::AllocatePages(fParent, &id, &pages, _SlabSize(),
				CACHE_ALIGN_TO_PAGE_TOTAL | flags) < B_OK)
			return NULL;

		_SlabInPages(pages)->id = id;

		return BaseCacheStrategy<Backend>::_ConstructSlab(_SlabInPages(pages),
			pages, sizeof(Slab), _Linkage, this);
	}

	void ReturnSlab(BaseCache::Slab *slab)
	{
		BaseCacheStrategy<Backend>::_DestructSlab(slab);
	}

private:
	size_t _SlabSize() const
	{
		return BaseCacheStrategy<Backend>::SlabSize(sizeof(Slab));
	}

	Link *_Linkage(void *object) const
	{
		return (Link *)(((uint8_t *)object)
			+ (fParent->ObjectSize() - sizeof(Link)));
	}

	Slab *_SlabInPages(const void *pages) const
	{
		return (Slab *)(((uint8_t *)pages) + _SlabSize() - sizeof(Slab));
	}

	static Link *_Linkage(void *_this, void *object)
	{
		return ((MergedLinkCacheStrategy *)_this)->_Linkage(object);
	}
};


template<typename Type, typename Backend>
class TypedCache : public Cache<MergedLinkCacheStrategy<Backend> > {
public:
	typedef MergedLinkCacheStrategy<Backend> Strategy;
	typedef Cache<Strategy> BaseType;

	TypedCache(const char *name, size_t alignment)
		: BaseType(name, sizeof(Type), alignment, _ConstructObject,
			_DestructObject, this) {}
	virtual ~TypedCache() {}

	Type *Alloc(uint32_t flags) { return (Type *)BaseType::AllocateObject(flags); }
	void Free(Type *object) { BaseType::ReturnObject(object); }

private:
	static void _ConstructObject(void *cookie, void *object)
	{
		((TypedCache *)cookie)->ConstructObject((Type *)object);
	}

	static void _DestructObject(void *cookie, void *object)
	{
		((TypedCache *)cookie)->DestructObject((Type *)object);
	}

	virtual void ConstructObject(Type *object) {}
	virtual void DestructObject(Type *object) {}
};


static inline int
Fls(size_t value)
{
	for (int i = 31; i >= 0; i--) {
		if ((value >> i) & 1)
			return i + 1;
	}

	return -1;
}


struct BaseHashCacheStrategy {
	typedef BaseCache::ObjectLink ObjectLink;
	typedef BaseCache::ObjectInfo ObjectInfo;

	struct Link : ObjectLink, HashTableLink<Link> {
		BaseCache::Slab *slab;
		void *buffer;
	};

	struct HashTableDefinition {
		typedef BaseHashCacheStrategy	ParentType;
		typedef void *					KeyType;
		typedef Link					ValueType;

		HashTableDefinition(BaseHashCacheStrategy *_parent) : parent(_parent) {}

		size_t HashKey(void *key) const
		{
			return (((uint8_t *)key) - ((uint8_t *)0)) >> parent->fLowerBoundary;
		}

		size_t Hash(Link *value) const { return HashKey(value->buffer); }

		bool Compare(void *key, Link *value) const
		{
			return value->buffer == key;
		}

		HashTableLink<Link> *GetLink(Link *value) const { return value; }

		BaseHashCacheStrategy *parent;
	};

	// for g++ 2.95
	friend class HashTableDefinition;

	typedef OpenHashTable<HashTableDefinition> HashTable;

	BaseHashCacheStrategy(BaseCache *parent)
		: fHashTable(this), fLowerBoundary(Fls(parent->ObjectSize()) - 1) {}

	void *Object(ObjectLink *link) const
	{
		return ((Link *)link)->buffer;
	}

	ObjectInfo ObjectInformation(void *object) const
	{
		Link *link = _Linkage(object);
		return ObjectInfo(link->slab, link);
	}

protected:
	Link *_Linkage(void *object) const
	{
		return fHashTable.Lookup(object);
	}

	static ObjectLink *_Linkage(void *_this, void *object)
	{
		return ((BaseHashCacheStrategy *)_this)->_Linkage(object);
	}

	HashTable fHashTable;
	const size_t fLowerBoundary;
};


template<typename Backend>
struct HashCacheStrategy : BaseCacheStrategy<Backend>, BaseHashCacheStrategy {
	typedef typename BaseCacheStrategy<Backend>::Slab Slab;
	typedef HashCacheStrategy<Backend> Strategy;

	HashCacheStrategy(BaseCache *parent)
		: BaseCacheStrategy<Backend>(parent), BaseHashCacheStrategy(parent),
			fSlabCache("slab cache", 0), fLinkCache("link cache", 0) {}

	static size_t RequiredSpace(size_t objectSize)
	{
		return objectSize;
	}

	BaseCache::Slab *NewSlab(uint32_t flags)
	{
		size_t byteCount = _SlabSize();

		Slab *slab = fSlabCache.Alloc(flags);
		if (slab == NULL)
			return NULL;

		void *pages;
		if (Backend::AllocatePages(fParent, &slab->id, &pages, byteCount,
				flags) < B_OK) {
			fSlabCache.Free(slab);
			return NULL;
		}

		if (_PrepareSlab(fParent, slab, pages, byteCount, flags) < B_OK) {
			Backend::FreePages(fParent, slab->id);
			fSlabCache.Free(slab);
			return NULL;
		}

		// it's very important that we cast this to BaseHashCacheStrategy
		// so we get the proper instance offset through void *
		return BaseCacheStrategy<Backend>::_ConstructSlab(slab, pages, 0,
			_Linkage, (BaseHashCacheStrategy *)this);
	}

	void ReturnSlab(BaseCache::Slab *slab)
	{
		_ClearSlab(fParent, slab->pages, _SlabSize());
		BaseCacheStrategy<Backend>::_DestructSlab(slab);
		fSlabCache.Free((Slab *)slab);
	}

private:
	size_t _SlabSize() const
	{
		return BaseCacheStrategy<Backend>::SlabSize(0);
	}

	status_t _PrepareSlab(BaseCache *parent, Slab *slab, void *pages,
		size_t byteCount, uint32_t flags)
	{
		uint8_t *data = (uint8_t *)pages;
		for (uint8_t *it = data;
				it < (data + byteCount); it += parent->ObjectSize()) {
			Link *link = fLinkCache.Alloc(flags);

			if (link == NULL) {
				_ClearSlabRange(parent, data, it);
				return B_NO_MEMORY;
			}

			link->slab = slab;
			link->buffer = it;
			fHashTable.Insert(link);
		}

		return B_OK;
	}

	void _ClearSlab(BaseCache *parent, void *pages, size_t byteCount)
	{
		_ClearSlabRange(parent, (uint8_t *)pages,
			((uint8_t *)pages) + byteCount);
	}

	void _ClearSlabRange(BaseCache *parent, uint8_t *data, uint8_t *end)
	{
		for (uint8_t *it = data; it < end; it += parent->ObjectSize()) {
			Link *link = fHashTable.Lookup(it);
			fHashTable.Remove(link);
			fLinkCache.Free(link);
		}
	}

	TypedCache<Slab, Backend> fSlabCache;
	TypedCache<Link, Backend> fLinkCache;
};


class BaseDepot {
public:
	struct Magazine {
		Magazine *next;
		uint16_t current_round, round_count;
		void *rounds[0];
	};

	struct CPUStore {
		benaphore lock;
		Magazine *loaded, *previous;
	};

protected:
	BaseDepot();
	virtual ~BaseDepot();

	status_t InitCheck() const;

	CPUStore *CPU() const { return &fStores[smp_get_current_cpu()]; }

	void *ObtainFromStore(CPUStore *store);
	bool ReturnToStore(CPUStore *store, void *object);
	void MakeEmpty();

	virtual void ReturnObject(void *object) = 0;

	bool _ExchangeWithFull(Magazine* &magazine);
	bool _ExchangeWithEmpty(Magazine* &magazine);
	void _EmptyMagazine(Magazine *magazine);

	Magazine *_AllocMagazine();
	void _FreeMagazine(Magazine *magazine);

	benaphore fLock;
	Magazine *fFull, *fEmpty;
	size_t fFullCount, fEmptyCount;
	CPUStore *fStores;
};


template<typename CacheType>
class LocalCache : public CacheType, protected BaseDepot {
public:
	typedef typename CacheType::Constructor Constructor;
	typedef typename CacheType::Destructor Destructor;

	LocalCache(const char *name, size_t objectSize, size_t alignment,
		Constructor _constructor, Destructor _destructor, void *_cookie)
		: CacheType(name, objectSize, alignment, _constructor, _destructor,
			_cookie) {}

	~LocalCache() { Destroy(); }

	void *Alloc(uint32_t flags)
	{
		void *object = BaseDepot::ObtainFromStore(CPU());
		if (object == NULL)
			object = CacheType::AllocateObject(flags);
		return object;
	}

	void Free(void *object)
	{
		if (!BaseDepot::ReturnToStore(CPU(), object))
			CacheType::ReturnObject(object);
	}

	void Destroy() { BaseDepot::MakeEmpty(); }

private:
	void ReturnObject(void *object)
	{
		CacheType::ReturnObject(object);
	}
};

#endif