psemek/libs/util/include/psemek/util/heterogeneous_container.hpp

#pragma once

#include <psemek/util/flat_list.hpp>
#include <psemek/util/at.hpp>

#include <tuple>
#include <variant>
#include <any>

namespace psemek::util
{

	namespace detail
	{

		constexpr std::size_t count_binary_digits(std::size_t n)
		{
			std::size_t result = 0;
			for (; n > 0; n >>= 1)
				++result;
			return result;
		}

		template <typename T, typename ... Types>
		struct one_of;

		template <typename T>
		struct one_of<T>
			: std::false_type
		{};

		template <typename T, typename H, typename ... Types>
		struct one_of<T, H, Types...>
			: std::bool_constant<
				std::is_same_v<T, H> ||
				one_of<T, Types...>::value
				>
		{};

		template <typename T, typename ... Types>
		static constexpr bool one_of_v = one_of<T, Types...>::value;

		template <typename T, typename ... Types>
		struct index_of;

		template <typename T>
		struct index_of<T>
		{
			// there is no good value here, but we have to break recursion
			static constexpr std::size_t value = -1;
		};

		template <typename T, typename H, typename ... Types>
		struct index_of<T, H, Types...>
		{
			static constexpr std::size_t value =
				std::is_same_v<T, H>
				? 0
				: 1 + index_of<T, Types...>::value;
		};

		template <typename T, typename ... Types>
		static constexpr std::size_t index_of_v = index_of<T, Types...>::value;

		template <typename F, typename T, std::size_t ... Is>
		void tuple_for_each_impl(F && f, T & t, std::index_sequence<Is...>)
		{
			using expander = int[];
			(void) expander { 0, (f(std::get<Is>(t), Is), 0) ... };
		}

		template <typename F, typename T>
		void tuple_for_each(F && f, T & t)
		{
			tuple_for_each_impl(std::forward<F>(f), t, std::make_index_sequence<std::tuple_size_v<T>>{});
		}

		template <typename F, typename T, typename Seq>
		struct tuple_apply_at_impl;

		template <typename F, typename T>
		struct tuple_apply_at_impl<F, T, std::index_sequence<>>
		{
			static decltype(auto) apply(F && f, T & t, std::size_t i)
			{
				throw key_error(i);
				// For proper type deduction
				return std::forward<F>(f)(std::get<0>(t));
			}
		};


		template <typename F, typename T, std::size_t I, std::size_t ... Is>
		struct tuple_apply_at_impl<F, T, std::index_sequence<I, Is...>>
		{
			static decltype(auto) apply(F && f, T & t, std::size_t i)
			{
				if (i == I)
				{
					return std::forward<F>(f)(std::get<I>(t));
				}
				return tuple_apply_at_impl<F, T, std::index_sequence<Is...>>::apply(std::forward<F>(f), t, i);
			}
		};

		template <typename F, typename T>
		decltype(auto) tuple_apply_at(F && f, T & t, std::size_t i)
		{
			if (i >= std::tuple_size_v<T>)
				throw key_error(i);

			return tuple_apply_at_impl<F, T, std::make_index_sequence<std::tuple_size_v<T>>>::apply(std::forward<F>(f), t, i);
		}

	}

	template <typename Handle, typename ... Types>
	struct heterogeneous_container
	{
		using types = std::tuple<Types...>;

		using any = std::variant<Types...>;

		using any_pointer = std::variant<Types *...>;

		using any_const_pointer = std::variant<Types const *...>;

		static constexpr std::size_t types_count = sizeof...(Types);
		static_assert(types_count > 0);

		using handle = Handle;
		static_assert(std::is_integral_v<handle> && std::is_unsigned_v<handle>);

		static std::size_t max_size()
		{
			return std::numeric_limits<handle>::max();
		}

		static std::size_t max_single_size()
		{
			return max_size() >> types_mask_bits;
		}

		template <typename T>
		std::enable_if_t<detail::one_of_v<T, Types...>, handle>
			insert(T value);

		handle insert(any value);

		void erase(handle h);

		void clear();

		any_pointer get(handle h);
		any_const_pointer get(handle h) const;

		template <typename T>
		std::enable_if_t<detail::one_of_v<T, Types...>, T const &>
			get(handle h) const;

		template <typename T>
		std::enable_if_t<detail::one_of_v<T, Types...>, T &>
			get(handle h);

		// TODO: arbitrary-arity visitors

		template <typename F>
		decltype(auto) visit(F && f, handle h1) const;

		template <typename F>
		decltype(auto) visit(F && f, handle h1, handle h2) const;

		template <typename F>
		decltype(auto) visit(F && f, handle h1);

		template <typename F>
		decltype(auto) visit(F && f, handle h1, handle h2);

		std::size_t size() const;

		bool empty() const;

	private:

		static constexpr std::size_t handle_bits = detail::count_binary_digits(static_cast<handle>(-1));

		static constexpr std::size_t types_mask_bits = detail::count_binary_digits(std::max<std::size_t>(1, types_count - 1));

		static constexpr handle type_mask = static_cast<handle>(-1) ^ (static_cast<handle>(-1) >> types_mask_bits);

		std::tuple<flat_list<Types, handle>...> storage_;

		static std::pair<std::size_t, handle> unpack_handle(handle h)
		{
			std::size_t const i = (h & type_mask) >> (handle_bits - types_mask_bits);
			handle const hh = h & (~type_mask);
			return {i, hh};
		}

		static handle pack_handle(std::size_t i, handle h)
		{
			return h | static_cast<handle>(i << (handle_bits - types_mask_bits));
		}
	};

	template <typename Handle, typename ... Types>
	template <typename T>
	std::enable_if_t<detail::one_of_v<T, Types...>, Handle>
	heterogeneous_container<Handle, Types...>::insert(T value)
	{
		static constexpr std::size_t i = detail::index_of_v<T, Types...>;
		auto h = std::get<i>(storage_).insert(std::move(value));
		return pack_handle(i, h);
	}

	template <typename Handle, typename ... Types>
	Handle heterogeneous_container<Handle, Types...>::insert(any value)
	{
		std::visit([this](auto && value){ this->insert(value); }, value);
	}

	template <typename Handle, typename ... Types>
	void heterogeneous_container<Handle, Types...>::erase(handle h)
	{
		auto u = unpack_handle(h);
		if (u.first >= types_count)
			throw exception("Broken handle");
		detail::tuple_apply_at([&u](auto & t) -> void {
			t.erase(u.second);
		}, storage_, u.first);
	}

	template <typename Handle, typename ... Types>
	void heterogeneous_container<Handle, Types...>::clear()
	{
		detail::tuple_for_each([](auto & t, auto){
			t.clear();
		}, storage_);
	}

	template <typename Handle, typename ... Types>
	typename heterogeneous_container<Handle, Types...>::any_pointer heterogeneous_container<Handle, Types...>::get(handle h)
	{
		auto u = unpack_handle(h);
		return detail::tuple_apply_at([&u](auto & t) -> any_pointer {
			return &t[u.second];
		}, storage_, u.first);
	}

	template <typename Handle, typename ... Types>
	typename heterogeneous_container<Handle, Types...>::any_const_pointer heterogeneous_container<Handle, Types...>::get(handle h) const
	{
		auto u = unpack_handle(h);
		return detail::tuple_apply_at([&u](auto & t) -> any_const_pointer {
			return &t[u.second];
		}, storage_, u.first);
	}

	template <typename Handle, typename ... Types>
	template <typename T>
	std::enable_if_t<detail::one_of_v<T, Types...>, T const &>
		heterogeneous_container<Handle, Types...>::get(handle h) const
	{
		T const * res = nullptr;
		std::pair<std::size_t, handle> const p = unpack_handle(h);
		std::size_t const i = p.first;
		handle const hh = p.second;
		detail::tuple_for_each([i, hh, &res](auto & t, std::size_t j){
			if constexpr (std::is_same_v<std::decay_t<decltype(t[hh])>, T>)
			{
				if (i != j)
					throw exception("Error in heterogeneous_container");
				res = &t[hh];
			}
		}, storage_);
		if (!res)
			throw exception("Error in heterogeneous_container");

		return *res;
	}

	template <typename Handle, typename ... Types>
	template <typename T>
	std::enable_if_t<detail::one_of_v<T, Types...>, T &>
		heterogeneous_container<Handle, Types...>::get(handle h)
	{
		T * res = nullptr;
		std::pair<std::size_t, handle> const p = unpack_handle(h);
		std::size_t const i = p.first;
		handle const hh = p.second;
		detail::tuple_for_each([i, hh, &res](auto & t, std::size_t j){
			if constexpr (std::is_same_v<std::decay_t<decltype(t[hh])>, T>)
			{
				if (i != j)
					throw exception("Error in heterogeneous_container");
				res = &t[hh];
			}
		}, storage_);
		if (!res)
			throw exception("Error in heterogeneous_container");

		return *res;
	}

	template <typename Handle, typename ... Types>
	template <typename F>
	decltype(auto) heterogeneous_container<Handle, Types...>::visit(F && f, handle h1) const
	{
		std::pair<std::size_t, handle> const p1 = unpack_handle(h1);
		std::size_t const i1 = p1.first;
		handle const hh1 = p1.second;

		return detail::tuple_apply_at([hh1, &f](auto & t1){
			return std::forward<F>(f)(t1[hh1]);
		}, storage_, i1);
	}

	template <typename Handle, typename ... Types>
	template <typename F>
	decltype(auto) heterogeneous_container<Handle, Types...>::visit(F && f, handle h1, handle h2) const
	{
		std::pair<std::size_t, handle> const p1 = unpack_handle(h1);
		std::size_t const i1 = p1.first;
		handle const hh1 = p1.second;

		std::pair<std::size_t, handle> const p2 = unpack_handle(h2);
		std::size_t const i2 = p2.first;
		handle const hh2 = p2.second;

		return detail::tuple_apply_at([i2, hh1, hh2, &f, this](auto & t1){
			return detail::tuple_apply_at([hh1, hh2, &f, &t1](auto & t2){
				return std::forward<F>(f)(t1[hh1], t2[hh2]);
			}, storage_, i2);
		}, storage_, i1);
	}

	template <typename Handle, typename ... Types>
	template <typename F>
	decltype(auto) heterogeneous_container<Handle, Types...>::visit(F && f, handle h1)
	{
		std::pair<std::size_t, handle> const p1 = unpack_handle(h1);
		std::size_t const i1 = p1.first;
		handle const hh1 = p1.second;

		return detail::tuple_apply_at([hh1, &f](auto & t1){
			return std::forward<F>(f)(t1[hh1]);
		}, storage_, i1);
	}

	template <typename Handle, typename ... Types>
	template <typename F>
	decltype(auto) heterogeneous_container<Handle, Types...>::visit(F && f, handle h1, handle h2)
	{
		std::pair<std::size_t, handle> const p1 = unpack_handle(h1);
		std::size_t const i1 = p1.first;
		handle const hh1 = p1.second;

		std::pair<std::size_t, handle> const p2 = unpack_handle(h2);
		std::size_t const i2 = p2.first;
		handle const hh2 = p2.second;

		return detail::tuple_apply_at([i2, hh1, hh2, &f, this](auto & t1){
			return detail::tuple_apply_at([hh1, hh2, &f, &t1](auto & t2){
				return std::forward<F>(f)(t1[hh1], t2[hh2]);
			}, storage_, i2);
		}, storage_, i1);
	}

	template <typename Handle, typename ... Types>
	std::size_t heterogeneous_container<Handle, Types...>::size() const
	{
		std::size_t res = 0;
		detail::tuple_for_each([&res](auto & t, auto){ res += t.size(); }, storage_);
		return res;
	}

	template <typename Handle, typename ... Types>
	bool heterogeneous_container<Handle, Types...>::empty() const
	{
		bool res = true;
		detail::tuple_for_each([&res](auto & t, auto){ res &= t.empty(); }, storage_);
		return res;
	}


}