1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

//! A container which can store up to a fixed number of uninitialized bytes on
//! the stack and read into and from it.

use core::fmt;
use core::marker;
use core::mem::MaybeUninit;
use core::ptr;

use musli::error::Error;

use crate::writer::Writer;

/// A fixed-size bytes storage which keeps track of how much has been initialized.
pub struct FixedBytes<const N: usize, E = FixedBytesWriterError> {
    /// Data storage.
    data: [MaybeUninit<u8>; N],
    /// How many bytes have been initialized.
    init: usize,
    /// Error type to raise when this is used as a `Writer` implementation.
    _marker: marker::PhantomData<E>,
}

impl<const N: usize, E> FixedBytes<N, E> {
    /// Construct a new fixed bytes array storage.
    pub const fn new() -> Self {
        Self {
            // SAFETY: MaybeUnint::uninit_array is not stable.
            data: unsafe { MaybeUninit::<[MaybeUninit<u8>; N]>::uninit().assume_init() },
            init: 0,
            _marker: marker::PhantomData,
        }
    }

    /// Get the length of the collection.
    pub const fn len(&self) -> usize {
        self.init
    }

    /// Coerce into the underlying bytes if all of them have been initialized.
    pub fn into_bytes(self) -> Option<[u8; N]> {
        if self.init == N {
            // SAFETY: All of the bytes in the sequence have been initialized
            // and can be safety transmuted.
            //
            // Method of transmuting comes from the implementation of
            // `MaybeUninit::array_assume_init` which is not yet stable.
            unsafe { Some((&self.data as *const _ as *const [u8; N]).read()) }
        } else {
            None
        }
    }

    /// Coerce into the slice of initialized memory which is present.
    pub fn as_bytes(&self) -> &[u8] {
        if self.init == 0 {
            return &[];
        }

        // SAFETY: We've asserted that `initialized` accounts for the number of
        // bytes that have been initialized.
        unsafe { std::slice::from_raw_parts(self.data.as_ptr() as *const u8, self.init) }
    }
}

decl_message_repr!(FixedBytesWriterErrorRepr, "failed to write to fixed bytes");

/// An error raised while decoding a slice.
#[derive(Debug)]
pub struct FixedBytesWriterError(FixedBytesWriterErrorRepr);

impl fmt::Display for FixedBytesWriterError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl Error for FixedBytesWriterError {
    fn custom<T>(message: T) -> Self
    where
        T: 'static + Send + Sync + fmt::Display + fmt::Debug,
    {
        Self(FixedBytesWriterErrorRepr::collect(message))
    }

    fn collect_from_display<T>(message: T) -> Self
    where
        T: fmt::Display,
    {
        Self(FixedBytesWriterErrorRepr::collect(message))
    }
}

#[cfg(feature = "std")]
impl std::error::Error for FixedBytesWriterError {}

impl<const N: usize, E> Writer for FixedBytes<N, E>
where
    E: Error,
{
    type Error = E;

    fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), Self::Error> {
        if bytes.len() > N.saturating_sub(self.init) {
            return Err(E::collect_from_display(BufferOverflow {
                at: self.init,
                additional: bytes.len(),
                capacity: N,
            }));
        }

        unsafe {
            let dst = (self.data.as_mut_ptr() as *mut u8).add(self.init);
            ptr::copy_nonoverlapping(bytes.as_ptr(), dst, bytes.len());
        }

        self.init += bytes.len();
        Ok(())
    }
}

struct BufferOverflow {
    at: usize,
    additional: usize,
    capacity: usize,
}

impl fmt::Display for BufferOverflow {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let Self {
            at,
            additional,
            capacity,
        } = *self;

        write!(
            f,
            "Overflow when writing {additional} bytes at {at} with capacity {capacity}"
        )
    }
}