1
#![deny(missing_docs)]
2

            
3
//! Write binary data
4

            
5
use std::marker::PhantomData;
6

            
7
use crate::binary::read::{ReadArray, ReadArrayCow, ReadScope, ReadUnchecked};
8
use crate::binary::{I16Be, I32Be, I64Be, U16Be, U24Be, U32Be, I8, U8};
9
use crate::error::WriteError;
10

            
11
/// An in-memory buffer that implements `WriteContext`.
12
pub struct WriteBuffer {
13
    data: Vec<u8>,
14
}
15

            
16
struct WriteSlice<'a> {
17
    offset: usize,
18
    data: &'a mut [u8],
19
}
20

            
21
/// A `WriteContext` implementation that just counts the bytes written.
22
pub struct WriteCounter {
23
    count: usize,
24
}
25

            
26
/// A `WriteContext` implementation that writes nothing.
27
struct NullWriter;
28

            
29
/// A placeholder for a value that will be filled in later using WriteContext::write_placeholder
30
pub struct Placeholder<T, HostType>
31
where
32
    T: WriteBinaryDep<HostType>,
33
{
34
    offset: usize,
35
    length: usize,
36
    marker: PhantomData<T>,
37
    host: PhantomData<HostType>,
38
}
39

            
40
/// Trait that describes a type that can be written to a `WriteContext` in binary form.
41
pub trait WriteBinary<HostType = Self> {
42
    /// The type of the value returned by `write`.
43
    type Output;
44

            
45
    /// Write the binary representation of Self to `ctxt`.
46
    fn write<C: WriteContext>(ctxt: &mut C, val: HostType) -> Result<Self::Output, WriteError>;
47
}
48

            
49
/// Trait that describes a type that can be written to a `WriteContext` in binary form with
50
/// dependent arguments.
51
pub trait WriteBinaryDep<HostType = Self> {
52
    /// The type of the arguments supplied to `write_dep`.
53
    type Args;
54
    /// The type of the value returned by `write_dep`.
55
    type Output;
56

            
57
    /// Write the binary representation of Self to `ctxt`.
58
    fn write_dep<C: WriteContext>(
59
        ctxt: &mut C,
60
        val: HostType,
61
        args: Self::Args,
62
    ) -> Result<Self::Output, WriteError>;
63
}
64

            
65
/// Trait for types that can have binary data written to them.
66
pub trait WriteContext {
67
    /// Write a `ReadArray` instance to a `WriteContext`.
68
    fn write_array<T>(&mut self, array: &ReadArray<'_, T>) -> Result<(), WriteError>
69
    where
70
        Self: Sized,
71
        T: ReadUnchecked + WriteBinary<<T as ReadUnchecked>::HostType>,
72
    {
73
        <&ReadArray<'_, _>>::write(self, array)
74
    }
75

            
76
    /// Write a `Vec` into a `WriteContext`.
77
    fn write_vec<T, HostType>(&mut self, vec: Vec<HostType>) -> Result<(), WriteError>
78
    where
79
        Self: Sized,
80
        T: WriteBinary<HostType>,
81
    {
82
        for val in vec {
83
            T::write(self, val)?;
84
        }
85

            
86
        Ok(())
87
    }
88

            
89
    /// Write a slice of values into a `WriteContext`.
90
    fn write_iter<T, HostType>(
91
        &mut self,
92
        iter: impl Iterator<Item = HostType>,
93
    ) -> Result<(), WriteError>
94
    where
95
        Self: Sized,
96
        T: WriteBinary<HostType>,
97
    {
98
        for val in iter {
99
            T::write(self, val)?;
100
        }
101

            
102
        Ok(())
103
    }
104

            
105
    /// Write a slice of bytes to a `WriteContext`.
106
    fn write_bytes(&mut self, data: &[u8]) -> Result<(), WriteError>;
107

            
108
    /// Write the specified number of zero bytes to the `WriteContext`.
109
    fn write_zeros(&mut self, count: usize) -> Result<(), WriteError>;
110

            
111
    /// The total number of bytes written so far.
112
    fn bytes_written(&self) -> usize;
113

            
114
    /// Return a placeholder to `T` in the context for filling in later.
115
    fn placeholder<T, HostType>(&mut self) -> Result<Placeholder<T, HostType>, WriteError>
116
    where
117
        T: WriteBinary<HostType> + ReadUnchecked,
118
    {
119
        let offset = self.bytes_written();
120
        self.write_zeros(T::SIZE)?;
121

            
122
        Ok(Placeholder {
123
            offset,
124
            length: T::SIZE,
125
            marker: PhantomData,
126
            host: PhantomData,
127
        })
128
    }
129

            
130
    /// Reserve space for `count` bytes in the context for filling in later.
131
    fn reserve<'a, T, HostType>(
132
        &mut self,
133
        count: usize,
134
    ) -> Result<Placeholder<T, &'a HostType>, WriteError>
135
    where
136
        T: WriteBinaryDep<&'a HostType>,
137
    {
138
        let offset = self.bytes_written();
139
        self.write_zeros(count)?;
140

            
141
        Ok(Placeholder {
142
            offset,
143
            length: count,
144
            marker: PhantomData,
145
            host: PhantomData,
146
        })
147
    }
148

            
149
    /// Return a `Vec` of `count` placeholders of type `T`.
150
    fn placeholder_array<T, HostType>(
151
        &mut self,
152
        count: usize,
153
    ) -> Result<Vec<Placeholder<T, HostType>>, WriteError>
154
    where
155
        T: WriteBinary<HostType> + ReadUnchecked,
156
    {
157
        (0..count)
158
            .map(|_| self.placeholder::<T, HostType>())
159
            .collect()
160
    }
161

            
162
    /// Consumes the placeholder and writes the supplied value into it
163
    fn write_placeholder<T, HostType>(
164
        &mut self,
165
        placeholder: Placeholder<T, HostType>,
166
        val: HostType,
167
    ) -> Result<T::Output, WriteError>
168
    where
169
        T: WriteBinary<HostType>;
170

            
171
    /// Consumes the placeholder and writes the supplied value into it.
172
    /// `WriteBinaryDep` version
173
    fn write_placeholder_dep<T, HostType>(
174
        &mut self,
175
        placeholder: Placeholder<T, HostType>,
176
        val: HostType,
177
        args: T::Args,
178
    ) -> Result<T::Output, WriteError>
179
    where
180
        T: WriteBinaryDep<HostType>;
181
}
182

            
183
/// Write `T` into a `WriteBuffer` and return it
184
pub fn buffer<HostType, T: WriteBinaryDep<HostType>>(
185
    writeable: HostType,
186
    args: T::Args,
187
) -> Result<(T::Output, WriteBuffer), WriteError> {
188
    let mut buffer = WriteBuffer::new();
189
    let output = T::write_dep(&mut buffer, writeable, args)?;
190
    Ok((output, buffer))
191
}
192

            
193
impl<T, HostType> WriteBinaryDep<HostType> for T
194
where
195
    T: WriteBinary<HostType>,
196
{
197
    type Args = ();
198
    type Output = T::Output;
199

            
200
    fn write_dep<C: WriteContext>(
201
        ctxt: &mut C,
202
        val: HostType,
203
        (): Self::Args,
204
    ) -> Result<Self::Output, WriteError> {
205
        T::write(ctxt, val)
206
    }
207
}
208

            
209
impl<T> WriteBinary<T> for U8
210
where
211
    T: Into<u8>,
212
{
213
    type Output = ();
214

            
215
    fn write<C: WriteContext>(ctxt: &mut C, t: T) -> Result<(), WriteError> {
216
        let val: u8 = t.into();
217
        ctxt.write_bytes(&[val])
218
    }
219
}
220

            
221
impl<T> WriteBinary<T> for I8
222
where
223
    T: Into<i8>,
224
{
225
    type Output = ();
226

            
227
    fn write<C: WriteContext>(ctxt: &mut C, t: T) -> Result<(), WriteError> {
228
        let val: i8 = t.into();
229
        ctxt.write_bytes(&val.to_be_bytes())
230
    }
231
}
232

            
233
impl<T> WriteBinary<T> for I16Be
234
where
235
    T: Into<i16>,
236
{
237
    type Output = ();
238

            
239
    fn write<C: WriteContext>(ctxt: &mut C, t: T) -> Result<(), WriteError> {
240
        let val: i16 = t.into();
241
        ctxt.write_bytes(&val.to_be_bytes())
242
    }
243
}
244

            
245
impl<T> WriteBinary<T> for U16Be
246
where
247
    T: Into<u16>,
248
{
249
    type Output = ();
250

            
251
    fn write<C: WriteContext>(ctxt: &mut C, t: T) -> Result<(), WriteError> {
252
        let val: u16 = t.into();
253
        ctxt.write_bytes(&val.to_be_bytes())
254
    }
255
}
256

            
257
impl<T> WriteBinary<T> for U24Be
258
where
259
    T: Into<u32>,
260
{
261
    type Output = ();
262

            
263
    fn write<C: WriteContext>(ctxt: &mut C, t: T) -> Result<(), WriteError> {
264
        let val: u32 = t.into();
265
        if val > 0xFF_FFFF {
266
            return Err(WriteError::BadValue);
267
        }
268
        ctxt.write_bytes(&val.to_be_bytes()[1..4])
269
    }
270
}
271

            
272
impl<T> WriteBinary<T> for I32Be
273
where
274
    T: Into<i32>,
275
{
276
    type Output = ();
277

            
278
    fn write<C: WriteContext>(ctxt: &mut C, t: T) -> Result<(), WriteError> {
279
        let val: i32 = t.into();
280
        ctxt.write_bytes(&val.to_be_bytes())
281
    }
282
}
283

            
284
impl<T> WriteBinary<T> for U32Be
285
where
286
    T: Into<u32>,
287
{
288
    type Output = ();
289

            
290
    fn write<C: WriteContext>(ctxt: &mut C, t: T) -> Result<(), WriteError> {
291
        let val: u32 = t.into();
292
        ctxt.write_bytes(&val.to_be_bytes())
293
    }
294
}
295

            
296
impl<T> WriteBinary<T> for I64Be
297
where
298
    T: Into<i64>,
299
{
300
    type Output = ();
301

            
302
    fn write<C: WriteContext>(ctxt: &mut C, t: T) -> Result<(), WriteError> {
303
        let val: i64 = t.into();
304
        ctxt.write_bytes(&val.to_be_bytes())
305
    }
306
}
307

            
308
impl WriteContext for WriteBuffer {
309
    fn write_bytes(&mut self, data: &[u8]) -> Result<(), WriteError> {
310
        self.data.extend(data.iter());
311
        Ok(())
312
    }
313

            
314
    fn write_zeros(&mut self, count: usize) -> Result<(), WriteError> {
315
        let zeros = std::iter::repeat_n(0, count);
316
        self.data.extend(zeros);
317
        Ok(())
318
    }
319

            
320
    fn bytes_written(&self) -> usize {
321
        self.data.len()
322
    }
323

            
324
    fn write_placeholder<T, HostType>(
325
        &mut self,
326
        placeholder: Placeholder<T, HostType>,
327
        val: HostType,
328
    ) -> Result<T::Output, WriteError>
329
    where
330
        T: WriteBinary<HostType>,
331
    {
332
        let data = &mut self.data[placeholder.offset..];
333
        let data = &mut data[0..placeholder.length];
334
        let mut slice = WriteSlice { offset: 0, data };
335
        T::write(&mut slice, val)
336
    }
337

            
338
    fn write_placeholder_dep<T, HostType>(
339
        &mut self,
340
        placeholder: Placeholder<T, HostType>,
341
        val: HostType,
342
        args: T::Args,
343
    ) -> Result<T::Output, WriteError>
344
    where
345
        T: WriteBinaryDep<HostType>,
346
    {
347
        let data = &mut self.data[placeholder.offset..];
348
        let data = &mut data[0..placeholder.length];
349
        let mut slice = WriteSlice { offset: 0, data };
350
        T::write_dep(&mut slice, val, args)
351
    }
352
}
353

            
354
impl WriteContext for WriteSlice<'_> {
355
    fn write_bytes(&mut self, data: &[u8]) -> Result<(), WriteError> {
356
        let data_len = data.len();
357
        let self_len = self.data.len();
358

            
359
        if data_len <= self_len {
360
            let subslice = &mut self.data[self.offset..][0..data_len];
361
            subslice.copy_from_slice(data);
362
            self.offset += data_len;
363
            Ok(())
364
        } else {
365
            Err(WriteError::PlaceholderMismatch)
366
        }
367
    }
368

            
369
    fn write_zeros(&mut self, count: usize) -> Result<(), WriteError> {
370
        for i in 0..count.min(self.data.len()) {
371
            self.data[i] = 0;
372
        }
373

            
374
        Ok(())
375
    }
376

            
377
    fn bytes_written(&self) -> usize {
378
        self.data.len()
379
    }
380

            
381
    fn write_placeholder<T, HostType>(
382
        &mut self,
383
        _placeholder: Placeholder<T, HostType>,
384
        _val: HostType,
385
    ) -> Result<T::Output, WriteError>
386
    where
387
        T: WriteBinary<HostType>,
388
    {
389
        unimplemented!()
390
    }
391

            
392
    fn write_placeholder_dep<T, HostType>(
393
        &mut self,
394
        _placeholder: Placeholder<T, HostType>,
395
        _val: HostType,
396
        _args: T::Args,
397
    ) -> Result<T::Output, WriteError>
398
    where
399
        T: WriteBinaryDep<HostType>,
400
    {
401
        unimplemented!()
402
    }
403
}
404

            
405
impl WriteContext for WriteCounter {
406
    fn write_bytes(&mut self, data: &[u8]) -> Result<(), WriteError> {
407
        self.count += data.len();
408
        Ok(())
409
    }
410

            
411
    fn write_zeros(&mut self, count: usize) -> Result<(), WriteError> {
412
        self.count += count;
413
        Ok(())
414
    }
415

            
416
    fn bytes_written(&self) -> usize {
417
        self.count
418
    }
419

            
420
    fn write_placeholder<T, HostType>(
421
        &mut self,
422
        _placeholder: Placeholder<T, HostType>,
423
        val: HostType,
424
    ) -> Result<T::Output, WriteError>
425
    where
426
        T: WriteBinary<HostType>,
427
    {
428
        let mut null = NullWriter;
429
        T::write(&mut null, val)
430
    }
431

            
432
    fn write_placeholder_dep<T, HostType>(
433
        &mut self,
434
        _placeholder: Placeholder<T, HostType>,
435
        val: HostType,
436
        args: T::Args,
437
    ) -> Result<T::Output, WriteError>
438
    where
439
        T: WriteBinaryDep<HostType>,
440
    {
441
        let mut null = NullWriter;
442
        T::write_dep(&mut null, val, args)
443
    }
444
}
445

            
446
impl WriteContext for NullWriter {
447
    fn write_bytes(&mut self, _data: &[u8]) -> Result<(), WriteError> {
448
        Ok(())
449
    }
450

            
451
    fn write_zeros(&mut self, _count: usize) -> Result<(), WriteError> {
452
        Ok(())
453
    }
454

            
455
    fn bytes_written(&self) -> usize {
456
        0
457
    }
458

            
459
    fn write_placeholder<T, HostType>(
460
        &mut self,
461
        _placeholder: Placeholder<T, HostType>,
462
        _val: HostType,
463
    ) -> Result<T::Output, WriteError>
464
    where
465
        T: WriteBinary<HostType>,
466
    {
467
        unimplemented!()
468
    }
469

            
470
    fn write_placeholder_dep<T, HostType>(
471
        &mut self,
472
        _placeholder: Placeholder<T, HostType>,
473
        _val: HostType,
474
        _args: T::Args,
475
    ) -> Result<T::Output, WriteError>
476
    where
477
        T: WriteBinaryDep<HostType>,
478
    {
479
        unimplemented!()
480
    }
481
}
482

            
483
impl<T> WriteBinary for &ReadArray<'_, T>
484
where
485
    T: ReadUnchecked + WriteBinary<<T as ReadUnchecked>::HostType>,
486
{
487
    type Output = ();
488

            
489
    fn write<C: WriteContext>(ctxt: &mut C, array: Self) -> Result<(), WriteError> {
490
        for val in array.into_iter() {
491
            T::write(ctxt, val)?;
492
        }
493

            
494
        Ok(())
495
    }
496
}
497

            
498
impl<T> WriteBinary<&Self> for ReadArrayCow<'_, T>
499
where
500
    T: ReadUnchecked + WriteBinary<<T as ReadUnchecked>::HostType>,
501
    T::HostType: Copy,
502
{
503
    type Output = ();
504

            
505
    fn write<C: WriteContext>(ctxt: &mut C, array: &Self) -> Result<(), WriteError> {
506
        for val in array.iter() {
507
            T::write(ctxt, val)?;
508
        }
509

            
510
        Ok(())
511
    }
512
}
513

            
514
impl WriteBinary for ReadScope<'_> {
515
    type Output = ();
516

            
517
    fn write<C: WriteContext>(ctxt: &mut C, scope: Self) -> Result<(), WriteError> {
518
        ctxt.write_bytes(scope.data())
519
    }
520
}
521

            
522
impl WriteBuffer {
523
    /// Create a new, empty `WriteBuffer`
524
    pub fn new() -> Self {
525
        WriteBuffer { data: Vec::new() }
526
    }
527

            
528
    /// Retrieve a slice of the data held by this buffer
529
    pub fn bytes(&self) -> &[u8] {
530
        &self.data
531
    }
532

            
533
    /// Clear the internal data so that this buffer can be reused
534
    pub fn clear(&mut self) {
535
        self.data.clear();
536
    }
537

            
538
    /// Returns the current size of the data held by this buffer
539
    pub fn len(&self) -> usize {
540
        self.data.len()
541
    }
542

            
543
    /// Consume `self` and return the inner buffer
544
    pub fn into_inner(self) -> Vec<u8> {
545
        self.data
546
    }
547
}
548

            
549
impl WriteCounter {
550
    /// Create a new, empty `WriteCounter`
551
    pub fn new() -> Self {
552
        WriteCounter { count: 0 }
553
    }
554
}
555

            
556
#[cfg(test)]
557
mod tests {
558
    use super::*;
559
    use crate::tag;
560

            
561
    struct TestTable {
562
        tag: u32,
563
    }
564

            
565
    struct BigStruct {
566
        tag: u32,
567
    }
568

            
569
    impl WriteBinary<Self> for TestTable {
570
        type Output = ();
571

            
572
        fn write<C: WriteContext>(ctxt: &mut C, val: Self) -> Result<(), WriteError> {
573
            U32Be::write(ctxt, val.tag)
574
        }
575
    }
576

            
577
    impl WriteBinary<&Self> for BigStruct {
578
        type Output = ();
579

            
580
        fn write<C: WriteContext>(ctxt: &mut C, val: &Self) -> Result<(), WriteError> {
581
            U32Be::write(ctxt, val.tag)
582
        }
583
    }
584

            
585
    #[test]
586
    fn test_basic() {
587
        let mut ctxt = WriteBuffer::new();
588
        let table = TestTable { tag: tag::GLYF };
589
        let big = BigStruct { tag: tag::BLOC };
590

            
591
        TestTable::write(&mut ctxt, table).unwrap();
592
        BigStruct::write(&mut ctxt, &big).unwrap();
593

            
594
        assert_eq!(ctxt.bytes(), b"glyfbloc")
595
    }
596

            
597
    #[test]
598
    fn test_write_u24be() {
599
        let mut ctxt = WriteBuffer::new();
600
        U24Be::write(&mut ctxt, 0x10203u32).unwrap();
601
        assert_eq!(ctxt.bytes(), &[1, 2, 3]);
602

            
603
        // Check out of range value
604
        match U24Be::write(&mut ctxt, std::u32::MAX) {
605
            Err(WriteError::BadValue) => {}
606
            _ => panic!("Expected WriteError::BadValue"),
607
        }
608
    }
609

            
610
    #[test]
611
    fn test_write_placeholder() {
612
        let mut ctxt = WriteBuffer::new();
613
        U8::write(&mut ctxt, 1).unwrap();
614
        let placeholder = ctxt.placeholder::<U16Be, u16>().unwrap();
615
        U8::write(&mut ctxt, 3).unwrap();
616
        ctxt.write_placeholder(placeholder, 2).unwrap();
617
        assert_eq!(ctxt.bytes(), &[1, 0, 2, 3]);
618
    }
619

            
620
    #[test]
621
    fn test_write_placeholder_overflow() {
622
        // Test that trying to write more data than reserved results in an error
623
        let mut ctxt = WriteBuffer::new();
624
        let placeholder = ctxt.reserve::<BigStruct, _>(1).unwrap();
625
        let value = BigStruct { tag: 1234 };
626
        assert!(ctxt.write_placeholder(placeholder, &value).is_err());
627
    }
628
}