1 files changed, 110 insertions, 0 deletions
diff --git a/src/coded_output_stream/buffer.rs b/src/coded_output_stream/buffer.rs
new file mode 100644
index 0000000..ec4c62f
--- /dev/null
+++ b/src/coded_output_stream/buffer.rs
@@ -0,0 +1,110 @@
+use std::fmt;
+use std::fmt::Formatter;
+use std::mem::MaybeUninit;
+use std::slice;
+
+use crate::misc::maybe_uninit_write_slice;
+
+pub(crate) struct OutputBuffer {
+    // Actual buffer is owned by `OutputTarget`,
+    // and here we alias the buffer so access to the buffer is branchless:
+    // access does not require switch by actual target type: `&[], `Vec`, `Write` etc.
+    // We don't access the actual buffer in `OutputTarget` except when
+    // we initialize `buffer` field here.
+    buffer: *mut [MaybeUninit<u8>],
+    /// Position within the buffer.
+    /// Always correct.
+    pos_within_buf: usize,
+}
+
+impl fmt::Debug for OutputBuffer {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("OutputBuffer")
+            .field("buffer.len", &self.buffer().len())
+            .field("pos_within_buf", &self.pos_within_buf)
+            .finish()
+    }
+}
+
+impl OutputBuffer {
+    #[inline]
+    pub(crate) fn new(buffer: *mut [MaybeUninit<u8>]) -> OutputBuffer {
+        Self {
+            buffer,
+            pos_within_buf: 0,
+        }
+    }
+
+    /// Whole buffer: written data + unwritten data.
+    #[inline]
+    pub(crate) fn buffer(&self) -> &[MaybeUninit<u8>] {
+        unsafe { &*self.buffer }
+    }
+
+    #[inline]
+    fn buffer_mut(&mut self) -> &mut [MaybeUninit<u8>] {
+        unsafe { &mut *self.buffer }
+    }
+
+    #[inline]
+    pub(crate) fn pos_within_buf(&self) -> usize {
+        self.pos_within_buf
+    }
+
+    #[inline]
+    pub(crate) fn filled(&self) -> &[u8] {
+        // SAFETY: This type invariant is data is filled up to `pos_within_buf`.
+        unsafe { slice::from_raw_parts_mut(self.buffer as *mut u8, self.pos_within_buf) }
+    }
+
+    #[inline]
+    pub(crate) fn unfilled(&mut self) -> &mut [MaybeUninit<u8>] {
+        // SAFETY: This type invariant is `pos_within_buf` is smaller than buffer length.
+        let pos_within_buf = self.pos_within_buf;
+        unsafe { self.buffer_mut().get_unchecked_mut(pos_within_buf..) }
+    }
+
+    #[inline]
+    pub(crate) fn unfilled_len(&self) -> usize {
+        self.buffer().len() - self.pos_within_buf
+    }
+
+    #[inline]
+    pub(crate) unsafe fn advance(&mut self, n: usize) {
+        debug_assert!(n <= self.unfilled_len());
+        self.pos_within_buf += n;
+    }
+
+    #[inline]
+    pub(crate) fn rewind(&mut self) {
+        self.pos_within_buf = 0;
+    }
+
+    #[inline]
+    pub(crate) fn replace_buffer_keep_pos(&mut self, buffer: *mut [MaybeUninit<u8>]) {
+        unsafe {
+            assert!(self.pos_within_buf <= (&*buffer).len());
+        }
+        self.buffer = buffer;
+    }
+
+    #[inline]
+    pub(crate) unsafe fn write_byte(&mut self, b: u8) {
+        debug_assert!(self.unfilled_len() >= 1);
+        // SAFETY: caller is responsible for ensuring that byte fits in the buffer.
+        let pos_within_buf = self.pos_within_buf;
+        self.buffer_mut().get_unchecked_mut(pos_within_buf).write(b);
+        self.pos_within_buf += 1;
+    }
+
+    #[inline]
+    pub(crate) unsafe fn write_bytes(&mut self, bytes: &[u8]) {
+        debug_assert!(self.unfilled_len() >= bytes.len());
+        let bottom = self.pos_within_buf as usize;
+        let top = bottom + (bytes.len() as usize);
+        // SAFETY: caller is responsible for ensuring that `bytes` fits in the buffer.
+        let buffer = self.buffer_mut().get_unchecked_mut(bottom..top);
+        maybe_uninit_write_slice(buffer, bytes);
+        self.pos_within_buf += bytes.len();
+    }
+}