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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
//! main structures and traits used to build serializers
use crate::lib::std::{
    fmt,
    io::{self, Seek as _, SeekFrom, Write},
};

/// Holds the result of serializing functions
///
/// The `Ok` case returns the `Write` used for writing, in the `Err` case an instance of
/// `cookie_factory::GenError` is returned.
pub type GenResult<W> = Result<WriteContext<W>, GenError>;

/// Base type for generator errors
#[derive(Debug)]
pub enum GenError {
    /// Input buffer is too small. Argument is the maximum index that is required
    BufferTooSmall(usize),
    /// We expected to fill the whole buffer but there is some space left
    BufferTooBig(usize),
    /// Operation asked for accessing an invalid index
    InvalidOffset,
    /// IoError returned by Write
    IoError(io::Error),

    /// Allocated for custom errors
    CustomError(u32),
    /// Generator or function not yet implemented
    NotYetImplemented,
}

impl fmt::Display for GenError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

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

impl From<io::Error> for GenError {
    fn from(err: io::Error) -> Self {
        GenError::IoError(err)
    }
}

/// Trait for serializing functions
///
/// Serializing functions take one input `W` that is the target of writing and return an instance
/// of `cookie_factory::GenResult`.
///
/// This trait is implemented for all `Fn(W) -> GenResult<W>`.
pub trait SerializeFn<W>: Fn(WriteContext<W>) -> GenResult<W> {}

impl<W, F: Fn(WriteContext<W>) -> GenResult<W>> SerializeFn<W> for F {}

/// Context around a `Write` impl that is passed through serializing functions
///
/// Currently this only keeps track of the current write position since the start of serialization.
pub struct WriteContext<W> {
    pub write: W,
    pub position: u64,
}

impl<W: Write> From<W> for WriteContext<W> {
    fn from(write: W) -> Self {
        Self { write, position: 0 }
    }
}

impl<W: Write> WriteContext<W> {
    /// Returns the contained `Write` and the current position
    pub fn into_inner(self) -> (W, u64) {
        (self.write, self.position)
    }
}

impl<W: Write> Write for WriteContext<W> {
    fn write(&mut self, data: &[u8]) -> crate::lib::std::io::Result<usize> {
        let amt = self.write.write(data)?;
        self.position += amt as u64;
        Ok(amt)
    }

    #[cfg(feature = "std")]
    fn flush(&mut self) -> io::Result<()> {
        self.write.flush()
    }
}

impl<W: Seek> io::Seek for WriteContext<W> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        let old_pos = self.write.stream_position()?;
        let new_pos = self.write.seek(pos)?;
        if new_pos >= old_pos {
            self.position += new_pos - old_pos;
        } else {
            self.position -= old_pos - new_pos;
        }
        Ok(new_pos)
    }
}

/// Runs the given serializer `f` with the `Write` impl `w` and returns the result
///
/// This internally wraps `w` in a `WriteContext`, starting at position 0.
///
/// ```rust
/// use cookie_factory::{gen, combinator::slice};
///
/// let mut buf = [0u8; 100];
///
/// {
///   let (buf, pos) = gen(slice(&b"abcd"[..]), &mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 4);
///   assert_eq!(pos, 4);
/// }
///
/// assert_eq!(&buf[..4], &b"abcd"[..]);
/// ```
pub fn gen<W: Write, F: SerializeFn<W>>(f: F, w: W) -> Result<(W, u64), GenError> {
    f(WriteContext::from(w)).map(|ctx| ctx.into_inner())
}

/// Runs the given serializer `f` with the `Write` impl `w` and returns the updated `w`
///
/// This internally wraps `w` in a `WriteContext`, starting at position 0.
///
/// ```rust
/// use cookie_factory::{gen_simple, combinator::slice};
///
/// let mut buf = [0u8; 100];
///
/// {
///   let buf = gen_simple(slice(&b"abcd"[..]), &mut buf[..]).unwrap();
///   assert_eq!(buf.len(), 100 - 4);
/// }
///
/// assert_eq!(&buf[..4], &b"abcd"[..]);
/// ```
pub fn gen_simple<W: Write, F: SerializeFn<W>>(f: F, w: W) -> Result<W, GenError> {
    f(WriteContext::from(w)).map(|ctx| ctx.into_inner().0)
}

/// Trait for `Write` types that allow skipping over the data
pub trait Skip: Write {
    fn skip(s: WriteContext<Self>, sz: usize) -> GenResult<Self>
    where
        Self: Sized;
}

/// Trait for `Write` types that allow skipping and reserving a slice, then writing some data,
/// then write something in the slice we reserved using the return for our data write.
pub trait BackToTheBuffer: Write {
    fn reserve_write_use<
        Tmp,
        Gen: Fn(WriteContext<Self>) -> Result<(WriteContext<Self>, Tmp), GenError>,
        Before: Fn(WriteContext<Self>, Tmp) -> GenResult<Self>,
    >(
        s: WriteContext<Self>,
        reserved: usize,
        gen: &Gen,
        before: &Before,
    ) -> Result<WriteContext<Self>, GenError>
    where
        Self: Sized;
}

/// Trait for `Seek` types that want to automatically implement `BackToTheBuffer`
pub trait Seek: Write + io::Seek {}
impl Seek for io::Cursor<&mut [u8]> {}

impl<W: Seek> BackToTheBuffer for W {
    fn reserve_write_use<
        Tmp,
        Gen: Fn(WriteContext<Self>) -> Result<(WriteContext<Self>, Tmp), GenError>,
        Before: Fn(WriteContext<Self>, Tmp) -> GenResult<Self>,
    >(
        mut s: WriteContext<Self>,
        reserved: usize,
        gen: &Gen,
        before: &Before,
    ) -> Result<WriteContext<Self>, GenError> {
        let start = s.stream_position()?;
        let begin = s.seek(SeekFrom::Current(reserved as i64))?;
        let (mut buf, tmp) = gen(s)?;
        let end = buf.stream_position()?;
        buf.seek(SeekFrom::Start(start))?;
        let mut buf = before(buf, tmp)?;
        let pos = buf.stream_position()?;
        if pos != begin {
            return Err(GenError::BufferTooBig((begin - pos) as usize));
        }
        buf.seek(SeekFrom::Start(end))?;
        Ok(buf)
    }
}

impl Skip for &mut [u8] {
    fn skip(s: WriteContext<Self>, len: usize) -> Result<WriteContext<Self>, GenError> {
        if s.write.len() < len {
            Err(GenError::BufferTooSmall(len - s.write.len()))
        } else {
            Ok(WriteContext {
                write: &mut s.write[len..],
                position: s.position + len as u64,
            })
        }
    }
}

impl Skip for io::Cursor<&mut [u8]> {
    fn skip(mut s: WriteContext<Self>, len: usize) -> GenResult<Self> {
        let remaining = s
            .write
            .get_ref()
            .len()
            .saturating_sub(s.write.position() as usize);
        if remaining < len {
            Err(GenError::BufferTooSmall(len - remaining))
        } else {
            let cursor_position = s.write.position();
            s.write.set_position(cursor_position + len as u64);
            s.position += len as u64;
            Ok(s)
        }
    }
}

impl BackToTheBuffer for &mut [u8] {
    fn reserve_write_use<
        Tmp,
        Gen: Fn(WriteContext<Self>) -> Result<(WriteContext<Self>, Tmp), GenError>,
        Before: Fn(WriteContext<Self>, Tmp) -> GenResult<Self>,
    >(
        s: WriteContext<Self>,
        reserved: usize,
        gen: &Gen,
        before: &Before,
    ) -> Result<WriteContext<Self>, GenError> {
        let WriteContext {
            write: slice,
            position: original_position,
        } = s;

        let (res, buf) = slice.split_at_mut(reserved);
        let (new_context, tmp) = gen(WriteContext {
            write: buf,
            position: original_position + reserved as u64,
        })?;

        let res = before(
            WriteContext {
                write: res,
                position: original_position,
            },
            tmp,
        )?;

        if !res.write.is_empty() {
            return Err(GenError::BufferTooBig(res.write.len()));
        }

        Ok(new_context)
    }
}

#[cfg(feature = "std")]
impl BackToTheBuffer for Vec<u8> {
    fn reserve_write_use<
        Tmp,
        Gen: Fn(WriteContext<Self>) -> Result<(WriteContext<Self>, Tmp), GenError>,
        Before: Fn(WriteContext<Self>, Tmp) -> GenResult<Self>,
    >(
        s: WriteContext<Self>,
        reserved: usize,
        gen: &Gen,
        before: &Before,
    ) -> Result<WriteContext<Self>, GenError> {
        let WriteContext {
            write: mut vec,
            position: original_position,
        } = s;

        let start_len = vec.len();
        vec.extend(std::iter::repeat(0).take(reserved));

        let (mut new_context, tmp) = gen(WriteContext {
            write: vec,
            position: original_position + reserved as u64,
        })?;

        let tmp_context = before(
            WriteContext {
                write: Vec::new(),
                position: original_position,
            },
            tmp,
        )?;

        let tmp_written = tmp_context.write.len();
        if tmp_written != reserved {
            return Err(GenError::BufferTooBig(reserved - tmp_written));
        }

        // FIXME?: find a way to do that without copying
        // Vec::from_raw_parts + core::mem::forget makes it work, but
        // if `before` writes more than `reserved`, realloc will cause troubles
        new_context.write[start_len..(start_len + reserved)]
            .copy_from_slice(&tmp_context.write[..]);

        Ok(new_context)
    }
}