guarded_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 "pw_allocator/allocator.h"
#include "pw_allocator/layout.h"
#include "pw_allocator/synchronized_allocator.h"
#include "pw_result/result.h"
#include "pw_sync/borrow.h"
 
namespace pw::allocator {
namespace internal {
 
class GenericGuardedAllocator : public pw::Allocator {
 protected:
  constexpr explicit GenericGuardedAllocator(const Capabilities& capabilities)
      : Allocator(capabilities) {}
 
  static Layout AdjustLayout(Layout layout);
 
  static size_t AdjustSize(void* ptr, size_t inner_size);
 
  static void* GetOriginal(void* ptr);
 
  static void* AddPrefix(void* ptr, size_t alignment);
 
  static void AddSuffix(void* ptr, size_t size);
 
  static bool CheckPrefixAndSuffix(void* ptr, size_t size);
};
 
}  // namespace internal
 
 
template <typename BlockAllocatorType, typename LockType = NoSync>
class GuardedAllocator : public internal::GenericGuardedAllocator {
 private:
  using Base = internal::GenericGuardedAllocator;
 
 public:
  using BlockType = typename BlockAllocatorType::BlockType;
 
  constexpr explicit GuardedAllocator(BlockAllocatorType& allocator)
      : Base(allocator.capabilities()), borrowable_(allocator, lock_) {
    block_ = *(allocator.blocks().end());
  }
 
  void* ValidateOne();
 
  void* ValidateAll();
 
 protected:
  void* DoAllocate(Layout layout) override;
 
  void DoDeallocate(void* ptr) override;
 
  bool DoResize(void* ptr, size_t new_size) override;
 
  Result<Layout> DoGetInfo(InfoType info_type, const void* ptr) const override;
 
 private:
  LockType lock_;
  sync::Borrowable<BlockAllocatorType, LockType> borrowable_;
 
  // Ideally, this would be annotated as being guarded `lock_`. However, the
  // lock is always acquired via `borrowable_`, which does not support thread
  // safety analysis.
  BlockType* block_;
};
 
 
// Template method implementations.
 
template <typename BlockAllocatorType, typename LockType>
void* GuardedAllocator<BlockAllocatorType, LockType>::ValidateOne() {
  auto allocator = borrowable_.acquire();
 
  // Find the bounds of the block range.
  auto range = allocator->blocks();
  BlockType* begin = *(range.begin());
  BlockType* end = *(range.end());
 
  // Ensure there is at least one block.
  if (begin == end) {
    return nullptr;
  }
 
  // Ensure we are starting from a block.
  if (block_ == nullptr || block_ == end) {
    block_ = begin;
  }
 
  // Find the next used block.
  BlockType* prev = block_;
  while (block_->IsFree()) {
    BlockType* next = block_->Next();
    if (next == end) {
      // Loop around.
      next = begin;
    }
    if (next == prev) {
      // All blocks are free.
      return nullptr;
    }
    block_ = next;
  }
 
  // Validate the block.
  void* ptr = block_->UsableSpace();
  size_t size = block_->InnerSize();
  return Base::CheckPrefixAndSuffix(ptr, size) ? nullptr : ptr;
}
 
template <typename BlockAllocatorType, typename LockType>
void* GuardedAllocator<BlockAllocatorType, LockType>::ValidateAll() {
  auto allocator = borrowable_.acquire();
  for (BlockType* block : allocator->blocks()) {
    if (block->IsFree()) {
      continue;
    }
    void* ptr = block->UsableSpace();
    size_t size = block->InnerSize();
    if (!Base::CheckPrefixAndSuffix(ptr, size)) {
      return ptr;
    }
  }
  return nullptr;
}
 
template <typename BlockAllocatorType, typename LockType>
void* GuardedAllocator<BlockAllocatorType, LockType>::DoAllocate(
    Layout layout) {
  layout = Base::AdjustLayout(layout);
  auto allocator = borrowable_.acquire();
  void* ptr = allocator->Allocate(layout);
  if (ptr == nullptr) {
    return nullptr;
  }
  auto* block = BlockType::FromUsableSpace(ptr);
 
  // Bytes may be shifted to the previous block.
  BlockType* prev = block->Prev();
  if (prev != nullptr && !prev->IsFree()) {
    Base::AddSuffix(prev->UsableSpace(), prev->InnerSize());
  }
 
  Base::AddSuffix(block->UsableSpace(), block->InnerSize());
  return Base::AddPrefix(ptr, layout.alignment());
}
 
template <typename BlockAllocatorType, typename LockType>
void GuardedAllocator<BlockAllocatorType, LockType>::DoDeallocate(void* ptr) {
  ptr = Base::GetOriginal(ptr);
 
  // This block and or the next one may be merged on deallocation.
  auto allocator = borrowable_.acquire();
  auto* block = BlockType::FromUsableSpace(ptr);
  BlockType* prev = block->Prev();
  if (block_ == block || block_ == block->Next()) {
    block_ = prev;
  }
  allocator->Deallocate(ptr);
 
  // Bytes may have been shifted from the previous block.
  if (prev != nullptr && !prev->IsFree()) {
    Base::AddSuffix(prev->UsableSpace(), prev->InnerSize());
  }
}
 
template <typename BlockAllocatorType, typename LockType>
bool GuardedAllocator<BlockAllocatorType, LockType>::DoResize(void* ptr,
                                                              size_t new_size) {
  ptr = Base::GetOriginal(ptr);
  new_size = Base::AdjustSize(ptr, new_size);
 
  // The next block may be recreated on resizing.
  auto allocator = borrowable_.acquire();
  auto* block = BlockType::FromUsableSpace(ptr);
  if (block_ == block->Next()) {
    block_ = block;
  }
  if (!allocator->Resize(ptr, new_size)) {
    return false;
  }
  Base::AddSuffix(block->UsableSpace(), block->InnerSize());
  return true;
}
 
template <typename BlockAllocatorType, typename LockType>
Result<Layout> GuardedAllocator<BlockAllocatorType, LockType>::DoGetInfo(
    InfoType info_type, const void* ptr) const {
  ptr = const_cast<const void*>(Base::GetOriginal(const_cast<void*>(ptr)));
  auto allocator = borrowable_.acquire();
  return BlockAllocatorType::GetInfo(*allocator, info_type, ptr);
}
 
}  // namespace pw::allocator