duckhandy/include/duckhandy/small_object_allocator.hpp

/* Copyright 2016-2024 Michele Santullo
 * This file is part of "duckhandy".
 *
 * "duckhandy" is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * "duckhandy" is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with "duckhandy".  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef idE77208CAFC79452DA12757DD0F6692D3
#define idE77208CAFC79452DA12757DD0F6692D3

#include "alignment.hpp"
#include <type_traits>
#include <cassert>
#include <climits>
#include <cstdint>
#include <ciso646>
#include <memory>
#include <strings.h>
#include <cstddef>
#include <cstring>
#include <iterator>
#include <vector>
#include <algorithm>

#if !defined(NDEBUG) && !defined(DEBUG_SMALL_OBJECT_ALLOCATOR)
#	define DEBUG_SMALL_OBJECT_ALLOCATOR
#endif

#if defined(DEBUG_SMALL_OBJECT_ALLOCATOR)
#	include "tiger_bt.hpp"
#	include "lengthof.h"
#endif

namespace dhandy {
namespace implem {
[[gnu::pure,gnu::always_inline]]
unsigned int ffs (unsigned int val) { return static_cast<unsigned int>(::ffs(static_cast<int>(val))); }

[[gnu::pure,gnu::always_inline]]
unsigned int ffs (unsigned long val) { return static_cast<unsigned int>(::ffsl(static_cast<long>(val))); }

[[gnu::pure,gnu::always_inline]]
unsigned int ffs (unsigned long long val) { return static_cast<unsigned int>(::ffsll(static_cast<long long>(val))); }

#if defined(DEBUG_SMALL_OBJECT_ALLOCATOR)
template <typename T, template <class> typename A>
consteval std::uint32_t make_signature() {
	return static_cast<std::uint32_t>(
		bt::tiger(__PRETTY_FUNCTION__, lengthof(__PRETTY_FUNCTION__), bt::TigerPaddingV2).a & 0xFFFFFFFF
	);
}
#endif
} //namespace implem

template <typename T>
struct AllocatorFunction {
	std::unique_ptr<T> operator()() {
		return std::make_unique<T>();
	}
};

template <typename T, template <class> typename A=AllocatorFunction>
class SmallObjectAllocator {
public:
	typedef T value_type;
	typedef T* pointer;
	typedef const T* const_pointer;
	typedef T& reference;
	typedef const T& const_reference;
	typedef std::size_t size_type;
	typedef std::ptrdiff_t difference_type;
	typedef std::true_type propagate_on_container_move_assignment;

private:
	typedef std::uint_fast32_t uint_freelist_t;
	typedef std::uint32_t index_t;

	struct TAndPtrFakeStruct { //only used for sizeof()
		TAndPtrFakeStruct() = delete;
		T t;
#if defined(DEBUG_SMALL_OBJECT_ALLOCATOR)
		std::uint32_t signature;
#endif
		index_t block_index; //index in a lookup table so that lookup[index-1]->next == owner
	};

	static constexpr std::uint32_t signature = implem::make_signature<T, A>();
	static constexpr size_type size = sizeof(TAndPtrFakeStruct);
	static constexpr size_type align = alignof(TAndPtrFakeStruct);
	static constexpr size_type object_size = sizeof(T); //(size < sizeof(T*) ? sizeof(T*) : size);
	static constexpr size_type object_align = alignof(T); //(align < alignof(T*) ? alignof(T*) : align);
	static constexpr size_type block_ptr_offset = offsetof(TAndPtrFakeStruct, block);

public:
	static constexpr size_type objects_per_block = CHAR_BIT * sizeof(uint_freelist_t);

private:
	struct Block {
		typedef typename std::aligned_storage<size, align>::type raw_t;

		raw_t data[objects_per_block];
		std::unique_ptr<Block> next;
		uint_freelist_t freelist{0};
	};

public:

	SmallObjectAllocator() = default;
	~SmallObjectAllocator() noexcept = default;

	[[nodiscard]] constexpr pointer allocate (size_type size);
	//[[nodiscard]] constexpr std::allocation_result<pointer, size_type> allocate_at_least (size_type size);
	constexpr void deallocate (pointer ptr, size_type size);

	//template <typename T1, typename T2>
	//constexpr bool operator== (const allocator<T1& lh, const allocator<T2>& rhs) noexcept;


private:
	static index_t* fetch_block_index (Block::raw_t* in_ptr);
	static void set_block_indices (Block* block, index_t new_index);
	static void set_debug_signatures (Block* block);

	std::vector<Block*> m_block_list;
	std::unique_ptr<Block> m_head;
	index_t m_prev_index;
	index_t m_curr_index{0};
};

template <typename T, template <class> typename A>
constexpr auto SmallObjectAllocator<T, A>::allocate (size_type size) -> pointer {
	assert(object_size == size);
#if defined(NDEBUG)
	static_cast<void>(size);
#endif

	assert(0 == m_curr_index or m_curr_index <= m_block_list.size());
	assert(m_head or m_block_list.empty());
	if (!m_curr_index or !~m_block_list[m_curr_index-1]->freelist) {
		A<Block> alloc_function{};
		auto new_block = alloc_function();
		set_debug_signatures(new_block.get());
		new_block->next = std::move(m_head);
		m_head.swap(new_block);
		m_block_list.push_back(m_head.get());
		m_prev_index = m_curr_index;
		m_curr_index = static_cast<index_t>(m_block_list.size());

		if (m_head->next)
			set_block_indices(m_head->next.get(), m_curr_index);
	}

	Block*const block = m_block_list[m_curr_index-1];
	assert(block != nullptr);
	const uint_freelist_t neg_freelist = static_cast<uint_freelist_t>(~block->freelist);
	assert(neg_freelist != 0);

	const unsigned int object_num = implem::ffs(neg_freelist);
	assert(object_num);
	const unsigned int object_index = object_num - 1;
	assert(object_index < objects_per_block);
	constexpr uint_freelist_t one = 1;
	block->freelist |= one << object_index;

	index_t* const index_ptr = fetch_block_index(block->data + object_index);
	assert(index_ptr);
	*index_ptr = m_prev_index;
	return reinterpret_cast<pointer>(block->data + object_index);
}

template <typename T, template <class> typename A>
constexpr void SmallObjectAllocator<T, A>::deallocate (pointer ptr, size_type size) {
	assert(object_size == size);
	auto* const block_ptr = reinterpret_cast<Block::raw_t*>(ptr);
	const index_t owner_prev_index = *fetch_block_index(block_ptr);
	Block* const owner = (0 == owner_prev_index ? m_head.get() : m_block_list[owner_prev_index-1]->next.get());

	assert(owner);
	const auto flag_index = std::distance(owner->data, block_ptr);
	assert(flag_index >= 0);
	assert(flag_index < objects_per_block);

	constexpr uint_freelist_t one = 1;
	const uint_freelist_t mask = one << flag_index;
	assert((owner->freelist & mask) != 0);
	owner->freelist &= ~mask;

	if (0 == owner->freelist) {
		if (owner_prev_index) {
			Block& prev = *m_block_list[owner_prev_index-1];
			auto found = std::find(m_block_list.begin(), m_block_list.end(), owner);
			assert(m_block_list.end() != found);
			it's all fucked up here, I can't delete the item because that would
			invalidate all the subsequent indices. It means all blocks following
			the one being deleted would have to get set_block_indices() invoked
			on them which is insane so this whole approach is fucke'

			auto empty_block = std::move(prev.next);
			assert(!prev.next);
			prev.next = std::move(empty_block->next);
			set_block_indices(prev.next.get(), owner_prev_index);

		}
		else {
			assert(m_block_list.empty());
			m_head.reset();
			m_curr_index = 0;
		}
	}
}

template <typename T, template <class> typename A>
auto SmallObjectAllocator<T, A>::fetch_block_index (Block::raw_t* in_ptr) -> index_t* {
	assert(in_ptr);

	char* ptr = reinterpret_cast<char*>(in_ptr);
#if defined(DEBUG_SMALL_OBJECT_ALLOCATOR)
	ptr = align_ptr_to<alignof(std::uint32_t)>(ptr, object_size);
	{
		std::uint32_t read_signature;
		std::memcpy(&read_signature, ptr, sizeof(signature));
		assert(signature == read_signature);
	}
	ptr = align_ptr_to<alignof(index_t)>(ptr, sizeof(signature));
#else
	ptr = align_ptr_to<alignof(index_t)>(ptr, object_size);
#endif

	index_t* const retval = reinterpret_cast<index_t*>(ptr);
	return retval;
}

template <typename T, template <class> typename A>
void SmallObjectAllocator<T, A>::set_block_indices (Block* block, index_t new_index) {
	uint_freelist_t flag = 1;
	for (unsigned int z = 0; z < objects_per_block; ++z, flag<<=1) {
		if (flag & block->freelist) {
			index_t* const dst_index = fetch_block_index(block->data + z);
			assert(dst_index);
			*dst_index = new_index;
		}
	}
}

template <typename T, template <class> typename A>
void SmallObjectAllocator<T, A>::set_debug_signatures (Block* block) {
#if defined(DEBUG_SMALL_OBJECT_ALLOCATOR)
	const auto sig = signature;
	for (unsigned int z = 0; z < objects_per_block; ++z) {
		std::memcpy(align_ptr_to<alignof(std::uint32_t)>(block->data + z, object_size), &sig, sizeof(signature));
	}
#else
	static_cast<void>(block);
#endif
}
} //namespace dhandy

#endif