src/io/writer.rs

// Copyright 2018 Brian Smith.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHORS DISCLAIM ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY
// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

use alloc::{boxed::Box, vec::Vec};

pub trait Accumulator {
    fn write_byte(&mut self, value: u8);
    fn write_bytes(&mut self, value: &[u8]);
}

pub(super) struct LengthMeasurement {
    len: usize,
}

impl Into<usize> for LengthMeasurement {
    fn into(self) -> usize {
        self.len
    }
}

impl LengthMeasurement {
    pub fn zero() -> Self {
        Self { len: 0 }
    }
}

impl Accumulator for LengthMeasurement {
    fn write_byte(&mut self, _value: u8) {
        self.len += 1;
    }
    fn write_bytes(&mut self, value: &[u8]) {
        self.len += value.len();
    }
}

pub(super) struct Writer {
    bytes: Vec<u8>,
    requested_capacity: usize,
}

impl Writer {
    pub(super) fn with_capacity(capacity: LengthMeasurement) -> Self {
        Self {
            bytes: Vec::with_capacity(capacity.len),
            requested_capacity: capacity.len,
        }
    }
}

impl Into<Box<[u8]>> for Writer {
    fn into(self) -> Box<[u8]> {
        assert_eq!(self.requested_capacity, self.bytes.len());
        self.bytes.into_boxed_slice()
    }
}

impl Accumulator for Writer {
    fn write_byte(&mut self, value: u8) {
        self.bytes.push(value);
    }
    fn write_bytes(&mut self, value: &[u8]) {
        self.bytes.extend(value);
    }
}

pub fn write_copy(accumulator: &mut dyn Accumulator, to_copy: untrusted::Input) {
    accumulator.write_bytes(to_copy.as_slice_less_safe())
}