SharedPtr.cpp

#ifndef SMART_POINTERS_H
#define SMART_POINTERS_H

#include <iostream>
#include <memory>

template <typename T>
class EnableSharedFromThis;

template <typename T>
class WeakPtr;

struct BaseControlBlock {
  size_t shared_count = 1;
  size_t weak_count = 0;

  virtual void deallocate_value() = 0;

  virtual void deallocate_myself() = 0;

  virtual ~BaseControlBlock() {}
};

template <typename U>
struct ControlBlockPtr final : BaseControlBlock {
  U* ptr;

  ControlBlockPtr(U* ptr) : ptr(ptr) {}

  void deallocate_value() override { delete ptr; }

  void deallocate_myself() override { delete this; }
};

template <typename U, typename Deleter, typename Alloc>
struct ControlBlockRegular final : BaseControlBlock {
  U* ptr;
  [[no_unique_address]] Deleter deleter;
  [[no_unique_address]] Alloc alloc;

  ControlBlockRegular(U* ptr, Deleter deleter) : ptr(ptr), deleter(deleter) {}
  ControlBlockRegular(U* ptr, Deleter deleter, Alloc alloc)
      : ptr(ptr), deleter(deleter), alloc(alloc) {}

  void deallocate_value() override { deleter(ptr); }

  void deallocate_myself() override {
    using ControlBlockAlloc = std::allocator_traits<
        Alloc>::template rebind_alloc<ControlBlockRegular>;
    ControlBlockAlloc control_block_alloc(alloc);
    this->~ControlBlockRegular();
    std::allocator_traits<ControlBlockAlloc>::deallocate(control_block_alloc,
                                                         this, 1);
  }
};

template <typename U, typename Alloc>
struct ControlBlockMakeShared final : BaseControlBlock {
  U value;
  [[no_unique_address]] Alloc alloc;

  template <typename... Args>
  ControlBlockMakeShared(Args... args) : value(std::forward<Args>(args)...) {}

  template <typename... Args>
  ControlBlockMakeShared(Alloc alloc, Args... args)
      : value(std::forward<Args>(args)...), alloc(alloc) {}

  void deallocate_value() override {
    using UAlloc = std::allocator_traits<Alloc>::template rebind_alloc<U>;
    UAlloc ualloc(alloc);
    std::allocator_traits<UAlloc>::destroy(ualloc, &value);
  }

  void deallocate_myself() override {
    using ControlBlockAlloc = std::allocator_traits<
        Alloc>::template rebind_alloc<ControlBlockMakeShared>;
    ControlBlockAlloc control_block_alloc(alloc);
    std::allocator_traits<ControlBlockAlloc>::deallocate(control_block_alloc,
                                                         this, 1);
    alloc.~Alloc();
  }
};

template <typename T>
class SharedPtr {
 private:
  T* ptr = nullptr;
  BaseControlBlock* cb = nullptr;

  void decrease_shared_count() {
    if (cb == nullptr) {
      return;
    }
    if (cb->shared_count == 1) {
      cb->deallocate_value();
      if (cb->weak_count == 0) {
        cb->deallocate_myself();
        return;
      }
    }
    --cb->shared_count;
  }

  SharedPtr(T* ptr, BaseControlBlock* cb) : ptr(ptr), cb(cb) {
    if (cb != nullptr) {
      ++cb->shared_count;
    }
  }

  template <typename U, typename... Args>
  friend SharedPtr<U> makeShared(Args... args);

  template <typename U, typename Alloc, typename... Args>
  friend SharedPtr<U> allocateShared(Alloc alloc, Args... args);

  template <typename U>
  friend class WeakPtr;

  template <typename U>
  friend class SharedPtr;

  template <typename U>
  friend class EnableSharedFromThis;

  template <typename U, typename... Args>
  friend SharedPtr<U> makeShared(Args... args);

  template<typename Alloc>
  SharedPtr(ControlBlockMakeShared<T, Alloc>* cb)
      : ptr(&cb->value), cb(cb) {
    if constexpr (std::is_base_of_v<EnableSharedFromThis<T>, T>) {
      ptr->wptr = *this;
    }
  }

  SharedPtr(const WeakPtr<T>& wptr) : ptr(wptr.ptr), cb(wptr.cb) {
    if (wptr.expired()) {
      throw std::bad_weak_ptr();
    }
    ++cb->shared_count;
  }

 public:
  SharedPtr() = default;

  template <typename U>
  SharedPtr(U* ptr) : ptr(ptr), cb(new ControlBlockPtr(ptr)) {
    if constexpr (std::is_base_of_v<EnableSharedFromThis<T>, T>) {
      ptr->wptr = *this;
    }
  }

  template <typename U>
  SharedPtr(const SharedPtr<U>& other) : ptr(other.ptr), cb(other.cb) {
    if (cb != nullptr) {
      ++cb->shared_count;
    }
  }

  SharedPtr(const SharedPtr& other) : ptr(other.ptr), cb(other.cb) {
    if (cb != nullptr) {
      ++cb->shared_count;
    }
  }

  template <typename U>
  SharedPtr& operator=(const SharedPtr<U>& other) {
    if (cb != other.cb) {
      decrease_shared_count();
      ptr = other.ptr;
      cb = other.cb;
      if (cb != nullptr) {
        ++cb->shared_count;
      }
    }
    return *this;
  }

  SharedPtr& operator=(const SharedPtr& other) { return operator= <T>(other); }

  template <typename U>
  SharedPtr(SharedPtr<U>&& other) noexcept : ptr(other.ptr), cb(other.cb) {
    other.ptr = nullptr;
    other.cb = nullptr;
  }

  SharedPtr(SharedPtr&& other) noexcept : ptr(other.ptr), cb(other.cb) {
    other.ptr = nullptr;
    other.cb = nullptr;
  }

  template <typename U>
  SharedPtr& operator=(SharedPtr<U>&& other) noexcept {
    decrease_shared_count();
    ptr = other.ptr;
    cb = other.cb;
    other.ptr = nullptr;
    other.cb = nullptr;
    return *this;
  }

  SharedPtr& operator=(SharedPtr&& other) noexcept {
    if (this != &other) {
      decrease_shared_count();
      ptr = other.ptr;
      cb = other.cb;
      other.ptr = nullptr;
      other.cb = nullptr;
    }
    return *this;
  }

  template <typename U, typename Deleter>
  SharedPtr(U* ptr, Deleter deleter)
      : ptr(ptr),
        cb(new ControlBlockRegular<U, Deleter, std::allocator<U>>(ptr,
                                                                  deleter)) {}

  template <typename U, typename Deleter, typename Alloc>
  SharedPtr(U* ptr, Deleter deleter, Alloc alloc) : ptr(ptr) {
    using ControlBlockAlloc = std::allocator_traits<
        Alloc>::template rebind_alloc<ControlBlockRegular<U, Deleter, Alloc>>;
    ControlBlockAlloc control_block_alloc(alloc);
    auto regular_cb = std::allocator_traits<ControlBlockAlloc>::allocate(
        control_block_alloc, 1);
    cb = regular_cb;
    new (regular_cb)
        ControlBlockRegular<U, Deleter, Alloc>(ptr, deleter, alloc);
  }

  ~SharedPtr() { decrease_shared_count(); }

  size_t use_count() const { return (cb != nullptr ? cb->shared_count : 0); }

  T* get() const { return ptr; }

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

  void swap(SharedPtr& other) noexcept {
    std::swap(ptr, other.ptr);
    std::swap(cb, other.cb);
  }

  void reset() {
    decrease_shared_count();
    ptr = nullptr;
    cb = nullptr;
  }

  template <typename U>
  void reset(U* new_ptr) {
    decrease_shared_count();
    ptr = new_ptr;
    cb = new ControlBlockPtr(new_ptr);
  }

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

template <typename T, typename... Args>
SharedPtr<T> makeShared(Args... args) {
  using ControlBlockType = ControlBlockMakeShared<T, std::allocator<T>>;
  auto cb = new ControlBlockType(std::forward<Args>(args)...);
  SharedPtr<T> res(cb);
  return res;
}

template <typename T, typename Alloc = std::allocator<T>, typename... Args>
SharedPtr<T> allocateShared(Alloc alloc, Args... args) {
  using ControlBlockType = ControlBlockMakeShared<T, Alloc>;
  using ControlBlockAlloc =
      std::allocator_traits<Alloc>::template rebind_alloc<ControlBlockType>;
  ControlBlockAlloc control_block_alloc(alloc);
  auto cb = std::allocator_traits<ControlBlockAlloc>::allocate(
      control_block_alloc, 1);
  std::allocator_traits<ControlBlockAlloc>::construct(
      control_block_alloc, cb, alloc, std::forward<Args>(args)...);
  return SharedPtr<T>(cb);
}

template <typename T>
class WeakPtr {
 private:
  T* ptr = nullptr;
  BaseControlBlock* cb = nullptr;

  void decrease_weak_count() {
    if (cb != nullptr && --cb->weak_count == 0 && cb->shared_count == 0) {
      cb->deallocate_myself();
    }
  }

  template <typename U>
  friend class WeakPtr;

  template <typename U>
  friend class SharedPtr;

 public:
  WeakPtr() {}

  template <typename U>
  WeakPtr(const SharedPtr<U>& sptr) : ptr(sptr.ptr), cb(sptr.cb) {
    if (cb != nullptr) {
      ++cb->weak_count;
    }
  }

  template <typename U>
  WeakPtr(const WeakPtr<U>& other) : ptr(other.ptr), cb(other.cb) {
    if (cb != nullptr) {
      ++cb->weak_count;
    }
  }

  WeakPtr(const WeakPtr& other) : ptr(other.ptr), cb(other.cb) {
    if (cb != nullptr) {
      ++cb->weak_count;
    }
  }

  template <typename U>
  WeakPtr& operator=(const WeakPtr<U>& other) {
    if (cb != other.cb) {
      decrease_weak_count();
      ptr = other.ptr;
      cb = other.cb;
      if (cb != nullptr) {
        ++cb->weak_count;
      }
    }
    return *this;
  }

  WeakPtr& operator=(const WeakPtr& other) { return operator= <T>(other); }

  template <typename U>
  WeakPtr(WeakPtr<U>&& other) noexcept : ptr(other.ptr), cb(other.cb) {
    other.ptr = nullptr;
    other.cb = nullptr;
  }

  WeakPtr(WeakPtr&& other) noexcept : ptr(other.ptr), cb(other.cb) {
    other.ptr = nullptr;
    other.cb = nullptr;
  }

  template <typename U>
  WeakPtr& operator=(WeakPtr<U>&& other) noexcept {
    if (this != &other) {
      decrease_weak_count();
      ptr = other.ptr;
      cb = other.cb;
      other.cb = nullptr;
    }
    return *this;
  }

  WeakPtr& operator=(WeakPtr&& other) noexcept {
    return operator= <T>(std::move(other));
  }

  ~WeakPtr() { decrease_weak_count(); }

  bool expired() const { return cb == nullptr || cb->shared_count == 0; }

  size_t use_count() const { return (cb == nullptr ? 0 : cb->shared_count); }

  SharedPtr<T> lock() const { return SharedPtr<T>(ptr, cb); }
};

template <typename T>
class EnableSharedFromThis {
 public:
  WeakPtr<T> wptr;
  SharedPtr<T> shared_from_this() { return wptr; }
  EnableSharedFromThis() {}
};

#endif  //SMART_POINTERS_H