base.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 <limits>
 
#include "pw_allocator/block/poisonable.h"
#include "pw_allocator/config.h"
#include "pw_allocator/hardening.h"
#include "pw_allocator/layout.h"
#include "pw_assert/assert.h"
#include "pw_bytes/alignment.h"
 
namespace pw::allocator::internal {
 
template <typename Derived, typename BlockType_, typename ItemType_>
class BucketBase {
 public:
  using BlockType = BlockType_;
  using ItemType = ItemType_;
 
  constexpr BucketBase() {
    if constexpr (is_poisonable_v<BlockType>) {
      static_assert(BlockType::kPoisonOffset >= sizeof(ItemType),
                    "Block type does not reserve sufficient space for an item");
    }
  }
 
  constexpr bool empty() const { return derived()->items_.empty(); }
 
  constexpr size_t max_inner_size() const { return max_inner_size_; }
 
  constexpr void set_max_inner_size(size_t max_inner_size) {
    if constexpr (Hardening::kIncludesDebugChecks) {
      PW_ASSERT(empty());
    }
    max_inner_size_ = max_inner_size;
  }
 
  [[nodiscard]] bool Add(BlockType& block) {
    if (block.InnerSize() < sizeof(ItemType)) {
      return false;
    }
    if constexpr (Hardening::kIncludesDebugChecks) {
      PW_ASSERT(block.InnerSize() <= max_inner_size_);
      PW_ASSERT(IsAlignedAs<ItemType>(block.UsableSpace()));
    }
    derived()->DoAdd(block);
    return true;
  }
 
  [[nodiscard]] BlockType* RemoveAny() {
    return empty() ? nullptr : derived()->DoRemoveAny();
  }
 
  [[nodiscard]] bool Remove(BlockType& block) {
    return block.InnerSize() >= sizeof(ItemType) && derived()->DoRemove(block);
  }
 
  [[nodiscard]] BlockType* RemoveCompatible(Layout layout) {
    return derived()->DoRemoveCompatible(layout);
  }
 
  void Clear() { derived()->items_.clear(); }
 
  template <typename Iterator, typename Predicate>
  static Iterator FindPrevIf(Iterator before_first,
                             Iterator last,
                             Predicate predicate) {
    Iterator prev = before_first;
    Iterator iter = prev;
    ++iter;
    for (; iter != last; ++iter) {
      if (predicate(*iter)) {
        break;
      }
      prev = iter;
    }
    return prev;
  }
 
  static auto MakeCanAllocPredicate(Layout layout) {
    return [layout](ItemType& item) {
      auto* block = BlockType::FromUsableSpace(&item);
      return block->CanAlloc(layout).ok();
    };
  }
 
  template <typename Iterator>
  constexpr BlockType* GetBlockFromIterator(Iterator iter, Iterator last) {
    return iter != last ? BlockType::FromUsableSpace(&(*iter)) : nullptr;
  }
 
  template <typename Iterator>
  constexpr BlockType* GetBlockFromPrev(Iterator prev, Iterator last) {
    return GetBlockFromIterator(++prev, last);
  }
 
 protected:
  static ItemType& GetItemFrom(BlockType& block) {
    return *(std::launder(reinterpret_cast<ItemType*>(block.UsableSpace())));
  }
 
 private:
  constexpr Derived* derived() { return static_cast<Derived*>(this); }
  constexpr const Derived* derived() const {
    return static_cast<const Derived*>(this);
  }
 
  size_t max_inner_size_ = std::numeric_limits<size_t>::max();
};
 
template <typename T, typename U = size_t>
constexpr U CountRZero(T t) {
  return static_cast<U>(cpp20::countr_zero(t));
}
 
template <typename T, typename U = size_t>
constexpr U CountLZero(T t) {
  return static_cast<U>(cpp20::countl_zero(t));
}
 
}  // namespace pw::allocator::internal