gplib/include/gp/allocator/buddy.hpp

#pragma once

#include "gp_config.hpp"

#include "gp/algorithm/modifiers.hpp"
#include "gp/algorithm/tmp_manip.hpp"
#include "gp/allocator/dummy.hpp"
#include "gp/array.hpp"
#include "gp/buffer.hpp"
#include "gp/math/integer_math.hpp"
#include "gp/allocator/allocator.hpp"

#include <type_traits>


namespace gp{
	/**
	 * @brief An allocator that uses the buddy algorithm to divide the space into allocatable memory.
	 * 
	 * This is not resilient to memory fragmentation, but the smallest memory unit should always be allocatable unless the memory is actually full.
	 * 
	 * @tparam max_msb The log2 rounded up of the maximum space you expect to address in bytes
	 * @tparam align The smallest size of memory you expect the allocator to allocate
	 */
	template<size_t max_msb = 24, size_t align = 8>
	class buddy : public allocator {
		struct twig {
			bool used : 1;
			bool used_children : 1;
			twig(uint8_t src) {
				used = 1 & src;
				used_children = 2 & src;
			}

			operator uint8_t() {
				return 1 * used + 2 * used_children;
			}
		};
		struct bundle {
			uint8_t a : 2;
			uint8_t b : 2;
			uint8_t c : 2;
			uint8_t d : 2;

			bundle() {
				a = 0; b = 0; c = 0; d = 0;
			}
		};
		gp::optional<gp::reference_wrapper<allocator>> allocator_v;
		gp::buffer<char> data;
		const size_t max_depth; 
		const size_t twig_explore_length;

		/**
		 * @brief The depth of the tree required to allocate
		 */
		static constexpr size_t max_theoric_depth = max_msb - gp::math::msb(align);

		/**
		 * @brief The actual number of twigs to support the specified depth
		 */
		static constexpr size_t required_twigs = (1 << (max_theoric_depth + 1)) - 1;

		/**
		 * @brief ((max allocatable size - min allocatable size) ** 2 - 1) / 4 twigs in a bundle
		**/
		static constexpr size_t span_size = required_twigs / 4 + (required_twigs % 4 != 0);

		/**
		 * @brief The array of twigs (in bundles)
		 */
		gp::array<bundle, span_size> stack;

/**
 * This code has been manually checked and will always return.
 * If you find a case where it doesn't, please file an issue.
**/
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wreturn-type"
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wreturn-type"
		twig get_twig(size_t idx) const {
			auto far = idx / 4;
			auto local = idx % 4;
			switch(local) {
				case 0:
					return stack[far].a;
				case 1:
					return stack[far].b;
				case 2:
					return stack[far].c;
				case 3:
					return stack[far].d;
			}
		}
#pragma GCC diagnostic pop
#pragma clang diagnostic pop

		void set_twig(size_t idx, twig v) {
			auto far = idx / 4;
			auto local = idx % 4;
			auto& group = stack[far];
			
			switch(local) {
				case 0:
					group.a = v;
					return;
				case 1:
					group.b = v;
					return;
				case 2:
					group.c = v;
					return;
				case 3:
					group.d = v;
					return;
			}
		}

		constexpr size_t size_to_depth(size_t sz) {
			size_t pow2 = gp::math::msb(sz) - gp::math::msb(align);
			return gp::clamp(
				(size_t)0 ,
				max_depth - pow2, 
				max_depth
			);
		}

		constexpr size_t depth_to_size(size_t depth) {
			return 1 << (max_depth - depth + gp::math::msb(align));
		}

		constexpr size_t get_left(size_t index) const {
			return ((index + 1) << 1) - 1;
		}

		constexpr size_t get_right(size_t index) const {
			return ((index + 1) << 1);
		}

		template<typename function>
		void all_under(size_t index, function func) {
			size_t left = get_left(index);
			size_t right = get_right(index);
			all_under(left, func);
			all_under(right, func);
			func(left);
			func(right);
		}

		template<typename function>
		void all_over(size_t index, function func) {
			if(index != 0) {
				size_t parent = ((index + 1) >> 1) - 1;
				func(parent);
				if(parent != 0)
					all_over(parent, func);
			}
		}

		template<typename function>
		bool is_any_child(size_t index, function func) const {
			size_t left = get_left(index);
			size_t right = get_right(index);
			if(left < twig_explore_length && right < twig_explore_length) {
				if(func(left)) return true;
				if(func(right)) return true;
				if(is_any_child(left, func)) return true;
				if(is_any_child(right, func)) return true;
			}
			return false;
		}

		static constexpr size_t no_twig = -1;

		size_t find_free_twig(size_t depth, size_t root = 0, size_t explored = 0) const {
			auto v = get_twig(root);
			if(depth == explored) {
				if(v.used || v.used_children)
				{
					return no_twig;
				} else {
					return root;
				}
			} else {
				if(v.used)
				{
					return no_twig;
				} 
				++explored;
				auto ret = find_free_twig(depth, get_right(root), explored);
				if(ret != no_twig)
				{
					return ret;
				} 
				ret = find_free_twig(depth, get_left(root), explored);
				if(ret != no_twig)
				{
					return ret;
				}
			}
			return no_twig;
		}

		size_t find_used_twig(size_t offset, size_t root = 0, size_t explored = 0) {
			auto v = get_twig(root);
			if(v.used && offset == 0)
			{
				return root;
			}
			++explored;
			if(explored > max_depth) return no_twig;
			size_t cut = (1 << (max_depth + gp::math::log2(align))) >> explored;
			if(offset >= cut)
			{
				return find_used_twig(offset-cut, get_right(root), explored);
			} else {
				return find_used_twig(offset, get_left(root), explored);
			}
		}


		static bool empty_node(const buddy* me, size_t node) {
			gp_config::assertion(node < me->twig_explore_length, "bad emptyness test");
			auto p = me->get_twig(node);
			return !(
				p.used | p.used_children
			);
		}
	public:
		/**
		 * @brief Construct a new empty buddy object
		 */
		buddy()
		: data(gp::buffer<char>(nullptr,nullptr))
		, max_depth(0)
		, twig_explore_length(1 << max_depth)
		{}

		/**
		 * @brief Construct a new buddy object from another allocator
		 * 
		 * @param sz the size to allocate
		 * @param allocator_p the source of the allocated memory
		 */
		buddy(size_t sz, allocator& allocator_p)
		: allocator_v(allocator_p)
		, data(nullptr,nullptr)
		, max_depth(gp::math::msb(sz)-gp::math::msb(align))
		, twig_explore_length(1 << max_depth)
		{
			if(sz!=0 && (sz & (sz - 1)) == 0)
			{
				auto v=allocator_v.value().get().allocate(sz);
				if(v!=nullptr)
				{
					if((reinterpret_cast<intptr_t>(v) % align) ==0)
					{
						data=gp::buffer<char>(reinterpret_cast<char*>(v),reinterpret_cast<char*>(v)+sz);
					}
					else
					{
						allocator_v.value().get().deallocate(v);
					}
				}
			}
		}

		/**
		 * @brief Construct a new buddy object from a memory location
		 * 
		 * @param pos the location of the memory to manage
		 * @param sz the size of the span to manage
		 */
		buddy(char* pos, size_t sz)
		: data(pos,pos+sz)
		, max_depth(gp::math::msb(sz)-gp::math::msb(align))
		, twig_explore_length(1 << max_depth)
		{
		}

		virtual void* allocate(size_t sz)
		{
			auto depth = size_to_depth(sz);
			auto index = find_free_twig(depth);
			if(index == no_twig) 
			{
				return nullptr;
			}

			auto pot = reinterpret_cast<char*>(
				(index - (1 << depth) + 1)*depth_to_size(depth)
				+ reinterpret_cast<intptr_t>(&*data.begin())
			);

			if(!data.contains(pot)) {
				return nullptr;
			}

			all_over(index, [&](size_t idx){
				auto t = get_twig(idx);
				t.used_children = true;
				set_twig(idx, t);
			});

			auto t = get_twig(index);
			t.used = true;
			set_twig(index, t);
			return pot;
		}

		virtual bool try_reallocate(void*, size_t) {
			return false;
		}

		virtual bool deallocate(void* ptr)
		{
			if(data.contains((char*)ptr))
			{
				size_t integral_offset = reinterpret_cast<intptr_t>(ptr) - reinterpret_cast<intptr_t>(&*data.begin());
				auto index = find_used_twig(integral_offset);
				if(index == no_twig)
				{
					return false;
				}
				twig v = get_twig(index);
				v.used = false;
				v.used_children = false;
				set_twig(index, v);
				all_over(index, [&](size_t idx){
					auto l = get_twig(get_left(idx));
					auto r = get_twig(get_right(idx));

					set_twig(idx, 2*(l.used | l.used_children | r.used | r.used_children));
				});

				return true;
			}
			return false;
		}

		/**
		 * @brief Checks if anything is still allocated in
		 * 
		 * @return true if the allocator is completely empty
		 * @return false if anything is still allocated in there
		 */
		bool empty() const {
			buddy* addr = (buddy*)this;
			auto prepred = not_fn(&buddy::empty_node);
			auto pred = bind_front(prepred, addr);
			return empty_node(addr, 0) && !is_any_child(0, pred);
		}

		virtual ~buddy()
		{
			if(allocator_v.has_value())
			{
				allocator_v.value().get().deallocate(data.begin().data);
			}
		}
	};
}