multibuf.h

// Copyright 2023 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 <iterator>
#include <tuple>
 
#include "pw_multibuf/chunk.h"
#include "pw_preprocessor/compiler.h"
#include "pw_status/status_with_size.h"
 
namespace pw::multibuf {
 
class MultiBuf;
 
class MultiBufChunks {
 public:
  using element_type = Chunk;
  using value_type = Chunk;
  using pointer = Chunk*;
  using reference = Chunk&;
  using const_pointer = const Chunk*;
  using difference_type = std::ptrdiff_t;
  using const_reference = const Chunk&;
  using size_type = std::size_t;
 
  class iterator {
   public:
    using value_type = Chunk;
    using difference_type = std::ptrdiff_t;
    using reference = Chunk&;
    using pointer = Chunk*;
    using iterator_category = std::forward_iterator_tag;
 
    constexpr iterator() = default;
 
    constexpr reference operator*() const { return *chunk_; }
    constexpr pointer operator->() const { return chunk_; }
 
    constexpr iterator& operator++() {
      chunk_ = chunk_->next_in_buf_;
      return *this;
    }
 
    constexpr iterator operator++(int) {
      iterator tmp = *this;
      ++(*this);
      return tmp;
    }
 
    constexpr bool operator==(const iterator& other) const {
      return chunk_ == other.chunk_;
    }
 
    constexpr bool operator!=(const iterator& other) const {
      return chunk_ != other.chunk_;
    }
 
   private:
    friend class MultiBufChunks;
    friend class MultiBuf;
 
    constexpr iterator(Chunk* chunk) : chunk_(chunk) {}
    static constexpr iterator end() { return iterator(nullptr); }
    Chunk* chunk_ = nullptr;
  };
 
  class const_iterator {
   public:
    using value_type = const Chunk;
    using difference_type = std::ptrdiff_t;
    using reference = const Chunk&;
    using pointer = const Chunk*;
    using iterator_category = std::forward_iterator_tag;
 
    constexpr const_iterator() = default;
 
    constexpr reference operator*() const { return *chunk_; }
    constexpr pointer operator->() const { return chunk_; }
 
    constexpr const_iterator& operator++() {
      chunk_ = chunk_->next_in_buf_;
      return *this;
    }
 
    constexpr const_iterator operator++(int) {
      const_iterator tmp = *this;
      ++(*this);
      return tmp;
    }
 
    constexpr bool operator==(const const_iterator& other) const {
      return chunk_ == other.chunk_;
    }
 
    constexpr bool operator!=(const const_iterator& other) const {
      return chunk_ != other.chunk_;
    }
 
   private:
    friend class MultiBufChunks;
    friend class MultiBuf;
 
    constexpr const_iterator(const Chunk* chunk) : chunk_(chunk) {}
    static constexpr const_iterator end() { return const_iterator(nullptr); }
    const Chunk* chunk_ = nullptr;
  };
 
  MultiBufChunks(const MultiBufChunks&) = delete;
  MultiBufChunks& operator=(const MultiBufChunks&) = delete;
 
  constexpr Chunk& front() { return *first_; }
  constexpr const Chunk& front() const { return *first_; }
 
  Chunk& back() { return const_cast<Chunk&>(std::as_const(*this).back()); }
  const Chunk& back() const;
 
  constexpr iterator begin() { return iterator(first_); }
  constexpr const_iterator begin() const { return cbegin(); }
  constexpr const_iterator cbegin() const { return const_iterator(first_); }
 
  constexpr iterator end() { return iterator::end(); }
  constexpr const_iterator end() const { return cend(); }
  constexpr const_iterator cend() const { return const_iterator::end(); }
 
  size_t size() const {
    return static_cast<size_t>(std::distance(begin(), end()));
  }
 
  size_t size_bytes() const;
 
  [[nodiscard]] bool empty() const { return first_ == nullptr; }
 
  void push_front(OwnedChunk&& chunk);
 
  void push_back(OwnedChunk&& chunk);
 
  OwnedChunk take_front();
 
  //
  // Implementation note: `Chunks().size()` should be remain relatively small,
  // but this could be made `O(1)` in the future by adding a `prev` pointer to
  // the `iterator`.
  iterator insert(iterator position, OwnedChunk&& chunk);
 
  //
  // Implementation note: `Chunks().size()` should be remain relatively small,
  // but this could be made `O(1)` in the future by adding a `prev` pointer to
  // the `iterator`.
  std::tuple<iterator, OwnedChunk> take(iterator position);
 
 protected:
  explicit constexpr MultiBufChunks(Chunk* first_chunk) : first_(first_chunk) {}
 
  ~MultiBufChunks() { Release(); }
 
  // Disable maybe-uninitialized: this check fails erroneously on Windows GCC.
  PW_MODIFY_DIAGNOSTICS_PUSH();
  PW_MODIFY_DIAGNOSTIC_GCC(ignored, "-Wmaybe-uninitialized");
  constexpr MultiBufChunks(MultiBufChunks&& other) noexcept
      : first_(other.first_) {
    other.first_ = nullptr;
  }
  PW_MODIFY_DIAGNOSTICS_POP();
 
  MultiBufChunks& operator=(MultiBufChunks&& other) noexcept {
    Release();
    first_ = other.first_;
    other.first_ = nullptr;
    return *this;
  }
 
  // Releases all chunks in the `MultiBuf`.
  void Release() noexcept;
 
  void PushSuffix(MultiBufChunks&& tail);
 
  Chunk* Previous(Chunk* chunk) const;
 
 private:
  Chunk* first_;
};
 
class MultiBuf : private MultiBufChunks {
 public:
  class iterator;
  class const_iterator;
 
  constexpr MultiBuf() : MultiBufChunks(nullptr) {}
 
  static MultiBuf FromChunk(OwnedChunk&& chunk) {
    return MultiBuf(std::move(chunk).Take());
  }
 
  MultiBuf(const MultiBuf&) = delete;
  MultiBuf& operator=(const MultiBuf&) = delete;
 
  constexpr MultiBuf(MultiBuf&& other) noexcept = default;
 
  MultiBuf& operator=(MultiBuf&& other) noexcept = default;
 
  void Release() noexcept { MultiBufChunks::Release(); }
 
  ~MultiBuf() = default;
 
  [[nodiscard]] size_t size() const { return MultiBufChunks::size_bytes(); }
 
  [[nodiscard]] bool empty() const;
 
  [[nodiscard]] bool IsContiguous() const {
    return ContiguousSpan().has_value();
  }
 
  std::optional<ByteSpan> ContiguousSpan() {
    auto result = std::as_const(*this).ContiguousSpan();
    if (result.has_value()) {
      return span(const_cast<std::byte*>(result->data()), result->size());
    }
    return std::nullopt;
  }
  std::optional<ConstByteSpan> ContiguousSpan() const;
 
  iterator begin() { return iterator(Chunks().begin().chunk_); }
  const_iterator begin() const {
    return const_iterator(Chunks().begin().chunk_);
  }
  const_iterator cbegin() const {
    return const_iterator(Chunks().begin().chunk_);
  }
 
  iterator end() { return iterator::end(); }
  const_iterator end() const { return const_iterator::end(); }
  const_iterator cend() const { return const_iterator::end(); }
 
  [[nodiscard]] bool ClaimPrefix(size_t bytes_to_claim);
 
  [[nodiscard]] bool ClaimSuffix(size_t bytes_to_claim);
 
  void DiscardPrefix(size_t bytes_to_discard);
 
  void Slice(size_t begin, size_t end);
 
  void Truncate(size_t len);
 
  void TruncateAfter(iterator pos);
 
  std::optional<MultiBuf> TakePrefix(size_t bytes_to_take);
 
  std::optional<MultiBuf> TakeSuffix(size_t bytes_to_take);
 
  void PushPrefix(MultiBuf&& front);
 
  void PushSuffix(MultiBuf&& tail) {
    return MultiBufChunks::PushSuffix(std::move(tail.Chunks()));
  }
 
  StatusWithSize CopyTo(ByteSpan dest, size_t position = 0) const;
 
  StatusWithSize CopyFrom(ConstByteSpan source, size_t position = 0) {
    return CopyFromAndOptionallyTruncate(source, position, /*truncate=*/false);
  }
 
  StatusWithSize CopyFromAndTruncate(ConstByteSpan source,
                                     size_t position = 0) {
    return CopyFromAndOptionallyTruncate(source, position, /*truncate=*/true);
  }
 
  //--------------------- Chunk manipulation ----------------------//
 
  void PushFrontChunk(OwnedChunk&& chunk) {
    MultiBufChunks::push_front(std::move(chunk));
  }
 
  void PushBackChunk(OwnedChunk&& chunk) {
    MultiBufChunks::push_back(std::move(chunk));
  }
 
  OwnedChunk TakeFrontChunk() { return MultiBufChunks::take_front(); }
 
  MultiBufChunks::iterator InsertChunk(MultiBufChunks::iterator position,
                                       OwnedChunk&& chunk) {
    return MultiBufChunks::insert(position, std::move(chunk));
  }
 
  std::tuple<MultiBufChunks::iterator, OwnedChunk> TakeChunk(
      MultiBufChunks::iterator position) {
    return MultiBufChunks::take(position);
  }
 
  constexpr MultiBufChunks& Chunks() { return *this; }
 
  constexpr const MultiBufChunks& Chunks() const { return *this; }
 
  constexpr const MultiBufChunks& ConstChunks() const { return *this; }
 
  //--------------------- Iterator details ------------------------//
 
  using element_type = std::byte;
  using value_type = std::byte;
  using pointer = std::byte*;
  using const_pointer = const std::byte*;
  using reference = std::byte&;
  using const_reference = const std::byte&;
  using difference_type = std::ptrdiff_t;
  using size_type = std::size_t;
 
  class const_iterator {
   public:
    using value_type = std::byte;
    using difference_type = std::ptrdiff_t;
    using reference = const std::byte&;
    using pointer = const std::byte*;
    using iterator_category = std::forward_iterator_tag;
 
    constexpr const_iterator() : chunk_(nullptr), byte_index_(0) {}
 
    reference operator*() const { return (*chunk_)[byte_index_]; }
    pointer operator->() const { return &(*chunk_)[byte_index_]; }
 
    const_iterator& operator++();
    const_iterator operator++(int) {
      const_iterator tmp = *this;
      ++(*this);
      return tmp;
    }
    const_iterator operator+(size_t rhs) const {
      const_iterator tmp = *this;
      tmp += rhs;
      return tmp;
    }
    const_iterator& operator+=(size_t advance);
 
    constexpr bool operator==(const const_iterator& other) const {
      return chunk_ == other.chunk_ && byte_index_ == other.byte_index_;
    }
 
    constexpr bool operator!=(const const_iterator& other) const {
      return !(*this == other);
    }
 
    constexpr const Chunk* chunk() const { return chunk_; }
 
    constexpr size_t byte_index() const { return byte_index_; }
 
   private:
    friend class MultiBuf;
 
    explicit constexpr const_iterator(const Chunk* chunk, size_t byte_index = 0)
        : chunk_(chunk), byte_index_(byte_index) {
      AdvanceToData();
    }
 
    static const_iterator end() { return const_iterator(nullptr); }
 
    constexpr void AdvanceToData() {
      while (chunk_ != nullptr && chunk_->empty()) {
        chunk_ = chunk_->next_in_buf_;
      }
    }
 
    const Chunk* chunk_;
    size_t byte_index_;
  };
 
  class iterator {
   public:
    using value_type = std::byte;
    using difference_type = std::ptrdiff_t;
    using reference = std::byte&;
    using pointer = std::byte*;
    using iterator_category = std::forward_iterator_tag;
 
    constexpr iterator() = default;
 
    reference operator*() const { return const_cast<std::byte&>(*const_iter_); }
    pointer operator->() const { return const_cast<std::byte*>(&*const_iter_); }
 
    iterator& operator++() {
      const_iter_++;
      return *this;
    }
    iterator operator++(int) {
      iterator tmp = *this;
      ++(*this);
      return tmp;
    }
    iterator operator+(size_t rhs) const {
      iterator tmp = *this;
      tmp += rhs;
      return tmp;
    }
    iterator& operator+=(size_t rhs) {
      const_iter_ += rhs;
      return *this;
    }
    constexpr bool operator==(const iterator& other) const {
      return const_iter_ == other.const_iter_;
    }
    constexpr bool operator!=(const iterator& other) const {
      return const_iter_ != other.const_iter_;
    }
 
    constexpr Chunk* chunk() const {
      return const_cast<Chunk*>(const_iter_.chunk());
    }
 
    constexpr size_t byte_index() const { return const_iter_.byte_index(); }
 
   private:
    friend class MultiBuf;
 
    explicit constexpr iterator(Chunk* chunk, size_t byte_index = 0)
        : const_iter_(chunk, byte_index) {}
 
    static iterator end() { return iterator(nullptr); }
 
    const_iterator const_iter_;
  };
 
 private:
  explicit constexpr MultiBuf(Chunk* first_chunk)
      : MultiBufChunks(first_chunk) {}
 
  StatusWithSize CopyFromAndOptionallyTruncate(ConstByteSpan source,
                                               size_t position,
                                               bool truncate);
};
 
}  // namespace pw::multibuf