fast_sorted.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/bucket/base.h"
#include "pw_containers/intrusive_multimap.h"
#include "pw_function/function.h"
 
namespace pw::allocator {
 
template <typename BlockType>
class FastSortedItem
    : public IntrusiveMultiMap<size_t, FastSortedItem<BlockType>>::Item {
 public:
  size_t key() const {
    const auto* block = BlockType::FromUsableSpace(this);
    return block->InnerSize();
  }
};
 
struct GenericFastSortedItem
    : public IntrusiveMultiMap<size_t, GenericFastSortedItem>::Item {};
 
template <typename BlockType>
class FastSortedBucket
    : public internal::BucketBase<FastSortedBucket<BlockType>,
                                  BlockType,
                                  FastSortedItem<BlockType>> {
 private:
  using Base = internal::BucketBase<FastSortedBucket<BlockType>,
                                    BlockType,
                                    FastSortedItem<BlockType>>;
  friend Base;
 
  template <typename>
  friend class ReverseFastSortedBucket;
 
  using Compare = Function<bool(size_t, size_t)>;
 
 public:
  constexpr FastSortedBucket() = default;
  ~FastSortedBucket() { Base::Clear(); }
 
 private:
  // Constructor used by `ReverseFastSortedBucket`.
  constexpr explicit FastSortedBucket(Compare&& compare)
      : items_(std::move(compare)) {}
 
  void DoAdd(BlockType& block) {
    auto* item = new (block.UsableSpace()) FastSortedItem<BlockType>();
    items_.insert(*item);
  }
 
  BlockType* DoRemoveAny() {
    auto iter = items_.begin();
    FastSortedItem<BlockType>& item = *iter;
    items_.erase(iter);
    return BlockType::FromUsableSpace(&item);
  }
 
  bool DoRemove(BlockType& block) {
    FastSortedItem<BlockType>& item_to_remove = Base::GetItemFrom(block);
    auto iters = items_.equal_range(block.InnerSize());
    auto iter =
        std::find_if(iters.first,
                     iters.second,
                     [&item_to_remove](FastSortedItem<BlockType>& item) {
                       return &item_to_remove == &item;
                     });
    if (iter == items_.end()) {
      return false;
    }
    items_.erase(iter);
    return true;
  }
 
  BlockType* DoRemoveCompatible(Layout layout) {
    auto iter = items_.lower_bound(layout.size());
    return RemoveImpl(iter, layout);
  }
 
  template <typename Iterator>
  BlockType* RemoveImpl(Iterator iter, Layout layout) {
    iter =
        std::find_if(iter, items_.end(), Base::MakeCanAllocPredicate(layout));
    auto* block = Base::GetBlockFromIterator(iter, items_.end());
    if (block != nullptr) {
      items_.erase(iter);
    }
    return block;
  }
 
  IntrusiveMultiMap<size_t, FastSortedItem<BlockType>> items_;
};
 
template <typename BlockType>
class ReverseFastSortedBucket
    : public internal::BucketBase<ReverseFastSortedBucket<BlockType>,
                                  BlockType,
                                  FastSortedItem<BlockType>> {
 private:
  using Base = internal::BucketBase<ReverseFastSortedBucket<BlockType>,
                                    BlockType,
                                    FastSortedItem<BlockType>>;
  friend Base;
 
 public:
  constexpr ReverseFastSortedBucket()
      : impl_(std::greater<>()), items_(impl_.items_) {}
 
 private:
  void DoAdd(BlockType& block) { impl_.DoAdd(block); }
 
  BlockType* DoRemoveAny() {
    auto iter = items_.begin();
    FastSortedItem<BlockType>& item = *iter;
    items_.erase(iter);
    return BlockType::FromUsableSpace(&item);
  }
 
  bool DoRemove(BlockType& block) { return impl_.DoRemove(block); }
 
  BlockType* DoRemoveCompatible(Layout layout) {
    return impl_.RemoveImpl(impl_.items_.begin(), layout);
  }
 
  FastSortedBucket<BlockType> impl_;
  IntrusiveMultiMap<size_t, FastSortedItem<BlockType>>& items_;
};
 
}  // namespace pw::allocator