nautilus_serialization/sbe/writer.rs
1// -------------------------------------------------------------------------------------------------
2// Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3// https://nautechsystems.io
4//
5// Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6// You may not use this file except in compliance with the License.
7// You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Pre-sized SBE byte writer for sequential encoding.
17//!
18//! The writer uses narrow `unsafe` blocks to avoid zero-initializing the backing
19//! buffer: safe-layout slice conversions between `[u8]` and `[MaybeUninit<u8>]`,
20//! and pointer-based aligned stores into the `MaybeUninit<u8>` slots. The crate
21//! otherwise forbids unsafe code; this module is the only file in the SBE
22//! codec that opts in, and every unsafe block carries a SAFETY comment
23//! documenting the invariant it relies on.
24
25#![allow(
26 unsafe_code,
27 reason = "SBE encoder needs uninit buffer writes to avoid a memset pass"
28)]
29
30use std::mem::MaybeUninit;
31
32/// Pre-sized SBE byte writer for sequential encoding.
33///
34/// Wraps a mutable `[MaybeUninit<u8>]` slice and tracks position, providing
35/// typed write methods that automatically advance the cursor. The caller is
36/// responsible for sizing the backing buffer to hold the full encoded
37/// message; writing past the end panics.
38///
39/// The writer uses `MaybeUninit<u8>` to avoid requiring the target buffer to
40/// be zero-initialized. Callers wrapping an already-initialized `&mut [u8]`
41/// can use [`SbeWriter::new`]; callers writing into uninit capacity (e.g.
42/// `Vec::spare_capacity_mut`) can use [`SbeWriter::new_uninit`].
43#[derive(Debug)]
44pub struct SbeWriter<'a> {
45 buf: &'a mut [MaybeUninit<u8>],
46 pos: usize,
47}
48
49impl<'a> SbeWriter<'a> {
50 /// Creates a new writer over an initialized byte buffer.
51 #[inline]
52 pub fn new(buf: &'a mut [u8]) -> Self {
53 // SAFETY: `MaybeUninit<u8>` has the same layout as `u8`, and going
54 // from `&mut [u8]` to `&mut [MaybeUninit<u8>]` is always valid (the
55 // elements are already initialized; we simply allow overwriting them
56 // without reading first).
57 let uninit = unsafe {
58 std::slice::from_raw_parts_mut(buf.as_mut_ptr().cast::<MaybeUninit<u8>>(), buf.len())
59 };
60 Self {
61 buf: uninit,
62 pos: 0,
63 }
64 }
65
66 /// Creates a new writer over an uninitialized byte buffer.
67 #[inline]
68 pub fn new_uninit(buf: &'a mut [MaybeUninit<u8>]) -> Self {
69 Self { buf, pos: 0 }
70 }
71
72 /// Current write position in the buffer.
73 #[inline]
74 #[must_use]
75 pub const fn pos(&self) -> usize {
76 self.pos
77 }
78
79 /// Total buffer length.
80 #[inline]
81 #[must_use]
82 pub const fn len(&self) -> usize {
83 self.buf.len()
84 }
85
86 /// Whether the buffer is empty.
87 #[inline]
88 #[must_use]
89 pub const fn is_empty(&self) -> bool {
90 self.buf.is_empty()
91 }
92
93 /// Writes a single byte and advances by 1.
94 #[inline]
95 pub fn write_u8(&mut self, value: u8) {
96 self.buf[self.pos].write(value);
97 self.pos += 1;
98 }
99
100 /// Writes a signed single byte and advances by 1.
101 #[inline]
102 pub fn write_i8(&mut self, value: i8) {
103 self.buf[self.pos].write(value as u8);
104 self.pos += 1;
105 }
106
107 /// Writes a u16 little-endian and advances by 2 bytes.
108 #[inline]
109 pub fn write_u16_le(&mut self, value: u16) {
110 self.write_array(value.to_le_bytes());
111 }
112
113 /// Writes an i16 little-endian and advances by 2 bytes.
114 #[inline]
115 pub fn write_i16_le(&mut self, value: i16) {
116 self.write_array(value.to_le_bytes());
117 }
118
119 /// Writes a u32 little-endian and advances by 4 bytes.
120 #[inline]
121 pub fn write_u32_le(&mut self, value: u32) {
122 self.write_array(value.to_le_bytes());
123 }
124
125 /// Writes an i32 little-endian and advances by 4 bytes.
126 #[inline]
127 pub fn write_i32_le(&mut self, value: i32) {
128 self.write_array(value.to_le_bytes());
129 }
130
131 /// Writes a u64 little-endian and advances by 8 bytes.
132 #[inline]
133 pub fn write_u64_le(&mut self, value: u64) {
134 self.write_array(value.to_le_bytes());
135 }
136
137 /// Writes an i64 little-endian and advances by 8 bytes.
138 #[inline]
139 pub fn write_i64_le(&mut self, value: i64) {
140 self.write_array(value.to_le_bytes());
141 }
142
143 /// Writes a u128 little-endian and advances by 16 bytes.
144 #[inline]
145 pub fn write_u128_le(&mut self, value: u128) {
146 self.write_array(value.to_le_bytes());
147 }
148
149 /// Writes an i128 little-endian and advances by 16 bytes.
150 #[inline]
151 pub fn write_i128_le(&mut self, value: i128) {
152 self.write_array(value.to_le_bytes());
153 }
154
155 /// Writes a slice of bytes and advances by its length.
156 #[inline]
157 pub fn write_bytes(&mut self, bytes: &[u8]) {
158 let end = self.pos + bytes.len();
159 // Bounds-check the destination via safe slice indexing, then copy
160 // through raw pointers. Constructing a `&mut [u8]` over the
161 // `MaybeUninit<u8>` slots would be UB even when only written; pointer
162 // writes do not impose the "reference refers to initialized memory"
163 // invariant.
164 let dst_ptr = self.buf[self.pos..end].as_mut_ptr().cast::<u8>();
165 unsafe {
166 // SAFETY: `dst_ptr` points to `bytes.len()` consecutive
167 // `MaybeUninit<u8>` slots (same layout as u8); we initialize
168 // every slot in the region. Source and destination cannot overlap
169 // because `bytes` is an immutable borrow distinct from `self.buf`.
170 std::ptr::copy_nonoverlapping(bytes.as_ptr(), dst_ptr, bytes.len());
171 }
172 self.pos = end;
173 }
174
175 // Mirrors `SbeCursor::read_array`: const-generic slice-to-array writes let
176 // LLVM lower each call to a single aligned store.
177 #[inline]
178 fn write_array<const N: usize>(&mut self, bytes: [u8; N]) {
179 let end = self.pos + N;
180 let dst_ptr = self.buf[self.pos..end].as_mut_ptr().cast::<[u8; N]>();
181 unsafe {
182 // SAFETY: slice bounds were checked by the indexing above, so
183 // `dst_ptr` points to an N-element `MaybeUninit<u8>` region with
184 // the same layout as `[u8; N]`. We overwrite all N bytes.
185 *dst_ptr = bytes;
186 }
187 self.pos = end;
188 }
189}
190
191#[cfg(test)]
192mod tests {
193 use rstest::rstest;
194
195 use super::*;
196
197 #[rstest]
198 fn test_new_starts_at_zero() {
199 let mut buf = [0u8; 4];
200 let writer = SbeWriter::new(&mut buf);
201 assert_eq!(writer.pos(), 0);
202 assert_eq!(writer.len(), 4);
203 }
204
205 #[rstest]
206 fn test_write_u8() {
207 let mut buf = [0u8; 2];
208 let mut writer = SbeWriter::new(&mut buf);
209 writer.write_u8(0x42);
210 writer.write_u8(0xFF);
211 assert_eq!(writer.pos(), 2);
212 assert_eq!(buf, [0x42, 0xFF]);
213 }
214
215 #[rstest]
216 fn test_write_u16_le() {
217 let mut buf = [0u8; 2];
218 let mut writer = SbeWriter::new(&mut buf);
219 writer.write_u16_le(0x1234);
220 assert_eq!(writer.pos(), 2);
221 assert_eq!(buf, [0x34, 0x12]);
222 }
223
224 #[rstest]
225 fn test_write_i64_le() {
226 let value: i64 = -1_234_567_890_123_456_789;
227 let mut buf = [0u8; 8];
228 let mut writer = SbeWriter::new(&mut buf);
229 writer.write_i64_le(value);
230 assert_eq!(writer.pos(), 8);
231 assert_eq!(buf, value.to_le_bytes());
232 }
233
234 #[rstest]
235 fn test_write_u128_le() {
236 let value: u128 = 0x0123_4567_89AB_CDEF_FEDC_BA98_7654_3210;
237 let mut buf = [0u8; 16];
238 let mut writer = SbeWriter::new(&mut buf);
239 writer.write_u128_le(value);
240 assert_eq!(writer.pos(), 16);
241 assert_eq!(buf, value.to_le_bytes());
242 }
243
244 #[rstest]
245 fn test_write_bytes() {
246 let mut buf = [0u8; 5];
247 let mut writer = SbeWriter::new(&mut buf);
248 writer.write_bytes(&[1, 2, 3]);
249 writer.write_bytes(&[4, 5]);
250 assert_eq!(writer.pos(), 5);
251 assert_eq!(buf, [1, 2, 3, 4, 5]);
252 }
253
254 #[rstest]
255 fn test_write_into_uninit() {
256 let mut buf: [MaybeUninit<u8>; 4] = [const { MaybeUninit::uninit() }; 4];
257 let mut writer = SbeWriter::new_uninit(&mut buf);
258 writer.write_u8(0xAA);
259 writer.write_u16_le(0x1234);
260 writer.write_u8(0xBB);
261 assert_eq!(writer.pos(), 4);
262 // SAFETY: all 4 bytes were written by the writer above.
263 let initialized: [u8; 4] = unsafe { std::mem::transmute(buf) };
264 assert_eq!(initialized, [0xAA, 0x34, 0x12, 0xBB]);
265 }
266
267 #[rstest]
268 #[should_panic(expected = "index out of bounds")]
269 fn test_write_past_end_panics() {
270 let mut buf = [0u8; 2];
271 let mut writer = SbeWriter::new(&mut buf);
272 writer.write_u8(0);
273 writer.write_u8(0);
274 writer.write_u8(0);
275 }
276}