#include <atomic>
#include <string>
#include <thread>
#include <cassert>
#include <iostream>
using namespace std;

/* SharedPtr changes
- atomic<int>* count
- no count if ptr is nullptr -> need null checks
    - we did this so move could be noexcept. in general move should be noexcept,
      and just in case thread uses move_ifnoexcept
- detach() factored out & release() so we dont have to assign to empty ptr
- added move so that thread doesnt hold a copy after passing it to functor
- relaxed inc and get_count() because no publishing

*/

template <typename T>
class SharedPtr {
public:
    explicit SharedPtr(T* p = nullptr)
        : ptr{p}, count{p ? new std::atomic<int>{1} : nullptr} {}
    ~SharedPtr() { detach(); }

    // Copy constructor and assignment
    SharedPtr(const SharedPtr<T>& sp) : ptr(sp.ptr), count(sp.count) {
        if (count) count->fetch_add(1, std::memory_order_relaxed);
    }
    SharedPtr<T>& operator=(const SharedPtr<T>& sp) {
        if (this != &sp) {
            // first, detach.
            detach();
            // attach to the new object.
            ptr = sp.ptr;
            count = sp.count;
            if (count) count->fetch_add(1, std::memory_order_relaxed);
        }
        return *this;
    }

    // Move constructor and assignment
    SharedPtr(SharedPtr<T>&& sp) noexcept : ptr(sp.ptr), count(sp.count) {
        sp.ptr = nullptr;
        sp.count = nullptr;
    }
    SharedPtr<T>& operator=(SharedPtr<T>&& sp) noexcept {
        if (this != &sp) {
            // first, detach.
            detach();
            // attach to the new object.
            ptr = sp.ptr;
            count = sp.count;
            sp.ptr = nullptr;
            sp.count = nullptr;
        }
        return *this;
    }

    void release() {
        detach();
        ptr = nullptr;
        count = nullptr;
    }

    T& operator*() const { return *ptr; }
    T* operator->() const { return ptr; }

    operator void*() const { return ptr; } // enables "if (sp) ..."
    T* get_ptr() const { return ptr; } // access to underlying ptr

    // access to reference count
    int get_count() const { 
        return count ? count->load(std::memory_order_relaxed) : 0;
    }

private:
    void detach() {
        // if (count && count->fetch_sub(1, std::memory_order_acq_rel) == 1) {
        if (count && count->fetch_sub(1, std::memory_order_relaxed) == 1) {
            delete ptr;
            delete count;
        }
    }

    T* ptr;                   // the underlying pointer
    std::atomic<int>* count;  // the reference count
};

struct X {
    int x;
    
    X() : x{0} {}
    // X should be mutated before destruction
    ~X() { assert(x == 42); }

    void mutate() { x = 42; }

    // Copy and move not needed for this demo 
    X(const X&) = delete;
    X& operator=(const X&) = delete;
    X(X&&) = delete;
    X& operator=(X&&) = delete;
};

int main() {
    for (int i = 0; i < 100'000'000; ++i) {
        SharedPtr<X> sp { new X{} };

        std::jthread t1(
            [](SharedPtr<X> sp1) {
                assert(sp1.get_count() <= 2);

                sp1->mutate();
                sp1.release();
            },
            sp);

        std::jthread t2(
            [](SharedPtr<X> sp2) {
                assert(sp2.get_count() <= 2);

                sp2.release();
            },
            std::move(sp));
    }
}
