bucket_block_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 "pw_allocator/block_allocator.h"
#include "pw_allocator/bucket.h"
#include "pw_assert/assert.h"
#include "pw_status/try.h"
 
namespace pw::allocator {
 
template <typename OffsetType = uintptr_t,
          size_t kMinBucketChunkSize = 32,
          size_t kNumBuckets = 5,
          size_t kAlign = std::max(alignof(OffsetType), alignof(std::byte*))>
class BucketBlockAllocator
    : public BlockAllocator<OffsetType,
                            0,
                            std::max(kAlign, alignof(std::byte*))> {
 public:
  using Base =
      BlockAllocator<OffsetType, 0, std::max(kAlign, alignof(std::byte*))>;
  using BlockType = typename Base::BlockType;
 
  constexpr BucketBlockAllocator() : Base() {}
 
  explicit BucketBlockAllocator(ByteSpan region) : BucketBlockAllocator() {
    Base::Init(region);
  }
 
  void Init(ByteSpan region) { Base::Init(region); }
 
  void Init(BlockType* begin) { Base::Init(begin); }
 
  void Init(BlockType* begin, BlockType* end) override {
    Base::Init(begin, end);
    internal::Bucket::Init(span(buckets_.data(), buckets_.size() - 1),
                           kMinBucketChunkSize);
    buckets_.back().Init();
    for (auto* block : Base::blocks()) {
      if (!block->Used()) {
        RecycleBlock(block);
      }
    }
  }
 
 private:
  BlockType* ChooseBlock(Layout layout) override {
    layout =
        Layout(std::max(layout.size(), sizeof(internal::Bucket::Chunk)),
               std::max(layout.alignment(), alignof(internal::Bucket::Chunk)));
    for (auto& bucket : buckets_) {
      if (bucket.chunk_size() < layout.size()) {
        continue;
      }
      void* chosen = bucket.RemoveIf([&layout](const std::byte* chunk) {
        const BlockType* block = BlockType::FromUsableSpace(chunk);
        return block->CanAllocLast(layout).ok();
      });
      if (chosen == nullptr) {
        continue;
      }
      BlockType* block = BlockType::FromUsableSpace(chosen);
      BlockType* prev = block;
      BlockType::AllocLast(block, layout).IgnoreError();
 
      // If the chunk was split, what we have is the leading free block.
      if (prev != block) {
        RecycleBlock(prev);
      }
      return block;
    }
    return nullptr;
  }
 
  void ReserveBlock(BlockType* block) override {
    PW_ASSERT(!block->Used());
    size_t inner_size = block->InnerSize();
    if (inner_size < sizeof(internal::Bucket::Chunk)) {
      return;
    }
    internal::Bucket::Remove(block->UsableSpace());
  }
 
  void RecycleBlock(BlockType* block) override {
    PW_ASSERT(!block->Used());
    size_t inner_size = block->InnerSize();
    if (inner_size < sizeof(internal::Bucket::Chunk)) {
      return;
    }
    for (auto& bucket : buckets_) {
      if (inner_size <= bucket.chunk_size()) {
        bucket.Add(block->UsableSpace());
        break;
      }
    }
  }
 
  std::array<internal::Bucket, kNumBuckets> buckets_;
};
 
}  // namespace pw::allocator