(cppcon)Practical memory pool based allocators for Modern C++

class bucket {
public:
	const std::size_t BlockSize;
	const std::size_t BlockCount;
	bucket(std::size_t block_size, std::size_t block_count);
	~bucket();
	// Tests if the pointer belongs to this bucket
	bool belongs(void * ptr) const noexcept;
	// Returns nullptr if failed
	[[nodiscard]] void * allocate(std::size_t bytes) noexcept;
	void deallocate(void * ptr, std::size_t bytes) noexcept;
private:
	// Finds n free contiguous blocks in the ledger and returns the first block’s index or BlockCount on failure
	std::size_t find_contiguous_blocks(std::size_t n) const noexcept;
	// Marks n blocks in the ledger as “in-use” starting at ‘index’
	void set_blocks_in_use(std::size_t index, std::size_t n) noexcept;
	// Marks n blocks in the ledger as “free” starting at ‘index’
	void set_blocks_free(std::size_t index, std::size_t n) noexcept;
	// Actual memory for allocations
	std::byte* m_data{nullptr};
	// Reserves one bit per block to indicate whether it is in-use
	std::byte* m_ledger{nullptr};
};

bucket::bucket(std::size_t block_size, std::size_t block_count)
: BlockSize{block_size}
, BlockCount{block_count}
{
	const auto data_size = BlockSize * BlockCount;
	m_data = static_cast<std::byte*>(std::malloc(data_size));
	assert(m_data != nullptr);
	const auto ledger_size = 1 + ((BlockCount - 1) / 8);
	m_ledger = static_cast<std::byte*>(std::malloc(ledger_size));
	assert(m_ledger != nullptr);
	std::memset(m_data, 0, data_size);
	std::memset(m_ledger, 0, ledger_size);
}
bucket::~bucket() {
	std::free(m_ledger);
	std::free(m_data);
}


void * bucket::allocate(std::size_t bytes) noexcept {
	// Calculate the required number of blocks
	const auto n = 1 + ((bytes - 1) / BlockSize);
	const auto index = find_contiguous_blocks(n);
	if (index == BlockCount) {
		return nullptr;
	}
	set_blocks_in_use(index, n);
	return m_data + (index * BlockSize);
}

void bucket::deallocate(void * ptr, std::size_t bytes) noexcept {
	const auto p = static_cast<const std::byte *>(ptr);
	const std::size_t dist = static_cast<std::size_t>(p - m_data);
	// Calculate block index from pointer distance
	const auto index = dist / BlockSize;
	// Calculate the required number of blocks
	const auto n = 1 + ((bytes - 1) / BlockSize);
	// Update the ledger
	set_blocks_free(index, n);
}

// The default implementation defines a pool with no buckets
template<std::size_t id>
struct bucket_descriptors {
	using type = std::tuple<>;
};

struct bucket_cfg16 {
	static constexpr std::size_t BlockSize = 16;
	static constexpr std::size_t BlockCount = 10000;
};
struct bucket_cfg32{
	static constexpr std::size_t BlockSize = 32;
	static constexpr std::size_t BlockCount = 10000;
};
struct bucket_cfg1024 {
	static constexpr std::size_t BlockSize = 1024;
	static constexpr std::size_t BlockCount = 50000;
};
template<>
struct bucket_descriptors<1> {
	using type = std::tuple<bucket_cfg16, bucket_cfg32, bucket_cfg1024>;
};

template<std::size_t id>
using bucket_descriptors_t = typename bucket_descriptors<id>::type;

template<std::size_t id>
static constexpr std::size_t bucket_count = std::tuple_size<bucket_descriptors_t<id>>::value;


template<std::size_t id>
using pool_type = std::array<bucket, bucket_count<id>>;

template<std::size_t id, std::size_t Idx>
struct get_size
    : std::integral_constant<std::size_t, std::tuple_element_t<Idx, bucket_descriptors_t<id>>::BlockSize>{\\
};
    
template<std::size_t id, std::size_t Idx>
struct get_count
    : std::integral_constant<std::size_t, std::tuple_element_t<Idx, bucket_descriptors_t<id>>::BlockCount>{\\
};

template<std::size_t id, std::size_t... Idx>
auto & get_instance(std::index_sequence<Idx...>) noexcept {
	static pool_type<id> instance{\{\{get_size<id, Idx>::value, get_count<id, Idx>::value} ...\}\};
	return instance;
}
template<std::size_t id>
auto & get_instance() noexcept {
	return get_instance<id>(std::make_index_sequence<bucket_count<id>>());
}

// Assuming buckets are sorted by their block sizes
template<std::size_t id>
[[nodiscard]] void * allocate(std::size_t bytes) {
	auto & pool = get_instance<id>();
	for (auto & bucket : pool) {
		if(bucket.BlockSize >= bytes) {
			if(auto ptr = bucket.allocate(bytes); ptr != nullptr) {
				return ptr;
			}
		}
	}
	throw std::bad_alloc{};
}

template<std::size_t id>
[[nodiscard]] void * allocate(std::size_t bytes) {
	auto & pool = get_instance<id>();
	std::array<info, bucket_count<id>> deltas;
	std::size_t index = 0;
	for (const auto & bucket : pool) {
		deltas[index].index = index;
		if (bucket.BlockSize >= bytes) {
			deltas[index].waste = bucket.BlockSize - bytes;
			deltas[index].block_count = 1;
		} else {
			const auto n = 1 + ((bytes - 1) / bucket.BlockSize);
			const auto storage_required = n * bucket.BlockSize;
			deltas[index].waste = storage_required - bytes;
			deltas[index].block_count = n;
		}
		++index;
	}

    sort(deltas.begin(), deltas.end()); // std::sort() is allowed to allocate
    
	for (const auto & d : deltas)
		if (auto ptr = pool[d.index].allocate(bytes); ptr != nullptr)
			return ptr;
	
    throw std::bad_alloc{};
}

	template<typename T = std::uint8_t, std::size_t id = 0>
class static_pool_allocator {
public:
	//rebind不用实现吧，我记得好像废弃了
    template<typename U>
	static_pool_allocator(const static_pool_allocator<U, id> & other) noexcept
		: m_upstream_resource{other.upstream_resource()} {}
	template<typename U>
	static_pool_allocator & operator=(const static_pool_allocator<U, id> & other) noexcept {
		m_upstream_resource = other.upstream_resource();
		return *this;
	}
	static bool initialize_memory_pool() noexcept { return memory_pool::initialize<id>(); }
private:
	pmr::memory_resource * m_upstream_resource;
};

class AllocListPass : public llvm::FunctionPass {
public:
	static char ID;
	AllocListPass() : llvm::FunctionPass(ID) {}

    bool runOnFunction(llvm::Function & f) override {
		const auto pretty_name = boost::core::demangle(f.getName().str().c_str());
		static const std::regex call_regex{R"(void instrument::type_reg<([^,]+),(.+),([^,]+)>\(\))"};
		std::smatch match;
		if (std::regex_match(pretty_name, match, call_regex)) {
			if (match.size() == 4) {
				const auto pool_id = std::atoi(match[1].str().c_str());
				const auto type = match[2].str();
				const auto size = std::atoi(match[3].str().c_str());
				std::cout << "Pool ID: " << pool_id << ", Size: " << size << ", Type: " << type << "\n";
			}
		}
		return false; // does not alter the code, a read-only pass
	}
};
char AllocListPass::ID = 0;
static llvm::RegisterPass<AllocListPass> dummy("alloc-list", "This pass lists memory pool allocations");