allocatable.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 "pw_allocator/block/contiguous.h"
#include "pw_allocator/block/result.h"
#include "pw_allocator/hardening.h"
#include "pw_allocator/layout.h"
#include "pw_bytes/alignment.h"
#include "pw_status/status.h"
#include "pw_status/status_with_size.h"
 
namespace pw::allocator {
namespace internal {
 
// Trivial base class for trait support.
struct AllocatableBase {};
 
}  // namespace internal
 
template <typename Derived>
class AllocatableBlock : public internal::AllocatableBase {
 protected:
  constexpr explicit AllocatableBlock() {
    // Assert within a function, since `Derived` is not complete when this type
    // is defined.
    static_assert(is_contiguous_v<Derived>,
                  "Types derived from AllocatableBlock must also derive from "
                  "ContiguousBlock");
  }
 
 public:
  constexpr bool IsFree() const;
 
  constexpr bool Used() const { return !IsFree(); }
 
  constexpr StatusWithSize CanAlloc(Layout layout) const;
 
  static constexpr BlockResult<Derived> AllocFirst(Derived*&& block,
                                                   Layout layout);
 
  static constexpr BlockResult<Derived> AllocLast(Derived*&& block,
                                                  Layout layout);
 
  constexpr BlockResult<Derived> Resize(size_t new_inner_size);
 
  static constexpr BlockResult<Derived> Free(Derived*&& block);
 
 protected:
  constexpr StatusWithSize DoCanAlloc(Layout layout) const;
 
  static constexpr BlockResult<Derived> DoAllocFirst(Derived*&& block,
                                                     Layout layout);
 
  static constexpr BlockResult<Derived> DoAllocLast(Derived*&& block,
                                                    Layout layout);
 
  constexpr BlockResult<Derived> DoResize(size_t new_inner_size,
                                          bool shifted = false);
 
  static constexpr BlockResult<Derived> DoFree(Derived*&& block);
 
 private:
  using BlockResultPrev = internal::GenericBlockResult::Prev;
  using BlockResultNext = internal::GenericBlockResult::Next;
 
  constexpr Derived* derived() { return static_cast<Derived*>(this); }
  constexpr const Derived* derived() const {
    return static_cast<const Derived*>(this);
  }
 
  // AlignableBlock calls DoCanAlloc, DoAllocLast, DoResize
  template <typename>
  friend class AlignableBlock;
 
  // BlockWithLayout calls DoFree, DoResize
  template <typename>
  friend class BlockWithLayout;
};
 
template <typename BlockType>
struct is_allocatable : std::is_base_of<internal::AllocatableBase, BlockType> {
};
 
template <typename BlockType>
constexpr bool is_allocatable_v = is_allocatable<BlockType>::value;
 
// Template method implementations.
 
template <typename Derived>
constexpr bool AllocatableBlock<Derived>::IsFree() const {
  if constexpr (Hardening::kIncludesDebugChecks) {
    derived()->CheckInvariants();
  }
  return derived()->IsFreeUnchecked();
}
 
template <typename Derived>
constexpr StatusWithSize AllocatableBlock<Derived>::CanAlloc(
    Layout layout) const {
  if constexpr (Hardening::kIncludesDebugChecks) {
    derived()->CheckInvariants();
  }
  return derived()->DoCanAlloc(layout);
}
 
template <typename Derived>
constexpr StatusWithSize AllocatableBlock<Derived>::DoCanAlloc(
    Layout layout) const {
  if (!derived()->IsFree()) {
    return StatusWithSize::FailedPrecondition();
  }
  if (layout.size() == 0) {
    return StatusWithSize::InvalidArgument();
  }
  size_t extra = derived()->InnerSize();
  size_t new_inner_size = AlignUp(layout.size(), Derived::kAlignment);
  if (PW_SUB_OVERFLOW(extra, new_inner_size, &extra)) {
    return StatusWithSize::ResourceExhausted();
  }
  return StatusWithSize(extra);
}
 
template <typename Derived>
constexpr BlockResult<Derived> AllocatableBlock<Derived>::AllocFirst(
    Derived*&& block, Layout layout) {
  if (block == nullptr || layout.size() == 0) {
    return BlockResult<Derived>(block, Status::InvalidArgument());
  }
  if constexpr (Hardening::kIncludesRobustChecks) {
    block->CheckInvariants();
  }
  if (!block->IsFree()) {
    return BlockResult<Derived>(block, Status::FailedPrecondition());
  }
  return Derived::DoAllocFirst(std::move(block), layout);
}
 
template <typename Derived>
constexpr BlockResult<Derived> AllocatableBlock<Derived>::DoAllocFirst(
    Derived*&& block, Layout layout) {
  size_t size = AlignUp(layout.size(), Derived::kAlignment);
  layout = Layout(size, layout.alignment());
  StatusWithSize can_alloc = block->DoCanAlloc(layout);
  if (!can_alloc.ok()) {
    return BlockResult<Derived>(block, can_alloc.status());
  }
  size_t extra = can_alloc.size();
  BlockResult<Derived> result(block);
  if (extra >= Derived::kMinOuterSize) {
    // Split the large padding off the back.
    block->DoSplitFirst(block->InnerSize() - extra);
    result = BlockResult<Derived>(block, BlockResultNext::kSplitNew);
  }
  block->SetFree(false);
  return result;
}
 
template <typename Derived>
constexpr BlockResult<Derived> AllocatableBlock<Derived>::AllocLast(
    Derived*&& block, Layout layout) {
  if (block == nullptr || layout.size() == 0) {
    return BlockResult<Derived>(block, Status::InvalidArgument());
  }
  if constexpr (Hardening::kIncludesRobustChecks) {
    block->CheckInvariants();
  }
  if (!block->IsFree()) {
    return BlockResult<Derived>(block, Status::FailedPrecondition());
  }
  return Derived::DoAllocLast(std::move(block), layout);
}
 
template <typename Derived>
constexpr BlockResult<Derived> AllocatableBlock<Derived>::DoAllocLast(
    Derived*&& block, Layout layout) {
  size_t size = AlignUp(layout.size(), Derived::kAlignment);
  layout = Layout(size, layout.alignment());
  StatusWithSize can_alloc = block->DoCanAlloc(layout);
  if (!can_alloc.ok()) {
    return BlockResult<Derived>(block, can_alloc.status());
  }
  size_t extra = can_alloc.size();
  BlockResult<Derived> result(block);
  Derived* prev = block->Prev();
  if (extra >= Derived::kMinOuterSize) {
    // Split the large padding off the front.
    block = block->DoSplitLast(layout.size());
    prev = block->Prev();
    result = BlockResult<Derived>(block, BlockResultPrev::kSplitNew);
 
  } else if (extra != 0 && prev != nullptr) {
    // The small amount of padding can be appended to the previous block.
    prev->DoResize(prev->InnerSize() + extra, true).IgnoreUnlessStrict();
    block = prev->Next();
    result =
        BlockResult<Derived>(block, BlockResultPrev::kResizedLarger, extra);
  }
  block->SetFree(false);
  return result;
}
 
template <typename Derived>
constexpr BlockResult<Derived> AllocatableBlock<Derived>::Resize(
    size_t new_inner_size) {
  if constexpr (Hardening::kIncludesRobustChecks) {
    derived()->CheckInvariants();
  }
  if (derived()->IsFree()) {
    return BlockResult<Derived>(derived(), Status::FailedPrecondition());
  }
  return derived()->DoResize(new_inner_size, /* shifted: */ false);
}
 
template <typename Derived>
constexpr BlockResult<Derived> AllocatableBlock<Derived>::DoResize(
    size_t new_inner_size, bool) {
  size_t old_inner_size = derived()->InnerSize();
  new_inner_size = AlignUp(new_inner_size, Derived::kAlignment);
  if (old_inner_size == new_inner_size) {
    return BlockResult<Derived>(derived());
  }
 
  // 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.
  derived()->SetFree(true);
  Derived* next = derived()->Next();
  BlockResult<Derived> result(derived());
  if (next != nullptr && next->IsFree()) {
    derived()->DoMergeNext();
    result = BlockResult<Derived>(derived(), BlockResultNext::kMerged);
  }
  size_t merged_inner_size = derived()->InnerSize();
  if (merged_inner_size < new_inner_size) {
    // The merged block is too small for the resized block. Restore the original
    // blocks as needed.
    if (merged_inner_size != old_inner_size) {
      derived()->DoSplitFirst(old_inner_size);
    }
    derived()->SetFree(false);
    return BlockResult<Derived>(derived(), Status::ResourceExhausted());
  }
  if (new_inner_size + Derived::kMinOuterSize <= merged_inner_size) {
    // There is enough room after the resized block for another trailing block.
    derived()->DoSplitFirst(new_inner_size);
    result = result.next() == BlockResultNext::kMerged
                 ? BlockResult<Derived>(derived(), BlockResultNext::kResized)
                 : BlockResult<Derived>(derived(), BlockResultNext::kSplitNew);
  }
  derived()->SetFree(false);
  return result;
}
 
template <typename Derived>
constexpr BlockResult<Derived> AllocatableBlock<Derived>::Free(
    Derived*&& block) {
  if (block == nullptr) {
    return BlockResult<Derived>(block, Status::InvalidArgument());
  }
  if constexpr (Hardening::kIncludesRobustChecks) {
    block->CheckInvariants();
  }
  return Derived::DoFree(std::move(block));
}
 
template <typename Derived>
constexpr BlockResult<Derived> AllocatableBlock<Derived>::DoFree(
    Derived*&& block) {
  block->SetFree(true);
  BlockResult<Derived> result(block);
 
  // Try to merge the previous block with this one.
  Derived* prev = block->Prev();
  if (prev != nullptr && prev->IsFree()) {
    prev->DoMergeNext();
    block = prev;
    result = BlockResult<Derived>(block, BlockResultNext::kMerged);
  }
 
  // Try to merge this block with the next one.
  Derived* next = block->Next();
  if (next != nullptr && next->IsFree()) {
    block->DoMergeNext();
    result = BlockResult<Derived>(block, BlockResultNext::kMerged);
  }
 
  if constexpr (Hardening::kIncludesDebugChecks) {
    block->CheckInvariants();
  }
  return result;
}
 
}  // namespace pw::allocator