simple_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 "pw_allocator/allocator.h"
#include "pw_containers/intrusive_list.h"
#include "pw_multibuf/allocator.h"
#include "pw_multibuf/multibuf.h"
 
namespace pw::multibuf {
 
class SimpleAllocator;
 
namespace internal {
 
class LinkedRegionTracker final
    : public ChunkRegionTracker,
      public IntrusiveList<LinkedRegionTracker>::Item {
 public:
  LinkedRegionTracker(SimpleAllocator& parent, ByteSpan region)
      : parent_(parent), region_(region) {}
  ~LinkedRegionTracker() final;
 
  // LinkedRegionTracker is not copyable nor movable.
  LinkedRegionTracker(const LinkedRegionTracker&) = delete;
  LinkedRegionTracker& operator=(const LinkedRegionTracker&) = delete;
  LinkedRegionTracker(LinkedRegionTracker&&) = delete;
  LinkedRegionTracker& operator=(LinkedRegionTracker&&) = delete;
 
 protected:
  void Destroy() final;
  ByteSpan Region() const final { return region_; }
  void* AllocateChunkClass() final;
  void DeallocateChunkClass(void*) final;
 
 private:
  SimpleAllocator& parent_;
  const ByteSpan region_;
 
  friend class ::pw::multibuf::SimpleAllocator;
};
 
}  // namespace internal
 
class SimpleAllocator : public MultiBufAllocator {
 public:
  SimpleAllocator(ByteSpan data_area,
                  pw::allocator::Allocator& metadata_alloc,
                  size_t alignment = 1);
 
 private:
  pw::Result<MultiBuf> DoAllocate(
      size_t min_size,
      size_t desired_size,
      ContiguityRequirement contiguity_requirement) final;
 
  std::optional<size_t> DoGetBackingCapacity() final {
    return data_area_.size();
  }
 
  pw::Result<MultiBuf> InternalAllocateContiguous(size_t size)
      PW_EXCLUSIVE_LOCKS_REQUIRED(lock_);
 
  struct FreeBlock final {
    IntrusiveList<internal::LinkedRegionTracker>::iterator iter;
 
    ByteSpan span;
  };
 
  pw::Result<OwnedChunk> InsertRegion(const FreeBlock&)
      PW_EXCLUSIVE_LOCKS_REQUIRED(lock_);
 
  struct AvailableMemorySize final {
    size_t total;
    size_t contiguous;
  };
 
  AvailableMemorySize GetAvailableMemorySize()
      PW_EXCLUSIVE_LOCKS_REQUIRED(lock_);
 
  enum class ControlFlow {
    Continue,
    Break,
  };
 
  template <typename Fn>
  void ForEachFreeBlock(Fn function) PW_EXCLUSIVE_LOCKS_REQUIRED(lock_) {
    std::byte* last_used_end = data_area_.data();
    // We need to track only the `prev_iter` in order to ensure that we don't
    // miss any blocks that ``function`` inserts.
    auto prev_iter = regions_.before_begin();
    while (true) {
      // Compute ``cur_iter`` by incrementing ``prev_iter``.
      auto cur_iter = prev_iter;
      cur_iter++;
      if (cur_iter == regions_.end()) {
        break;
      }
      size_t unused = cur_iter->region_.data() - last_used_end;
      if (unused != 0) {
        ControlFlow cf = function({prev_iter, ByteSpan(last_used_end, unused)});
        if (cf == ControlFlow::Break) {
          return;
        }
      }
      last_used_end = cur_iter->region_.data() + cur_iter->region_.size();
      prev_iter = cur_iter;
    }
    size_t unused = (data_area_.data() + data_area_.size()) - last_used_end;
    if (unused != 0) {
      function({prev_iter, ByteSpan(last_used_end, unused)});
    }
  }
 
  pw::sync::Mutex lock_;
  IntrusiveList<internal::LinkedRegionTracker> regions_ PW_GUARDED_BY(lock_);
  pw::allocator::Allocator& metadata_alloc_;
  const ByteSpan data_area_;
  const size_t alignment_;
 
  friend class internal::LinkedRegionTracker;
};
 
}  // namespace pw::multibuf