bucket_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 <array>
#include <cstdint>
#include <utility>
 
#include "pw_allocator/block/detailed_block.h"
#include "pw_allocator/block_allocator.h"
#include "pw_allocator/bucket/unordered.h"
#include "pw_status/try.h"
 
namespace pw::allocator {
 
 
template <typename OffsetType = uintptr_t>
using BucketBlock = DetailedBlock<OffsetType, UnorderedItem>;
 
template <typename BlockType = BucketBlock<>>
class BasicBucketAllocator : public BlockAllocator<BlockType> {
 private:
  using Base = BlockAllocator<BlockType>;
 
 public:
  constexpr BasicBucketAllocator() = default;
 
  constexpr void InitBuckets(span<UnorderedBucket<BlockType>> buckets,
                             size_t max_inner_size);
 
 private:
  size_t DoGetMaxAllocatable() override;
 
  BlockResult<BlockType> ChooseBlock(Layout layout) override;
 
  void ReserveBlock(BlockType& block) override;
 
  void RecycleBlock(BlockType& block) override;
 
  span<UnorderedBucket<BlockType>> buckets_;
};
 
template <typename BlockType = BucketBlock<>,
          size_t kMinInnerSize = 32,
          size_t kNumBuckets = 5>
class BucketAllocator : public BasicBucketAllocator<BlockType> {
 private:
  using Base = BasicBucketAllocator<BlockType>;
 
 public:
  constexpr BucketAllocator() { Base::InitBuckets(buckets_, kMinInnerSize); }
 
  explicit BucketAllocator(ByteSpan region) : BucketAllocator() {
    Base::Init(region);
  }
 
  ~BucketAllocator() override;
 
 private:
  std::array<UnorderedBucket<BlockType>, kNumBuckets> buckets_;
};
 
 
// Template method implementations.
 
template <typename BlockType>
constexpr void BasicBucketAllocator<BlockType>::InitBuckets(
    span<UnorderedBucket<BlockType>> buckets, size_t max_inner_size) {
  buckets_ = buckets;
  for (auto& bucket : buckets_.subspan(0, buckets_.size() - 1)) {
    bucket.set_max_inner_size(max_inner_size);
    max_inner_size <<= 1;
  }
}
 
template <typename BlockType>
size_t BasicBucketAllocator<BlockType>::DoGetMaxAllocatable() {
  for (auto b = buckets_.rbegin(); b != buckets_.rend(); ++b) {
    const BlockType* largest = b->FindLargest();
    if (largest != nullptr) {
      return largest->InnerSize();
    }
  }
  return 0;
}
 
template <typename BlockType>
BlockResult<BlockType> BasicBucketAllocator<BlockType>::ChooseBlock(
    Layout layout) {
  for (auto& bucket : buckets_) {
    if (bucket.max_inner_size() < layout.size()) {
      continue;
    }
    BlockType* block = bucket.RemoveCompatible(layout);
    if (block != nullptr) {
      return BlockType::AllocFirst(std::move(block), layout);
    }
  }
  return BlockResult<BlockType>(nullptr, Status::NotFound());
}
 
template <typename BlockType>
void BasicBucketAllocator<BlockType>::ReserveBlock(BlockType& block) {
  for (auto& bucket : buckets_) {
    if (block.InnerSize() <= bucket.max_inner_size()) {
      std::ignore = bucket.Remove(block);
      break;
    }
  }
}
 
template <typename BlockType>
void BasicBucketAllocator<BlockType>::RecycleBlock(BlockType& block) {
  for (auto& bucket : buckets_) {
    if (block.InnerSize() <= bucket.max_inner_size()) {
      std::ignore = bucket.Add(block);
      break;
    }
  }
}
 
template <typename BlockType, size_t kMinInnerSize, size_t kNumBuckets>
BucketAllocator<BlockType, kMinInnerSize, kNumBuckets>::~BucketAllocator() {
  for (auto& bucket : buckets_) {
    bucket.Clear();
  }
}
 
}  // namespace pw::allocator