// 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.

#include "pw_allocator/internal/control_block.h"

#include "lib/stdcompat/bit.h"
#include "pw_allocator/allocator.h"
#include "pw_allocator/capability.h"
#include "pw_allocator/deallocator.h"
#include "pw_allocator/hardening.h"
#include "pw_assert/check.h"
#include "pw_bytes/alignment.h"

namespace pw::allocator::internal {

const ControlBlockHandle& ControlBlockHandle::GetInstance_DO_NOT_USE() {
  static const ControlBlockHandle kControlBlockHandle;
  return kControlBlockHandle;
}

ControlBlock* ControlBlock::Create(const ControlBlockHandle&,
                                   pw::Allocator* allocator,
                                   Layout layout) {
  size_t size = layout.size();
  layout = layout.Extend(AlignUp(sizeof(ControlBlock), layout.alignment()));
  void* ptr = allocator->Allocate(layout);
  if (ptr == nullptr) {
    return nullptr;
  }
  auto addr = cpp20::bit_cast<uintptr_t>(ptr);
  addr = AlignUp(addr + sizeof(ControlBlock), layout.alignment());
  auto* data = cpp20::bit_cast<std::byte*>(addr);
  return new (ptr) ControlBlock(allocator, data, size, true);
}

ControlBlock* ControlBlock::Create(const ControlBlockHandle&,
                                   Deallocator* deallocator,
                                   void* data,
                                   size_t size) {
  if (!deallocator->HasCapability(
          allocator::Capability::kCanAllocateArbitraryLayout)) {
    return nullptr;
  }
  auto* allocator = static_cast<pw::Allocator*>(deallocator);
  void* ptr = allocator->Allocate(Layout::Of<ControlBlock>());
  if (ptr == nullptr) {
    return nullptr;
  }
  return new (ptr) ControlBlock(allocator, data, size, false);
}

ControlBlock::ControlBlock(pw::Allocator* allocator,
                           void* data,
                           size_t size,
                           bool coallocated)
    : allocator_(allocator),
      data_(data),
      coallocated_and_size_((size << 1) | (coallocated ? 1 : 0)),
      num_weak_and_shared_(Pack(1, 1)) {
  PW_CHECK_UINT_EQ(size, GetSize());
}

ControlBlock::~ControlBlock() {
  if (!IsCoallocated()) {
    allocator_->Deallocate(data_);
  }
}

int32_t ControlBlock::num_shared() const noexcept {
  return UnpackShared(num_weak_and_shared_.load(std::memory_order_relaxed));
}

bool ControlBlock::IncrementShared() {
  uint32_t num = num_weak_and_shared_.load(std::memory_order_relaxed);
  while (true) {
    uint16_t num_weak = UnpackWeak(num);
    uint16_t num_shared = UnpackShared(num);
    if constexpr (Hardening::kIncludesDebugChecks) {
      PW_CHECK_UINT_GE(num_weak, num_shared);
    }
    if (num_shared == 0 || num_shared == 0xFFFF) {
      return false;
    }
    uint32_t packed = Pack(num_weak + 1, num_shared + 1);
    if (num_weak_and_shared_.compare_exchange_weak(
            num, packed, std::memory_order_relaxed)) {
      return true;
    }
  }
}

bool ControlBlock::IncrementWeak() {
  uint32_t num = num_weak_and_shared_.load(std::memory_order_relaxed);
  while (true) {
    uint16_t num_weak = UnpackWeak(num);
    uint16_t num_shared = UnpackShared(num);
    if constexpr (Hardening::kIncludesDebugChecks) {
      PW_CHECK_UINT_GE(num_weak, num_shared);
    }
    if (num_weak == 0 || num_weak == 0xFFFF) {
      return false;
    }
    uint32_t packed = Pack(num_weak + 1, num_shared);
    if (num_weak_and_shared_.compare_exchange_weak(
            num, packed, std::memory_order_relaxed)) {
      return true;
    }
  }
}

ControlBlock::Action ControlBlock::DecrementShared() {
  uint32_t prev =
      num_weak_and_shared_.fetch_sub(Pack(1, 1), std::memory_order_acq_rel);
  uint16_t prev_weak = UnpackWeak(prev);
  uint16_t prev_shared = UnpackShared(prev);
  if constexpr (Hardening::kIncludesDebugChecks) {
    PW_CHECK_UINT_NE(prev_weak, 0);
    PW_CHECK_UINT_NE(prev_shared, 0);
    PW_CHECK_UINT_GE(prev_weak, prev_shared);
  }
  if (prev_weak == 1 && prev_shared == 1) {
    return ControlBlock::Action::kFree;
  }
  if (prev_shared == 1) {
    return ControlBlock::Action::kExpire;
  }
  return ControlBlock::Action::kNone;
}

ControlBlock::Action ControlBlock::DecrementWeak() {
  uint32_t prev =
      num_weak_and_shared_.fetch_sub(Pack(1, 0), std::memory_order_acq_rel);
  uint16_t prev_weak = UnpackWeak(prev);
  if constexpr (Hardening::kIncludesDebugChecks) {
    uint16_t prev_shared = UnpackShared(prev);
    PW_CHECK_UINT_NE(prev_weak, 0);
    PW_CHECK_UINT_GE(prev_weak, prev_shared);
  }
  if (prev_weak == 1) {
    return ControlBlock::Action::kFree;
  }
  return ControlBlock::Action::kNone;
}

}  // namespace pw::allocator::internal