block.h

// Copyright 2020 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 <algorithm>
#include <climits>
#include <cstddef>
#include <cstdint>
#include <cstring>
 
#include "pw_allocator/allocator.h"
#include "pw_allocator/buffer.h"
#include "pw_bytes/alignment.h"
#include "pw_bytes/span.h"
#include "pw_result/result.h"
#include "pw_status/status.h"
#include "pw_status/try.h"
 
namespace pw::allocator {
namespace internal {
 
// Types of corrupted blocks, and functions to crash with an error message
// corresponding to each type. These functions are implemented independent of
// any template parameters to allow them to use `PW_CHECK`.
enum BlockStatus {
  kValid,
  kMisaligned,
  kPrevMismatched,
  kNextMismatched,
  kPoisonCorrupted,
};
void CrashMisaligned(uintptr_t addr);
void CrashNextMismatched(uintptr_t addr, uintptr_t next_prev);
void CrashPrevMismatched(uintptr_t addr, uintptr_t prev_next);
void CrashPoisonCorrupted(uintptr_t addr);
 
}  // namespace internal
 
template <typename OffsetType = uintptr_t,
          size_t kAlign = alignof(OffsetType),
          bool kCanPoison = false>
class Block {
 public:
  using offset_type = OffsetType;
  static_assert(std::is_unsigned_v<offset_type>,
                "offset type must be unsigned");
 
  static constexpr size_t kAlignment = std::max(kAlign, alignof(offset_type));
  static constexpr size_t kBlockOverhead = AlignUp(sizeof(Block), kAlignment);
 
  // No copy or move.
  Block(const Block& other) = delete;
  Block& operator=(const Block& other) = delete;
 
  static Result<Block*> Init(ByteSpan region);
 
  template <int&... DeducedTypesOnly,
            typename PtrType,
            bool is_const_ptr = std::is_const_v<std::remove_pointer_t<PtrType>>,
            typename BytesPtr =
                std::conditional_t<is_const_ptr, const std::byte*, std::byte*>,
            typename BlockPtr =
                std::conditional_t<is_const_ptr, const Block*, Block*>>
  static BlockPtr FromUsableSpace(PtrType usable_space) {
    // Perform memory laundering to prevent the compiler from tracing the memory
    // used to store the block and to avoid optimaztions that may be invalidated
    // by the use of placement-new to create blocks in `Init` and `Split`.
    auto* bytes = reinterpret_cast<BytesPtr>(usable_space);
    return std::launder(reinterpret_cast<BlockPtr>(bytes - kBlockOverhead));
  }
 
  size_t OuterSize() const { return next_ * kAlignment; }
 
  size_t InnerSize() const { return OuterSize() - kBlockOverhead; }
 
  size_t RequestedSize() const { return InnerSize() - padding_; }
 
  std::byte* UsableSpace() {
    return std::launder(reinterpret_cast<std::byte*>(this) + kBlockOverhead);
  }
  const std::byte* UsableSpace() const {
    return std::launder(reinterpret_cast<const std::byte*>(this) +
                        kBlockOverhead);
  }
 
  static Status AllocFirst(Block*& block, Layout layout);
 
  StatusWithSize CanAllocLast(Layout layout) const;
 
  static Status AllocLast(Block*& block, Layout layout);
 
  static void Free(Block*& block);
 
  static Status Resize(Block*& block, size_t new_inner_size);
 
  Block* Next() const;
 
  static Block* NextBlock(const Block* block) {
    return block == nullptr ? nullptr : block->Next();
  }
 
  Block* Prev() const;
 
  static Block* PrevBlock(const Block* block) {
    return block == nullptr ? nullptr : block->Prev();
  }
 
  size_t Alignment() const { return Used() ? info_.alignment : 1; }
 
  bool Used() const { return info_.used; }
 
  bool Last() const { return info_.last; }
 
  void MarkUsed() { info_.used = 1; }
 
  void MarkFree() { info_.used = 0; }
 
  void MarkLast() { info_.last = 1; }
 
  void ClearLast() { info_.last = 1; }
 
  void Poison(bool should_poison = true);
 
  bool IsValid() const { return CheckStatus() == internal::kValid; }
 
  void CrashIfInvalid() const;
 
 private:
  static constexpr uint8_t kPoisonByte = 0xf7;
 
  static ByteSpan AsBytes(Block*&& block);
 
  static Block* AsBlock(size_t prev_outer_size, ByteSpan bytes);
 
  Block(size_t prev_outer_size, size_t outer_size);
 
  internal::BlockStatus CheckStatus() const;
 
  static void ShiftBlock(Block*& block, size_t pad_size);
 
  static Block* Split(Block*& block, size_t new_inner_size);
 
  static void MergeNext(Block*& block);
 
  bool CheckPoison() const;
 
  offset_type prev_ = 0;
 
  offset_type next_ = 0;
 
  struct {
    uint16_t used : 1;
    uint16_t poisoned : 1;
    uint16_t last : 1;
    uint16_t alignment : 13;
  } info_;
 
  uint16_t padding_ = 0;
 
 public:
  // Associated types.
 
  class Iterator final {
   public:
    Iterator(Block* block) : block_(block) {}
    Iterator& operator++() {
      block_ = Block::NextBlock(block_);
      return *this;
    }
    bool operator!=(const Iterator& other) { return block_ != other.block_; }
    Block* operator*() { return block_; }
 
   private:
    Block* block_;
  };
 
  class ReverseIterator final {
   public:
    ReverseIterator(Block* block) : block_(block) {}
    ReverseIterator& operator++() {
      block_ = Block::PrevBlock(block_);
      return *this;
    }
    bool operator!=(const ReverseIterator& other) {
      return block_ != other.block_;
    }
    Block* operator*() { return block_; }
 
   private:
    Block* block_;
  };
 
  class Range final {
   public:
    explicit Range(Block* begin) : begin_(begin), end_(nullptr) {}
 
    Range(Block* begin_inclusive, Block* end_inclusive)
        : begin_(begin_inclusive), end_(end_inclusive->Next()) {}
 
    Iterator& begin() { return begin_; }
    Iterator& end() { return end_; }
 
   private:
    Iterator begin_;
    Iterator end_;
  };
 
  class ReverseRange final {
   public:
    explicit ReverseRange(Block* rbegin) : begin_(rbegin), end_(nullptr) {}
 
    ReverseRange(Block* rbegin_inclusive, Block* rend_inclusive)
        : begin_(rbegin_inclusive), end_(rend_inclusive->Prev()) {}
 
    ReverseIterator& begin() { return begin_; }
    ReverseIterator& end() { return end_; }
 
   private:
    ReverseIterator begin_;
    ReverseIterator end_;
  };
} __attribute__((packed, aligned(kAlign)));
 
// Public template method implementations.
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Result<Block<OffsetType, kAlign, kCanPoison>*>
Block<OffsetType, kAlign, kCanPoison>::Init(ByteSpan region) {
  Result<ByteSpan> result = GetAlignedSubspan(region, kAlignment);
  if (!result.ok()) {
    return result.status();
  }
  region = result.value();
  if (region.size() < kBlockOverhead) {
    return Status::ResourceExhausted();
  }
  if (std::numeric_limits<OffsetType>::max() < region.size() / kAlignment) {
    return Status::OutOfRange();
  }
  Block* block = AsBlock(0, region);
  block->MarkLast();
  return block;
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Status Block<OffsetType, kAlign, kCanPoison>::AllocFirst(Block*& block,
                                                         Layout layout) {
  if (block == nullptr || layout.size() == 0) {
    return Status::InvalidArgument();
  }
  block->CrashIfInvalid();
  if (block->Used()) {
    return Status::FailedPrecondition();
  }
  Block* prev = block->Prev();
 
  // Check if padding will be needed at the front to align the usable space.
  size_t alignment = std::max(layout.alignment(), kAlignment);
  auto addr = reinterpret_cast<uintptr_t>(block->UsableSpace());
  size_t pad_size = AlignUp(addr, alignment) - addr;
  bool should_poison = block->info_.poisoned;
  if (pad_size == 0) {
    // No padding needed.
  } else if (pad_size > kBlockOverhead) {
    // Enough padding is needed that it can be split off as a new free block.
  } else if (prev == nullptr) {
    // First block; increase padding to at least the minimum block size.
    pad_size += AlignUp(kBlockOverhead, alignment);
  }
 
  // Make sure everything fits.
  size_t inner_size = AlignUp(layout.size(), kAlignment);
  if (block->InnerSize() < pad_size + inner_size) {
    return Status::OutOfRange();
  }
  ShiftBlock(block, pad_size);
 
  // If the block is large enough to have a trailing block, split it to get the
  // requested usable space.
  if (inner_size + kBlockOverhead <= block->InnerSize()) {
    Block* trailing = Split(block, inner_size);
    trailing->Poison(should_poison);
  }
 
  block->MarkUsed();
  block->info_.alignment = alignment;
  block->padding_ = block->InnerSize() - layout.size();
  return OkStatus();
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
StatusWithSize Block<OffsetType, kAlign, kCanPoison>::CanAllocLast(
    Layout layout) const {
  if (Used()) {
    return StatusWithSize::FailedPrecondition();
  }
  CrashIfInvalid();
  // Find the last address that is aligned and is followed by enough space for
  // block overhead and the requested size.
  if (InnerSize() < layout.size()) {
    return StatusWithSize::OutOfRange();
  }
  auto addr = reinterpret_cast<uintptr_t>(UsableSpace());
  size_t alignment = std::max(layout.alignment(), kAlignment);
  uintptr_t next = AlignDown(addr + (InnerSize() - layout.size()), alignment);
 
  if (next < addr) {
    // Requested size does not fit.
    return StatusWithSize::ResourceExhausted();
  }
  return StatusWithSize(next - addr);
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Status Block<OffsetType, kAlign, kCanPoison>::AllocLast(Block*& block,
                                                        Layout layout) {
  if (block == nullptr || layout.size() == 0) {
    return Status::InvalidArgument();
  }
  size_t pad_size = 0;
  PW_TRY_ASSIGN(pad_size, block->CanAllocLast(layout));
  ShiftBlock(block, pad_size);
 
  block->MarkUsed();
  block->info_.alignment = layout.alignment();
  block->padding_ = block->InnerSize() - layout.size();
  return OkStatus();
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
void Block<OffsetType, kAlign, kCanPoison>::ShiftBlock(Block*& block,
                                                       size_t pad_size) {
  if (pad_size == 0) {
    return;
  }
 
  // Check if this is the first block.
  Block* prev = block->Prev();
  if (prev == nullptr && pad_size <= kBlockOverhead) {
    pad_size += kBlockOverhead;
  }
 
  bool should_poison = block->info_.poisoned;
  if (pad_size <= kBlockOverhead) {
    // The small amount of padding can be appended to the previous block.
    Block::Resize(prev, prev->InnerSize() + pad_size).IgnoreError();
    prev->padding_ += pad_size;
    block = prev->Next();
 
  } else if (kBlockOverhead < pad_size) {
    // Split the large padding off the front.
    Block* leading = block;
    block = Split(leading, pad_size - kBlockOverhead);
    leading->Poison(should_poison);
  }
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
void Block<OffsetType, kAlign, kCanPoison>::Free(Block*& block) {
  if (block == nullptr) {
    return;
  }
  block->MarkFree();
  Block* prev = block->Prev();
  if (prev != nullptr && !prev->Used()) {
    MergeNext(prev);
    block = prev;
  }
  MergeNext(block);
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Status Block<OffsetType, kAlign, kCanPoison>::Resize(Block*& block,
                                                     size_t new_inner_size) {
  if (block == nullptr) {
    return Status::InvalidArgument();
  }
  if (!block->Used()) {
    return Status::FailedPrecondition();
  }
  size_t requested_inner_size = new_inner_size;
  size_t old_inner_size = block->InnerSize();
  size_t alignment = block->Alignment();
  uint16_t padding = block->padding_;
  new_inner_size = AlignUp(new_inner_size, kAlignment);
 
  if (old_inner_size == new_inner_size) {
    return OkStatus();
  }
 
  // Treat the block as free and try to combine it with the next block. At most
  // one free block is expected to follow this block.
  block->MarkFree();
  MergeNext(block);
 
  Status status = OkStatus();
 
  if (block->InnerSize() < new_inner_size) {
    // The merged block is too small for the resized block.
    status = Status::OutOfRange();
 
  } else if (new_inner_size + kBlockOverhead <= block->InnerSize()) {
    // There is enough room after the resized block for another trailing block.
    bool should_poison = block->info_.poisoned;
    Block* trailing = Split(block, new_inner_size);
    trailing->Poison(should_poison);
  }
 
  if (status.ok()) {
    padding = block->InnerSize() - requested_inner_size;
  } else if (block->InnerSize() != old_inner_size) {
    // Restore the original block on failure.
    Split(block, old_inner_size);
  }
  block->MarkUsed();
  block->info_.alignment = alignment;
  block->padding_ = padding;
  return status;
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Block<OffsetType, kAlign, kCanPoison>*
Block<OffsetType, kAlign, kCanPoison>::Split(Block*& block,
                                             size_t new_inner_size) {
  size_t prev_outer_size = block->prev_ * kAlignment;
  size_t outer_size1 = new_inner_size + kBlockOverhead;
  bool is_last = block->Last();
  ByteSpan bytes = AsBytes(std::move(block));
  Block* block1 = AsBlock(prev_outer_size, bytes.subspan(0, outer_size1));
  Block* block2 = AsBlock(outer_size1, bytes.subspan(outer_size1));
  if (is_last) {
    block2->MarkLast();
  } else {
    block2->Next()->prev_ = block2->next_;
  }
  block = std::move(block1);
  return block2;
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
void Block<OffsetType, kAlign, kCanPoison>::MergeNext(Block*& block) {
  if (block->Last()) {
    return;
  }
  Block* next = block->Next();
  if (block->Used() || next->Used()) {
    return;
  }
  size_t prev_outer_size = block->prev_ * kAlignment;
  bool is_last = next->Last();
  ByteSpan prev_bytes = AsBytes(std::move(block));
  ByteSpan next_bytes = AsBytes(std::move(next));
  size_t outer_size = prev_bytes.size() + next_bytes.size();
  std::byte* merged = ::new (prev_bytes.data()) std::byte[outer_size];
  block = AsBlock(prev_outer_size, ByteSpan(merged, outer_size));
  if (is_last) {
    block->MarkLast();
  } else {
    block->Next()->prev_ = block->next_;
  }
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Block<OffsetType, kAlign, kCanPoison>*
Block<OffsetType, kAlign, kCanPoison>::Next() const {
  uintptr_t addr = Last() ? 0 : reinterpret_cast<uintptr_t>(this) + OuterSize();
  // See the note in `FromUsableSpace` about memory laundering.
  return std::launder(reinterpret_cast<Block*>(addr));
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Block<OffsetType, kAlign, kCanPoison>*
Block<OffsetType, kAlign, kCanPoison>::Prev() const {
  uintptr_t addr =
      (prev_ == 0) ? 0
                   : reinterpret_cast<uintptr_t>(this) - (prev_ * kAlignment);
  // See the note in `FromUsableSpace` about memory laundering.
  return std::launder(reinterpret_cast<Block*>(addr));
}
 
// Private template method implementations.
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Block<OffsetType, kAlign, kCanPoison>::Block(size_t prev_outer_size,
                                             size_t outer_size) {
  prev_ = prev_outer_size / kAlignment;
  next_ = outer_size / kAlignment;
  info_.used = 0;
  info_.poisoned = 0;
  info_.last = 0;
  info_.alignment = kAlignment;
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
ByteSpan Block<OffsetType, kAlign, kCanPoison>::AsBytes(Block*&& block) {
  size_t block_size = block->OuterSize();
  std::byte* bytes = ::new (std::move(block)) std::byte[block_size];
  return {bytes, block_size};
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
Block<OffsetType, kAlign, kCanPoison>*
Block<OffsetType, kAlign, kCanPoison>::AsBlock(size_t prev_outer_size,
                                               ByteSpan bytes) {
  return ::new (bytes.data()) Block(prev_outer_size, bytes.size());
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
void Block<OffsetType, kAlign, kCanPoison>::Poison(bool should_poison) {
  if constexpr (kCanPoison) {
    if (!Used() && should_poison) {
      std::memset(UsableSpace(), kPoisonByte, InnerSize());
      info_.poisoned = true;
    }
  }
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
bool Block<OffsetType, kAlign, kCanPoison>::CheckPoison() const {
  if constexpr (kCanPoison) {
    if (!info_.poisoned) {
      return true;
    }
    const std::byte* begin = UsableSpace();
    return std::all_of(begin, begin + InnerSize(), [](std::byte b) {
      return static_cast<uint8_t>(b) == kPoisonByte;
    });
  } else {
    return true;
  }
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
internal::BlockStatus Block<OffsetType, kAlign, kCanPoison>::CheckStatus()
    const {
  if (reinterpret_cast<uintptr_t>(this) % kAlignment != 0) {
    return internal::kMisaligned;
  }
  if (!Last() && (this >= Next() || this != Next()->Prev())) {
    return internal::kNextMismatched;
  }
  if (Prev() && (this <= Prev() || this != Prev()->Next())) {
    return internal::kPrevMismatched;
  }
  if (!Used() && !CheckPoison()) {
    return internal::kPoisonCorrupted;
  }
  return internal::kValid;
}
 
template <typename OffsetType, size_t kAlign, bool kCanPoison>
void Block<OffsetType, kAlign, kCanPoison>::CrashIfInvalid() const {
  uintptr_t addr = reinterpret_cast<uintptr_t>(this);
  switch (CheckStatus()) {
    case internal::kValid:
      break;
    case internal::kMisaligned:
      internal::CrashMisaligned(addr);
      break;
    case internal::kNextMismatched:
      internal::CrashNextMismatched(
          addr, reinterpret_cast<uintptr_t>(Next()->Prev()));
      break;
    case internal::kPrevMismatched:
      internal::CrashPrevMismatched(
          addr, reinterpret_cast<uintptr_t>(Prev()->Next()));
      break;
    case internal::kPoisonCorrupted:
      internal::CrashPoisonCorrupted(addr);
      break;
  }
}
 
}  // namespace pw::allocator