// -*- C++ -*-
//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

// Modifications Copyright (c) 2025 Advanced Micro Devices, Inc.
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#ifndef __LIBHIPTHREADS_MUTEX___
#define __LIBHIPTHREADS_MUTEX___

/*
    mutex synopsis

namespace std
{

class mutex
{
public:
     constexpr mutex() noexcept;
     ~mutex();

    mutex(const mutex&) = delete;
    mutex& operator=(const mutex&) = delete;

    void lock();
    bool try_lock();
    void unlock();

    typedef pthread_mutex_t* native_handle_type;
    native_handle_type native_handle();
};

class recursive_mutex
{
public:
     recursive_mutex();
     ~recursive_mutex();

    recursive_mutex(const recursive_mutex&) = delete;
    recursive_mutex& operator=(const recursive_mutex&) = delete;

    void lock();
    bool try_lock() noexcept;
    void unlock();

    typedef pthread_mutex_t* native_handle_type;
    native_handle_type native_handle();
};

class timed_mutex
{
public:
     timed_mutex();
     ~timed_mutex();

    timed_mutex(const timed_mutex&) = delete;
    timed_mutex& operator=(const timed_mutex&) = delete;

    void lock();
    bool try_lock();
    template <class Rep, class Period>
        bool try_lock_for(const chrono::duration<Rep, Period>& rel_time);
    template <class Clock, class Duration>
        bool try_lock_until(const chrono::time_point<Clock, Duration>& abs_time);
    void unlock();
};

class recursive_timed_mutex
{
public:
     recursive_timed_mutex();
     ~recursive_timed_mutex();

    recursive_timed_mutex(const recursive_timed_mutex&) = delete;
    recursive_timed_mutex& operator=(const recursive_timed_mutex&) = delete;

    void lock();
    bool try_lock() noexcept;
    template <class Rep, class Period>
        bool try_lock_for(const chrono::duration<Rep, Period>& rel_time);
    template <class Clock, class Duration>
        bool try_lock_until(const chrono::time_point<Clock, Duration>& abs_time);
    void unlock();
};

struct defer_lock_t { explicit defer_lock_t() = default; };
struct try_to_lock_t { explicit try_to_lock_t() = default; };
struct adopt_lock_t { explicit adopt_lock_t() = default; };

inline constexpr defer_lock_t  defer_lock{};
inline constexpr try_to_lock_t try_to_lock{};
inline constexpr adopt_lock_t  adopt_lock{};

template <class Mutex>
class lock_guard
{
public:
    typedef Mutex mutex_type;

    explicit lock_guard(mutex_type& m);
    lock_guard(mutex_type& m, adopt_lock_t);
    ~lock_guard();

    lock_guard(lock_guard const&) = delete;
    lock_guard& operator=(lock_guard const&) = delete;
};

template <class... MutexTypes>
class scoped_lock // C++17
{
public:
    using mutex_type = Mutex;  // Only if sizeof...(MutexTypes) == 1

    explicit scoped_lock(MutexTypes&... m);
    scoped_lock(adopt_lock_t, MutexTypes&... m);
    ~scoped_lock();
    scoped_lock(scoped_lock const&) = delete;
    scoped_lock& operator=(scoped_lock const&) = delete;
private:
    tuple<MutexTypes&...> pm; // exposition only
};

template <class Mutex>
class unique_lock
{
public:
    typedef Mutex mutex_type;
    unique_lock() noexcept;
    explicit unique_lock(mutex_type& m);
    unique_lock(mutex_type& m, defer_lock_t) noexcept;
    unique_lock(mutex_type& m, try_to_lock_t);
    unique_lock(mutex_type& m, adopt_lock_t);
    template <class Clock, class Duration>
        unique_lock(mutex_type& m, const chrono::time_point<Clock, Duration>& abs_time);
    template <class Rep, class Period>
        unique_lock(mutex_type& m, const chrono::duration<Rep, Period>& rel_time);
    ~unique_lock();

    unique_lock(unique_lock const&) = delete;
    unique_lock& operator=(unique_lock const&) = delete;

    unique_lock(unique_lock&& u) noexcept;
    unique_lock& operator=(unique_lock&& u) noexcept;

    void lock();
    bool try_lock();

    template <class Rep, class Period>
        bool try_lock_for(const chrono::duration<Rep, Period>& rel_time);
    template <class Clock, class Duration>
        bool try_lock_until(const chrono::time_point<Clock, Duration>& abs_time);

    void unlock();

    void swap(unique_lock& u) noexcept;
    mutex_type* release() noexcept;

    bool owns_lock() const noexcept;
    explicit operator bool () const noexcept;
    mutex_type* mutex() const noexcept;
};

template <class Mutex>
  void swap(unique_lock<Mutex>& x, unique_lock<Mutex>& y) noexcept;

template <class L1, class L2, class... L3>
  int try_lock(L1&, L2&, L3&...);
template <class L1, class L2, class... L3>
  void lock(L1&, L2&, L3&...);

struct once_flag
{
    constexpr once_flag() noexcept;

    once_flag(const once_flag&) = delete;
    once_flag& operator=(const once_flag&) = delete;
};

template<class Callable, class ...Args>
  void call_once(once_flag& flag, Callable&& func, Args&&... args);

}  // std

*/

#include "hip/thread_config"

#include <cstdint>
#include <mutex>

// #include "hip/__condition_variable/condition_variable.h"
// #include "hip/__memory/shared_ptr.h"
#include "hip/__mutex/lock_guard.h"
#include "hip/__mutex/spin_mutex.h"
#include "hip/__mutex/unique_lock.h"

/**
 * @file
 * @brief Multi-mutex locking helpers (device-side analogues of `std::try_lock` / `std::lock`).
 * @ingroup mutex
 *
 * Provides overloads of:
 *  - try_lock(m1, m2[, ...])  : Attempts to acquire all mutexes without indefinite blocking.
 *       Returns -1 on success (all locked) or the (0-based) index of the first mutex
 *       that could not be locked. On failure, all successfully locked earlier mutexes
 *       are unlocked before return.
 *  - lock(m1, m2[, ...])      : Acquires all mutexes (spinning) using a deadlock‑avoidance
 *       strategy that alternates starting points until success.
 *
 * Implementation strategy (adapted from libc++):
 *  - try_lock variants use RAII (unique_lock) for early resources and release on failure.
 *  - lock variants loop, attempting acquisition order permutations until all succeed.
 *
 * Notes:
 *  - Intended for short critical sections (spin semantics).
 *  - Mutex types must model BasicLockable (lock(), try_lock(), unlock()).
 *  - Overloads with >2 locks are variadic templates; helpers (__lock_first / __unlock)
 *    are internal.
 */

namespace cuda {

/**
 * @brief try_lock for two mutex-like objects.
 *
 * Attempts to lock @p __l0 (non-blocking) then @p __l1 (non-blocking).
 * On success both are locked and -1 is returned. On failure:
 *  - Return 0 if @p __l0 could not be locked.
 *  - Return 1 if @p __l0 locked but @p __l1 failed (and @p __l0 is released).
 *
 * @return int Index of the first mutex that failed, or -1 if all locked.
 */
template <class _L0, class _L1>
__device__ _LIBHIPTHREADS_HIDE_FROM_ABI int
try_lock(_L0& __l0, _L1& __l1)
{
    unique_lock<_L0> __u0(__l0, ::std::try_to_lock);
    if (__u0.owns_lock())
    {
        if (__l1.try_lock())
        {
            __u0.release();
            return -1;
        }
        else
            return 1;
    }
    return 0;
}

#ifndef _LIBHIPTHREADS_CXX03_LANG

/**
 * @brief try_lock for N >= 3 mutex-like objects.
 *
 * Attempts to lock the first mutex then recursively tries the remainder.
 * On failure the already-locked subset is released (strong exception / failure safety).
 *
 * @return int Index (0-based) of first mutex that failed, or -1 if all locked.
 */
template <class _L0, class _L1, class _L2, class... _L3>
__device__ _LIBHIPTHREADS_HIDE_FROM_ABI int
try_lock(_L0& __l0, _L1& __l1, _L2& __l2, _L3&... __l3)
{
    int __r = 0;
    unique_lock<_L0> __u0(__l0, ::std::try_to_lock);
    if (__u0.owns_lock())
    {
        __r = hip::try_lock(__l1, __l2, __l3...);
        if (__r == -1)
            __u0.release();
        else
            ++__r;
    }
    return __r;
}

#endif // _LIBHIPTHREADS_CXX03_LANG

/**
 * @brief lock for two mutex-like objects (deadlock avoidance).
 *
 * Repeatedly attempts to acquire both without holding them simultaneously on failure,
 * alternating acquisition order until both succeed.
 */
template <class _L0, class _L1>
__device__ _LIBHIPTHREADS_HIDE_FROM_ABI void
lock(_L0& __l0, _L1& __l1)
{
    while (true)
    {
        {
            unique_lock<_L0> __u0(__l0);
            if (__l1.try_lock())
            {
                __u0.release();
                break;
            }
        }
        // __libcpp_thread_yield();
        {
            unique_lock<_L1> __u1(__l1);
            if (__l0.try_lock())
            {
                __u1.release();
                break;
            }
        }
        // __libcpp_thread_yield();
    }
}

#ifndef _LIBHIPTHREADS_CXX03_LANG

/**
 * @internal
 * @brief Recursive helper that acquires all mutexes using a rotating start index.
 * @param __i Starting index permutation.
 */
template <class _L0, class _L1, class _L2, class ..._L3>
__device__ void
__lock_first(int __i, _L0& __l0, _L1& __l1, _L2& __l2, _L3& ...__l3)
{
    while (true)
    {
        switch (__i)
        {
        case 0:
            {
                unique_lock<_L0> __u0(__l0);
                __i = hip::try_lock(__l1, __l2, __l3...);
                if (__i == -1)
                {
                    __u0.release();
                    return;
                }
            }
            ++__i;
            // __libcpp_thread_yield();
            break;
        case 1:
            {
                unique_lock<_L1> __u1(__l1);
                __i = hip::try_lock(__l2, __l3..., __l0);
                if (__i == -1)
                {
                    __u1.release();
                    return;
                }
            }
            if (__i == sizeof...(_L3) + 1)
                __i = 0;
            else
                __i += 2;
            // __libcpp_thread_yield();
            break;
        default:
            hip::__lock_first(__i - 2, __l2, __l3..., __l0, __l1);
            return;
        }
    }
}

/**
 * @brief lock for N >= 3 mutex-like objects.
 *
 * Employs permutation / rotation strategy to avoid classic deadlock when multiple
 * threads attempt to lock the same set in different orders.
 *
 * Spins until all locks are held.
 */
template <class _L0, class _L1, class _L2, class ..._L3>
__device__ inline _LIBHIPTHREADS_INLINE_VISIBILITY
void
lock(_L0& __l0, _L1& __l1, _L2& __l2, _L3& ...__l3)
{
    hip::__lock_first(0, __l0, __l1, __l2, __l3...);
}

/**
 * @internal
 * @brief Unlock helper (single).
 */
template <class _L0>
__device__ inline _LIBHIPTHREADS_INLINE_VISIBILITY
void __unlock(_L0& __l0) {
    __l0.unlock();
}

/**
 * @internal
 * @brief Unlock helper (two).
 */
template <class _L0, class _L1>
__device__ inline _LIBHIPTHREADS_INLINE_VISIBILITY
void __unlock(_L0& __l0, _L1& __l1) {
    __l0.unlock();
    __l1.unlock();
}

/**
 * @internal
 * @brief Variadic unlock helper (N >= 3).
 */
template <class _L0, class _L1, class _L2, class ..._L3>
__device__ inline _LIBHIPTHREADS_INLINE_VISIBILITY
void __unlock(_L0& __l0, _L1& __l1, _L2& __l2, _L3&... __l3) {
    __l0.unlock();
    __l1.unlock();
    hip::__unlock(__l2, __l3...);
}

#endif // _LIBHIPTHREADS_CXX03_LANG

} // namespace cuda

#endif // __LIBHIPTHREADS_MUTEX___