block_allocator.h

// Copyright 2024 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 <optional>
 
#include "pw_allocator/allocator.h"
#include "pw_allocator/block/basic.h"
#include "pw_allocator/block/iterable.h"
#include "pw_allocator/block/poisonable.h"
#include "pw_allocator/block/result.h"
#include "pw_allocator/block/with_layout.h"
#include "pw_allocator/capability.h"
#include "pw_allocator/config.h"
#include "pw_allocator/fragmentation.h"
#include "pw_allocator/hardening.h"
#include "pw_assert/assert.h"
#include "pw_bytes/span.h"
#include "pw_result/result.h"
#include "pw_status/status.h"
 
namespace pw::allocator {
namespace internal {
 
class GenericBlockAllocator : public pw::Allocator {
 public:
  // Not copyable or movable.
  GenericBlockAllocator(const GenericBlockAllocator&) = delete;
  GenericBlockAllocator& operator=(const GenericBlockAllocator&) = delete;
  GenericBlockAllocator(GenericBlockAllocator&&) = delete;
  GenericBlockAllocator& operator=(GenericBlockAllocator&&) = delete;
 
 protected:
  template <typename BlockType>
  static constexpr Capabilities GetCapabilities();
 
  constexpr explicit GenericBlockAllocator(Capabilities capabilities)
      : pw::Allocator(capabilities) {}
 
  [[noreturn]] static void CrashOnAllocated(const void* allocated);
 
  [[noreturn]] static void CrashOnOutOfRange(const void* freed);
 
  [[noreturn]] static void CrashOnDoubleFree(const void* freed);
};
 
}  // namespace internal
 
namespace test {
 
// Forward declaration for friending.
template <typename, size_t>
class BlockAllocatorTest;
 
}  // namespace test
 
 
template <typename BlockType_>
class BlockAllocator : public internal::GenericBlockAllocator {
 private:
  using Base = internal::GenericBlockAllocator;
 
 public:
  using BlockType = BlockType_;
  using Range = typename BlockType::Range;
 
  static constexpr Capabilities kCapabilities =
      Base::GetCapabilities<BlockType>();
  static constexpr size_t kPoisonInterval = PW_ALLOCATOR_BLOCK_POISON_INTERVAL;
 
  ~BlockAllocator() override;
 
  Range blocks() const { return Range(first_); }
 
  void Init(ByteSpan region);
 
  size_t GetMaxAllocatable() { return DoGetMaxAllocatable(); }
 
 protected:
  constexpr explicit BlockAllocator() : Base(kCapabilities) {}
 
  void* DoAllocate(Layout layout) override;
 
  void DoDeallocate(void* ptr) override;
 
  bool DoResize(void* ptr, size_t new_size) override;
 
  size_t DoGetAllocated() const override { return allocated_; }
 
  std::optional<Fragmentation> DoMeasureFragmentation() const override;
 
  Result<Layout> DoGetInfo(InfoType info_type, const void* ptr) const override;
 
  void Init(BlockType* begin);
 
  template <typename Ptr>
  internal::copy_const_ptr_t<Ptr, BlockType*> FromUsableSpace(Ptr ptr) const;
 
  virtual void DeallocateBlock(BlockType*&& block);
 
 private:
  using BlockResultPrev = internal::GenericBlockResult::Prev;
  using BlockResultNext = internal::GenericBlockResult::Next;
 
  // Let unit tests call internal methods in order to "preallocate" blocks..
  template <typename, size_t>
  friend class test::BlockAllocatorTest;
 
  virtual size_t DoGetMaxAllocatable() = 0;
 
  virtual BlockResult<BlockType> ChooseBlock(Layout layout) = 0;
 
  virtual void ReserveBlock(BlockType&) {}
 
  virtual void RecycleBlock(BlockType&) {}
 
  virtual void Flush() {}
 
  static bool PrevIsFree(const BlockType* block) {
    auto* prev = block->Prev();
    return prev != nullptr && prev->IsFree();
  }
 
  static bool NextIsFree(const BlockType* block) {
    auto* next = block->Next();
    return next != nullptr && next->IsFree();
  }
 
  void UpdateLast(BlockType* block);
 
  // Represents the range of blocks for this allocator.
  size_t capacity_ = 0;
  size_t allocated_ = 0;
  BlockType* first_ = nullptr;
  BlockType* last_ = nullptr;
  uint16_t unpoisoned_ = 0;
};
 
 
// Template method implementations
 
namespace internal {
 
template <typename BlockType>
constexpr Capabilities GenericBlockAllocator::GetCapabilities() {
  Capabilities common = kImplementsGetUsableLayout |
                        kImplementsGetAllocatedLayout | kImplementsGetCapacity |
                        kImplementsRecognizes;
  if constexpr (has_layout_v<BlockType>) {
    return common | kImplementsGetRequestedLayout;
  } else {
    return common;
  }
}
 
}  // namespace internal
 
template <typename BlockType>
BlockAllocator<BlockType>::~BlockAllocator() {
  if constexpr (Hardening::kIncludesRobustChecks) {
    for (auto* block : blocks()) {
      if (!block->IsFree()) {
        CrashOnAllocated(block);
      }
    }
  }
}
 
template <typename BlockType>
void BlockAllocator<BlockType>::Init(ByteSpan region) {
  Result<BlockType*> result = BlockType::Init(region);
  PW_ASSERT(result.ok());
  Init(*result);
}
 
template <typename BlockType>
void BlockAllocator<BlockType>::Init(BlockType* begin) {
  if constexpr (Hardening::kIncludesRobustChecks) {
    PW_ASSERT(begin != nullptr);
    PW_ASSERT(begin->Prev() == nullptr);
  }
  first_ = begin;
  for (auto* block : blocks()) {
    last_ = block;
    capacity_ += block->OuterSize();
    if (block->IsFree()) {
      RecycleBlock(*block);
    }
  }
}
 
template <typename BlockType>
void* BlockAllocator<BlockType>::DoAllocate(Layout layout) {
  if (capacity_ == 0) {
    // Not initialized.
    return nullptr;
  }
 
  if constexpr (Hardening::kIncludesDebugChecks) {
    PW_ASSERT(last_->Next() == nullptr);
  }
  auto result = ChooseBlock(layout);
  if (!result.ok()) {
    // No valid block for request.
    return nullptr;
  }
  BlockType* block = result.block();
  allocated_ += block->OuterSize();
  switch (result.prev()) {
    case BlockResultPrev::kSplitNew:
      // New free blocks may be created when allocating.
      RecycleBlock(*(block->Prev()));
      break;
    case BlockResultPrev::kResizedLarger:
      // Extra bytes may be appended to the previous block.
      allocated_ += result.size();
      break;
    case BlockResultPrev::kUnchanged:
    case BlockResultPrev::kResizedSmaller:
      break;
  }
  if (result.next() == BlockResultNext::kSplitNew) {
    RecycleBlock(*(block->Next()));
  }
 
  UpdateLast(block);
  if constexpr (Hardening::kIncludesDebugChecks) {
    PW_ASSERT(block <= last_);
  }
 
  return block->UsableSpace();
}
 
template <typename BlockType>
void BlockAllocator<BlockType>::DoDeallocate(void* ptr) {
  BlockType* block = FromUsableSpace(ptr);
  if (block->IsFree()) {
    if constexpr (Hardening::kIncludesBasicChecks) {
      CrashOnDoubleFree(block);
    } else {
      return;
    }
  }
  DeallocateBlock(std::move(block));
}
 
template <typename BlockType>
void BlockAllocator<BlockType>::DeallocateBlock(BlockType*&& block) {
  // Neighboring blocks may be merged when freeing.
  if (auto* prev = block->Prev(); prev != nullptr && prev->IsFree()) {
    ReserveBlock(*prev);
  }
  if (auto* next = block->Next(); next != nullptr && next->IsFree()) {
    ReserveBlock(*next);
  }
 
  // Free the block and merge it with its neighbors, if possible.
  allocated_ -= block->OuterSize();
  auto free_result = BlockType::Free(std::move(block));
  block = free_result.block();
  UpdateLast(block);
 
  if (free_result.prev() == BlockResultPrev::kResizedSmaller) {
    // Bytes were reclaimed from the previous block.
    allocated_ -= free_result.size();
  }
 
  if constexpr (is_poisonable_v<BlockType> && kPoisonInterval != 0) {
    ++unpoisoned_;
    if (unpoisoned_ >= kPoisonInterval) {
      block->Poison();
      unpoisoned_ = 0;
    }
  }
 
  RecycleBlock(*block);
}
 
template <typename BlockType>
bool BlockAllocator<BlockType>::DoResize(void* ptr, size_t new_size) {
  BlockType* block = FromUsableSpace(ptr);
 
  // Neighboring blocks may be merged when resizing.
  if (auto* next = block->Next(); next != nullptr && next->IsFree()) {
    ReserveBlock(*next);
  }
 
  size_t old_size = block->OuterSize();
  if (!block->Resize(new_size).ok()) {
    return false;
  }
  allocated_ -= old_size;
  allocated_ += block->OuterSize();
  UpdateLast(block);
 
  if (auto* next = block->Next(); next != nullptr && next->IsFree()) {
    RecycleBlock(*next);
  }
 
  return true;
}
 
template <typename BlockType>
Result<Layout> BlockAllocator<BlockType>::DoGetInfo(InfoType info_type,
                                                    const void* ptr) const {
  // Handle types not related to a block first.
  if (info_type == InfoType::kCapacity) {
    return Layout(capacity_);
  }
  // Get a block from the given pointer.
  if (ptr < first_->UsableSpace() || last_->UsableSpace() < ptr) {
    return Status::NotFound();
  }
  const auto* block = BlockType::FromUsableSpace(ptr);
  if (!block->IsValid()) {
    return Status::DataLoss();
  }
  if (block->IsFree()) {
    return Status::FailedPrecondition();
  }
  if constexpr (kCapabilities.has(kImplementsGetRequestedLayout)) {
    if (info_type == InfoType::kRequestedLayoutOf) {
      return block->RequestedLayout();
    }
  }
  switch (info_type) {
    case InfoType::kUsableLayoutOf:
      return Layout(block->InnerSize(), BlockType::kAlignment);
    case InfoType::kAllocatedLayoutOf:
      return Layout(block->OuterSize(), BlockType::kAlignment);
    case InfoType::kRecognizes:
      return Layout();
    case InfoType::kCapacity:
    case InfoType::kRequestedLayoutOf:
    default:
      return Status::Unimplemented();
  }
}
 
template <typename BlockType>
std::optional<Fragmentation> BlockAllocator<BlockType>::DoMeasureFragmentation()
    const {
  Fragmentation fragmentation;
  for (auto block : blocks()) {
    if (block->IsFree()) {
      fragmentation.AddFragment(block->InnerSize() / BlockType::kAlignment);
    }
  }
  return fragmentation;
}
 
template <typename BlockType>
template <typename Ptr>
internal::copy_const_ptr_t<Ptr, BlockType*>
BlockAllocator<BlockType>::FromUsableSpace(Ptr ptr) const {
  if (ptr < first_->UsableSpace() || last_->UsableSpace() < ptr) {
    if constexpr (Hardening::kIncludesBasicChecks) {
      CrashOnOutOfRange(ptr);
    }
    return nullptr;
  }
  auto* block = BlockType::FromUsableSpace(ptr);
  if (!block->IsValid()) {
    if constexpr (Hardening::kIncludesBasicChecks) {
      block->CheckInvariants();
    }
    return nullptr;
  }
  return block;
}
 
template <typename BlockType>
void BlockAllocator<BlockType>::UpdateLast(BlockType* block) {
  BlockType* next = block->Next();
  if (next == nullptr) {
    last_ = block;
  } else if (next->Next() == nullptr) {
    last_ = next;
  }
}
 
}  // namespace pw::allocator