1 // Copyright 2024 The Chromium Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #ifndef BASE_CONTAINERS_SPAN_WRITER_H_
6 #define BASE_CONTAINERS_SPAN_WRITER_H_
7
8 #include <optional>
9
10 #include "base/containers/span.h"
11 #include "base/memory/raw_span.h"
12 #include "base/numerics/byte_conversions.h"
13
14 namespace base {
15
16 // A Writer to write into and consume elements from the front of a span
17 // dynamically.
18 //
19 // SpanWriter is used to split off prefix spans from a larger span, reporting
20 // errors if there's not enough room left (instead of crashing, as would happen
21 // with span directly).
22 template <typename T>
23 class SpanWriter {
24 static_assert(!std::is_const_v<T>,
25 "SpanWriter needs mutable access to its buffer");
26
27 public:
28 // Construct SpanWriter that writes to `buf`.
SpanWriter(span<T> buf)29 constexpr explicit SpanWriter(span<T> buf)
30 : buf_(buf), original_size_(buf_.size()) {}
31
32 // Returns true and writes the span `data` into the front of the inner span,
33 // if there is enough room left. Otherwise, it returns false and does
34 // nothing.
Write(span<const T> data)35 constexpr bool Write(span<const T> data) {
36 if (data.size() > remaining()) {
37 return false;
38 }
39 auto [lhs, rhs] = buf_.split_at(data.size());
40 lhs.copy_from(data);
41 buf_ = rhs;
42 return true;
43 }
44
45 // Returns true and writes `value` into the front of the inner span if there
46 // is space remaining. Otherwise, it returns false and does nothing.
47 template <typename V>
requires(std::same_as<T,std::remove_cvref_t<V>>)48 requires(std::same_as<T, std::remove_cvref_t<V>>)
49 bool Write(V&& value) {
50 if (!remaining()) {
51 return false;
52 }
53 buf_[0] = std::forward<V>(value);
54 buf_ = buf_.last(remaining() - 1);
55 return true;
56 }
57
58 // Skips over the next `n` objects, and returns a span that points to the
59 // skipped objects, if there are enough objects left. Otherwise, it returns
60 // nullopt and does nothing.
Skip(StrictNumeric<size_t> n)61 constexpr std::optional<span<T>> Skip(StrictNumeric<size_t> n) {
62 if (n > remaining()) {
63 return std::nullopt;
64 }
65 auto [lhs, rhs] = buf_.split_at(n);
66 buf_ = rhs;
67 return lhs;
68 }
69 template <size_t N>
Skip()70 constexpr std::optional<span<T, N>> Skip() {
71 if (N > remaining()) {
72 return std::nullopt;
73 }
74 auto [lhs, rhs] = buf_.template split_at<N>();
75 buf_ = rhs;
76 return lhs;
77 }
78
79 // For a SpanWriter over bytes, we can write integer values directly to those
80 // bytes as a memcpy. Returns true if there was room remaining and the bytes
81 // were written.
82 //
83 // This provides big, little, and native endian writing orders. Note that
84 // "native" order is almost never what you want; it only makes sense for byte
85 // buffers that stay in memory and are never written to the disk or network.
86 #define BASE_SPANWRITER_WRITE(signchar, bitsize, endian, typeprefix) \
87 constexpr bool Write##signchar##bitsize##endian##Endian( \
88 typeprefix##int##bitsize##_t value) \
89 requires(std::same_as<T, uint8_t>) \
90 { \
91 return Write(signchar##bitsize##To##endian##Endian(value)); \
92 }
93 #define BASE_SPANWRITER_WRITE_BOTH_SIGNS(bitsize, endian) \
94 BASE_SPANWRITER_WRITE(U, bitsize, endian, u) \
95 BASE_SPANWRITER_WRITE(I, bitsize, endian, )
96 #define BASE_SPANWRITER_WRITE_BOTH_SIGNS_ALL_SIZES(endian) \
97 BASE_SPANWRITER_WRITE_BOTH_SIGNS(8, endian) \
98 BASE_SPANWRITER_WRITE_BOTH_SIGNS(16, endian) \
99 BASE_SPANWRITER_WRITE_BOTH_SIGNS(32, endian) \
100 BASE_SPANWRITER_WRITE_BOTH_SIGNS(64, endian)
101
102 BASE_SPANWRITER_WRITE_BOTH_SIGNS_ALL_SIZES(Big)
BASE_SPANWRITER_WRITE_BOTH_SIGNS_ALL_SIZES(Little)103 BASE_SPANWRITER_WRITE_BOTH_SIGNS_ALL_SIZES(Little)
104 BASE_SPANWRITER_WRITE_BOTH_SIGNS_ALL_SIZES(Native)
105
106 #undef BASE_SPANWRITER_WRITE_BOTH_SIGNS_ALL_SIZES
107 #undef BASE_SPANWRITER_WRITE_BOTH_SIGNS
108 #undef BASE_SPANWRITER_WRITE
109
110 // Returns the remaining not-yet-written-to object count.
111 constexpr size_t remaining() const { return buf_.size(); }
112
113 // Returns the remaining not-yet-written-to objects.
remaining_span()114 constexpr span<T> remaining_span() const { return buf_; }
115
116 // Returns the number of objects already written (or skipped).
num_written()117 constexpr size_t num_written() const { return original_size_ - buf_.size(); }
118
119 private:
120 raw_span<T> buf_;
121 size_t original_size_;
122 };
123
124 template <class T, size_t N>
125 SpanWriter(span<T, N>) -> SpanWriter<T>;
126
127 } // namespace base
128
129 #endif // BASE_CONTAINERS_SPAN_WRITER_H_
130