#pragma once

#include <cinttypes>
#include <atomic>

/*
 * Code taken Daniel Anderson's Talk "When Lock-Free Still Isn't Enough" Slide 28:
 * https://github.com/CppCon/CppCon2024/blob/main/Presentations/When_Lock-Free_Still_Isn't_Enough.pdf
 */

struct reference_counter
{
	using size_type = std::uint64_t;

	static constexpr size_type bit_count = sizeof(size_type) * 8;
	static constexpr size_type was_zero_flag = size_type{ 1 } << (bit_count - 1);
	static constexpr size_type has_helped_flag = size_type{ 1 } << (bit_count - 2);

	bool increment()
	{
		const auto prev_count = m_counter.fetch_add(1);
		return (prev_count & was_zero_flag) == 0;
	}

	bool decrement()
	{
		if (m_counter.fetch_sub(1) == 1)
		{
			auto e = size_type{};
			if (m_counter.compare_exchange_strong(e, was_zero_flag))
			{
				return true;
			}

			if ((e & has_helped_flag) and (m_counter.exchange(was_zero_flag) & has_helped_flag))
			{
				return true;
			}
		}
		return false;
	}

	size_type count()
	{
		auto value = m_counter.load();
		if (value == 0 and m_counter.compare_exchange_strong(value, was_zero_flag | has_helped_flag))
		{
			return 0;
		}
		return value & was_zero_flag ? 0 : value;
	}

private:
	std::atomic<size_type> m_counter{ 1 };
};