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 "pw_allocator/allocator.h"
#include "pw_allocator/block.h"
#include "pw_allocator/capability.h"
#include "pw_allocator/fragmentation.h"
#include "pw_bytes/span.h"
#include "pw_result/result.h"
#include "pw_status/status.h"
 
namespace pw::allocator {
namespace internal {
 
class GenericBlockAllocator : public Allocator {
 public:
  static constexpr Capabilities kCapabilities =
      kImplementsGetRequestedLayout | kImplementsGetUsableLayout |
      kImplementsGetAllocatedLayout | kImplementsGetCapacity |
      kImplementsRecognizes;
 
  // Not copyable or movable.
  GenericBlockAllocator(const GenericBlockAllocator&) = delete;
  GenericBlockAllocator& operator=(const GenericBlockAllocator&) = delete;
  GenericBlockAllocator(GenericBlockAllocator&&) = delete;
  GenericBlockAllocator& operator=(GenericBlockAllocator&&) = delete;
 
 protected:
  constexpr GenericBlockAllocator() : Allocator(kCapabilities) {}
 
  static void CrashOnAllocated(void* allocated);
 
  static void CrashOnInvalidFree(void* freed);
 
  static void CrashOnDoubleFree(void* freed);
};
 
}  // namespace internal
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
class BlockAllocator : public internal::GenericBlockAllocator {
 public:
  using BlockType = Block<OffsetType, kAlign, kPoisonInterval != 0>;
  using Range = typename BlockType::Range;
 
  constexpr BlockAllocator() : internal::GenericBlockAllocator() {}
 
  explicit BlockAllocator(ByteSpan region) : BlockAllocator() { Init(region); }
 
  ~BlockAllocator() override { Reset(); }
 
  Range blocks() const;
 
  void Init(ByteSpan region);
 
  void Init(BlockType* begin) { return Init(begin, nullptr); }
 
  virtual void Init(BlockType* begin, BlockType* end);
 
  Fragmentation MeasureFragmentation() const;
 
  void Reset();
 
 protected:
  using ReverseRange = typename BlockType::ReverseRange;
 
  ReverseRange rblocks();
 
  template <typename PtrType,
            typename BlockPtrType = std::conditional_t<
                std::is_const_v<std::remove_pointer_t<PtrType>>,
                const BlockType*,
                BlockType*>>
  Result<BlockPtrType> FromUsableSpace(PtrType ptr) const;
 
  virtual void ReserveBlock(BlockType*) {}
 
  virtual void RecycleBlock(BlockType*) {}
 
 private:
  void* DoAllocate(Layout layout) override;
 
  void DoDeallocate(void* ptr) override;
 
  void DoDeallocate(void* ptr, Layout) override { DoDeallocate(ptr); }
 
  bool DoResize(void* ptr, size_t new_size) override;
 
  Result<Layout> DoGetInfo(InfoType info_type, const void* ptr) const override;
 
  virtual BlockType* ChooseBlock(Layout layout) = 0;
 
  bool PrevIsFree(const BlockType* block) const {
    return block->Prev() != nullptr && !block->Prev()->Used();
  }
 
  bool NextIsFree(const BlockType* block) const {
    return !block->Last() && !block->Next()->Used();
  }
 
  void UpdateLast(BlockType* block);
 
  // Represents the range of blocks for this allocator.
  size_t capacity_ = 0;
  BlockType* first_ = nullptr;
  BlockType* last_ = nullptr;
  uint16_t unpoisoned_ = 0;
};
 
// Template method implementations
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
typename BlockAllocator<OffsetType, kPoisonInterval, kAlign>::Range
BlockAllocator<OffsetType, kPoisonInterval, kAlign>::blocks() const {
  return Range(first_);
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
typename BlockAllocator<OffsetType, kPoisonInterval, kAlign>::ReverseRange
BlockAllocator<OffsetType, kPoisonInterval, kAlign>::rblocks() {
  while (last_ != nullptr && !last_->Last()) {
    last_ = last_->Next();
  }
  return ReverseRange(last_);
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
void BlockAllocator<OffsetType, kPoisonInterval, kAlign>::Init(
    ByteSpan region) {
  Result<BlockType*> result = BlockType::Init(region);
  Init(*result, nullptr);
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
void BlockAllocator<OffsetType, kPoisonInterval, kAlign>::Init(BlockType* begin,
                                                               BlockType* end) {
  PW_ASSERT(begin != nullptr);
  if (end == nullptr) {
    end = begin;
    while (!end->Last()) {
      end = end->Next();
    }
  } else {
    PW_ASSERT(begin < end);
    end->MarkLast();
  }
  first_ = begin;
  last_ = end;
  for (const auto& block : blocks()) {
    capacity_ += block->OuterSize();
  }
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
void BlockAllocator<OffsetType, kPoisonInterval, kAlign>::Reset() {
  for (auto* block : blocks()) {
    if (block->Used()) {
      CrashOnAllocated(block);
    }
  }
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
void* BlockAllocator<OffsetType, kPoisonInterval, kAlign>::DoAllocate(
    Layout layout) {
  BlockType* block = ChooseBlock(layout);
  if (block == nullptr) {
    return nullptr;
  }
  UpdateLast(block);
  return block->UsableSpace();
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
void BlockAllocator<OffsetType, kPoisonInterval, kAlign>::DoDeallocate(
    void* ptr) {
  auto result = FromUsableSpace(ptr);
  if (!result.ok()) {
    CrashOnInvalidFree(ptr);
  }
  BlockType* block = *result;
  if (!block->Used()) {
    CrashOnDoubleFree(block);
  }
 
  // Neighboring blocks may be merged when freeing.
  if (PrevIsFree(block)) {
    ReserveBlock(block->Prev());
  }
  if (NextIsFree(block)) {
    ReserveBlock(block->Next());
  }
 
  // Free the block and merge it with its neighbors, if possible.
  BlockType::Free(block);
  UpdateLast(block);
 
  if constexpr (kPoisonInterval != 0) {
    ++unpoisoned_;
    if (unpoisoned_ >= kPoisonInterval) {
      block->Poison();
      unpoisoned_ = 0;
    }
  }
 
  RecycleBlock(block);
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
bool BlockAllocator<OffsetType, kPoisonInterval, kAlign>::DoResize(
    void* ptr, size_t new_size) {
  auto result = FromUsableSpace(ptr);
  if (!result.ok()) {
    return false;
  }
  BlockType* block = *result;
 
  // Neighboring blocks may be merged when resizing.
  if (NextIsFree(block)) {
    ReserveBlock(block->Next());
  }
 
  if (!BlockType::Resize(block, new_size).ok()) {
    return false;
  }
  UpdateLast(block);
 
  if (NextIsFree(block)) {
    RecycleBlock(block->Next());
  }
 
  return true;
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
Result<Layout> BlockAllocator<OffsetType, kPoisonInterval, kAlign>::DoGetInfo(
    InfoType info_type, const void* ptr) const {
  // Handle types not related to a block first.
  if (info_type == InfoType::kCapacity) {
    return Layout(capacity_, kAlign);
  }
  // Get a block from the given pointer.
  auto result = FromUsableSpace(ptr);
  if (!result.ok()) {
    return Status::NotFound();
  }
  const BlockType* block = result.value();
  if (!block->Used()) {
    return Status::FailedPrecondition();
  }
  switch (info_type) {
    case InfoType::kRequestedLayoutOf:
      return Layout(block->RequestedSize(), block->Alignment());
    case InfoType::kUsableLayoutOf:
      return Layout(block->InnerSize(), block->Alignment());
    case InfoType::kAllocatedLayoutOf:
      return Layout(block->OuterSize(), block->Alignment());
    case InfoType::kRecognizes:
      return Layout();
    case InfoType::kCapacity:
    default:
      return Status::Unimplemented();
  }
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
Fragmentation
BlockAllocator<OffsetType, kPoisonInterval, kAlign>::MeasureFragmentation()
    const {
  Fragmentation fragmentation;
  for (auto block : blocks()) {
    if (!block->Used()) {
      fragmentation.AddFragment(block->InnerSize() / BlockType::kAlignment);
    }
  }
  return fragmentation;
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
template <typename PtrType, typename BlockPtrType>
Result<BlockPtrType>
BlockAllocator<OffsetType, kPoisonInterval, kAlign>::FromUsableSpace(
    PtrType ptr) const {
  if (ptr < first_->UsableSpace() || last_->UsableSpace() < ptr) {
    return Status::OutOfRange();
  }
  BlockPtrType block = BlockType::FromUsableSpace(ptr);
  block->CrashIfInvalid();
  return block;
}
 
template <typename OffsetType, uint16_t kPoisonInterval, uint16_t kAlign>
void BlockAllocator<OffsetType, kPoisonInterval, kAlign>::UpdateLast(
    BlockType* block) {
  if (block->Last()) {
    last_ = block;
  } else if (block->Next()->Last()) {
    last_ = block->Next();
  }
}
 
}  // namespace pw::allocator