weak_ptr.h

// Copyright 2025 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.
#pragma once
 
#include <cstddef>
#include <utility>
 
#include "pw_allocator/internal/control_block.h"
#include "pw_allocator/internal/managed_ptr.h"
#include "pw_allocator/shared_ptr.h"
 
namespace pw {
 
template <typename T>
class WeakPtr final : public ::pw::allocator::internal::WeakManagedPtr<T> {
 private:
  using Base = ::pw::allocator::internal::WeakManagedPtr<T>;
  using ControlBlock = ::pw::allocator::internal::ControlBlock;
 
 public:
  using element_type = typename Base::element_type;
 
  constexpr WeakPtr() noexcept = default;
 
  constexpr WeakPtr(std::nullptr_t) noexcept : WeakPtr() {}
 
  WeakPtr(const WeakPtr& other) noexcept { *this = other; }
 
  template <typename U>
  WeakPtr(const WeakPtr<U>& other) noexcept {
    *this = other;
  }
 
  template <typename U>
  WeakPtr(const SharedPtr<U>& other) noexcept {
    *this = other;
  }
 
  template <typename U>
  WeakPtr(WeakPtr<U>&& other) noexcept {
    *this = other;
  }
 
  ~WeakPtr() { reset(); }
 
  constexpr WeakPtr& operator=(const WeakPtr& other) noexcept {
    operator= <T>(other);
    return *this;
  }
 
  template <typename U>
  WeakPtr& operator=(const WeakPtr<U>& other) noexcept;
 
  template <typename U>
  WeakPtr& operator=(const SharedPtr<U>& other) noexcept;
 
  template <typename U>
  WeakPtr& operator=(WeakPtr<U>&& other) noexcept;
 
  void reset() noexcept;
 
  void swap(WeakPtr& other) noexcept {
    std::swap(control_block_, other.control_block_);
  }
 
  uint32_t use_count() const noexcept {
    return control_block_ == nullptr ? 0 : control_block_->num_shared();
  }
 
  bool expired() const noexcept { return use_count() == 0; }
 
  SharedPtr<T> Lock() const noexcept;
 
  template <typename PtrType>
  bool owner_before(const PtrType& other) const noexcept {
    return control_block_ < other.control_block();
  }
 
 private:
  // Allow WeakPtr<T> to access WeakPtr<U> and vice versa.
  template <typename>
  friend class WeakPtr;
 
  ControlBlock* control_block() const { return control_block_; }
 
  ControlBlock* control_block_ = nullptr;
};
 
// Template implementations.
 
template <typename T>
template <typename U>
WeakPtr<T>& WeakPtr<T>::operator=(const WeakPtr<U>& other) noexcept {
  Base::template CheckAssignable<U>();
  control_block_ = other.control_block_;
  if (control_block_ != nullptr) {
    control_block_->IncrementWeak();
  }
  return *this;
}
 
template <typename T>
template <typename U>
WeakPtr<T>& WeakPtr<T>::operator=(const SharedPtr<U>& other) noexcept {
  Base::template CheckAssignable<U>();
  control_block_ = other.control_block_;
  if (control_block_ != nullptr) {
    control_block_->IncrementWeak();
  }
  return *this;
}
 
template <typename T>
template <typename U>
WeakPtr<T>& WeakPtr<T>::operator=(WeakPtr<U>&& other) noexcept {
  Base::template CheckAssignable<U>();
  control_block_ = other.control_block_;
  other.control_block_ = nullptr;
  return *this;
}
 
template <typename T>
void WeakPtr<T>::reset() noexcept {
  if (control_block_ == nullptr ||
      control_block_->DecrementWeak() != ControlBlock::Action::kFree) {
    return;
  }
  Allocator* allocator = control_block_->allocator();
  Base::Deallocate(allocator, control_block_);
  control_block_ = nullptr;
}
 
template <typename T>
SharedPtr<T> WeakPtr<T>::Lock() const noexcept {
  if (control_block_ == nullptr || !control_block_->IncrementShared()) {
    return SharedPtr<T>();
  }
  void* data = control_block_->data();
  auto* t = std::launder(reinterpret_cast<element_type*>(data));
  return SharedPtr<T>(t, control_block_);
}
 
}  // namespace pw