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<>,
          size_t kMinInnerSize = 32,
          size_t kNumBuckets = 5>
class BucketAllocator : public BlockAllocator<BlockType> {
 private:
  using Base = BlockAllocator<BlockType>;
 
 public:
  constexpr BucketAllocator() {
    size_t max_inner_size = kMinInnerSize;
    for (auto& bucket : span(buckets_.data(), kNumBuckets - 1)) {
      bucket.set_max_inner_size(max_inner_size);
      max_inner_size <<= 1;
    }
  }
 
  explicit BucketAllocator(ByteSpan region) : BucketAllocator() {
    Base::Init(region);
  }
 
  ~BucketAllocator() override {
    for (auto& bucket : buckets_) {
      bucket.Clear();
    }
  }
 
 private:
  BlockResult<BlockType> ChooseBlock(Layout layout) override {
    for (auto& bucket : buckets_) {
      if (layout.size() <= bucket.max_inner_size()) {
        BlockType* block = bucket.RemoveCompatible(layout);
        if (block != nullptr) {
          return BlockType::AllocFirst(std::move(block), layout);
        }
      }
    }
    return BlockResult<BlockType>(nullptr, Status::NotFound());
  }
 
  void ReserveBlock(BlockType& block) override {
    for (auto& bucket : buckets_) {
      if (block.InnerSize() <= bucket.max_inner_size()) {
        std::ignore = bucket.Remove(block);
        break;
      }
    }
  }
 
  void RecycleBlock(BlockType& block) override {
    for (auto& bucket : buckets_) {
      if (block.InnerSize() <= bucket.max_inner_size()) {
        std::ignore = bucket.Add(block);
        break;
      }
    }
  }
 
  std::array<UnorderedBucket<BlockType>, kNumBuckets> buckets_;
};
 
}  // namespace pw::allocator