diff options
Diffstat (limited to 'src/coded_output_stream/mod.rs')
| -rw-r--r-- | src/coded_output_stream/mod.rs | 1225 |
1 files changed, 1225 insertions, 0 deletions
diff --git a/src/coded_output_stream/mod.rs b/src/coded_output_stream/mod.rs new file mode 100644 index 0000000..5010df9 --- /dev/null +++ b/src/coded_output_stream/mod.rs @@ -0,0 +1,1225 @@ +mod buffer; +mod output_target; +pub(crate) mod with; + +use std::io; +use std::io::Write; +use std::mem::MaybeUninit; +use std::ptr; +use std::slice; + +use crate::byteorder::LITTLE_ENDIAN; +use crate::coded_output_stream::buffer::OutputBuffer; +use crate::coded_output_stream::output_target::OutputTarget; +use crate::error::ProtobufError; +use crate::rt::packed::vec_packed_enum_or_unknown_data_size; +use crate::rt::packed::vec_packed_fixed_data_size; +use crate::rt::packed::vec_packed_varint_data_size; +use crate::rt::packed::vec_packed_varint_zigzag_data_size; +use crate::varint::encode::encode_varint32; +use crate::varint::encode::encode_varint64; +use crate::varint::encode::encoded_varint64_len; +use crate::varint::MAX_VARINT_ENCODED_LEN; +use crate::wire_format; +use crate::wire_format::check_message_size; +use crate::wire_format::WireType; +use crate::wire_format::MAX_MESSAGE_SIZE; +use crate::zigzag::encode_zig_zag_32; +use crate::zigzag::encode_zig_zag_64; +use crate::Enum; +use crate::EnumOrUnknown; +use crate::Message; +use crate::MessageDyn; +use crate::MessageFull; +use crate::UnknownFields; +use crate::UnknownValueRef; + +// Equal to the default buffer size of `BufWriter`, so when +// `CodedOutputStream` wraps `BufWriter`, it often skips double buffering. +const OUTPUT_STREAM_BUFFER_SIZE: usize = 8 * 1024; + +/// Buffered write with handy utilities +#[derive(Debug)] +pub struct CodedOutputStream<'a> { + target: OutputTarget<'a>, + buffer: OutputBuffer, + /// Absolute position of the buffer start. + pos_of_buffer_start: u64, +} + +impl<'a> CodedOutputStream<'a> { + /// Construct from given `Write`. + /// + /// `CodedOutputStream` is buffered even if `Write` is not + pub fn new(writer: &'a mut dyn Write) -> CodedOutputStream<'a> { + let buffer_len = OUTPUT_STREAM_BUFFER_SIZE; + + let mut buffer_storage = Vec::with_capacity(buffer_len); + + // SAFETY: we are not using the `buffer_storage` + // except for initializing the `buffer` field. + // See `buffer` field documentation. + let buffer = OutputBuffer::new(buffer_storage.spare_capacity_mut()); + + CodedOutputStream { + target: OutputTarget::Write(writer, buffer_storage), + buffer, + pos_of_buffer_start: 0, + } + } + + /// `CodedOutputStream` which writes directly to bytes. + /// + /// Attempt to write more than bytes capacity results in error. + pub fn bytes(bytes: &'a mut [u8]) -> CodedOutputStream<'a> { + // SAFETY: it is safe to cast from &mut [u8] to &mut [MaybeUninit<u8>]. + let buffer = + ptr::slice_from_raw_parts_mut(bytes.as_mut_ptr() as *mut MaybeUninit<u8>, bytes.len()); + let buffer = OutputBuffer::new(buffer); + CodedOutputStream { + target: OutputTarget::Bytes, + buffer, + pos_of_buffer_start: 0, + } + } + + /// `CodedOutputStream` which writes directly to `Vec<u8>`. + pub fn vec(vec: &'a mut Vec<u8>) -> CodedOutputStream<'a> { + let buffer = OutputBuffer::new(vec.spare_capacity_mut()); + CodedOutputStream { + target: OutputTarget::Vec(vec), + buffer, + pos_of_buffer_start: 0, + } + } + + pub(crate) fn reserve_additional( + &mut self, + additional: u32, + message: &str, + ) -> crate::Result<()> { + if additional as usize <= self.buffer.unfilled_len() { + return Ok(()); + } + match &mut self.target { + OutputTarget::Write(..) => Ok(()), + OutputTarget::Vec(v) => { + let reserve = (additional as usize) + .checked_add(self.buffer.pos_within_buf()) + .unwrap(); + v.reserve(reserve); + // `pos_within_buf` remains unchanged. + self.buffer.replace_buffer_keep_pos(v.spare_capacity_mut()); + Ok(()) + } + OutputTarget::Bytes => { + Err(ProtobufError::BufferHasNotEnoughCapacity(message.to_owned()).into()) + } + } + } + + pub(crate) fn reserve_additional_for_length_delimited( + &mut self, + size: u32, + message: &str, + ) -> crate::Result<()> { + debug_assert!( + size <= MAX_MESSAGE_SIZE as u32, + "Caller of this function is responsible to guarantee \ + that message size does not exceed; size: {}, MAX_MESSAGE_SIZE: {}", + size, + MAX_MESSAGE_SIZE, + ); + let reserve = size + encoded_varint64_len(size as u64) as u32; + self.reserve_additional(reserve, message) + } + + /// Total number of bytes written to this stream. + /// + /// This number may be larger than the actual number of bytes written to the underlying stream, + /// if the buffer was not flushed. + /// + /// The number may be inaccurate if there was an error during the write. + pub fn total_bytes_written(&self) -> u64 { + self.pos_of_buffer_start + self.buffer.pos_within_buf() as u64 + } + + /// Check if EOF is reached. + /// + /// # Panics + /// + /// If underlying write has no EOF + pub fn check_eof(&self) { + match self.target { + OutputTarget::Bytes => { + assert_eq!( + self.buffer.buffer().len() as u64, + self.buffer.pos_within_buf() as u64 + ); + } + OutputTarget::Write(..) | OutputTarget::Vec(..) => { + panic!("must not be called with Writer or Vec"); + } + } + } + + fn refresh_buffer(&mut self) -> crate::Result<()> { + match self.target { + OutputTarget::Write(ref mut write, _) => { + write.write_all(self.buffer.filled())?; + self.pos_of_buffer_start += self.buffer.pos_within_buf() as u64; + self.buffer.rewind(); + } + OutputTarget::Vec(ref mut vec) => unsafe { + let vec_len = vec.len(); + assert!(vec_len + self.buffer.pos_within_buf() <= vec.capacity()); + vec.set_len(vec_len + self.buffer.pos_within_buf()); + vec.reserve(1); + self.pos_of_buffer_start += self.buffer.pos_within_buf() as u64; + self.buffer = OutputBuffer::new(vec.spare_capacity_mut()); + }, + OutputTarget::Bytes => { + return Err(ProtobufError::IoError(io::Error::new( + io::ErrorKind::Other, + "given slice is too small to serialize the message", + )) + .into()); + } + } + Ok(()) + } + + /// Flush to buffer to the underlying buffer. + /// Note that `CodedOutputStream` does `flush` in the destructor, + /// however, if `flush` in destructor fails, then destructor panics + /// and program terminates. So it's advisable to explicitly call flush + /// before destructor. + pub fn flush(&mut self) -> crate::Result<()> { + match &mut self.target { + OutputTarget::Bytes => Ok(()), + OutputTarget::Vec(vec) => { + let vec_len = vec.len(); + assert!(vec_len + self.buffer.pos_within_buf() <= vec.capacity()); + unsafe { + vec.set_len(vec_len + self.buffer.pos_within_buf()); + } + self.pos_of_buffer_start += self.buffer.pos_within_buf() as u64; + self.buffer = OutputBuffer::new(vec.spare_capacity_mut()); + Ok(()) + } + OutputTarget::Write(..) => self.refresh_buffer(), + } + } + + /// Write a byte + pub fn write_raw_byte(&mut self, byte: u8) -> crate::Result<()> { + if self.buffer.unfilled_len() == 0 { + self.refresh_buffer()?; + } + unsafe { self.buffer.write_byte(byte) }; + Ok(()) + } + + /// Write bytes + pub fn write_raw_bytes(&mut self, bytes: &[u8]) -> crate::Result<()> { + if bytes.len() <= self.buffer.unfilled_len() { + // SAFETY: we've just checked that there's enough space in the buffer. + unsafe { self.buffer.write_bytes(bytes) }; + return Ok(()); + } + + self.refresh_buffer()?; + + assert!(self.buffer.pos_within_buf() == 0); + + if bytes.len() <= self.buffer.unfilled_len() { + // SAFETY: we've just checked that there's enough space in the buffer. + unsafe { self.buffer.write_bytes(bytes) }; + return Ok(()); + } + + match self.target { + OutputTarget::Bytes => { + unreachable!(); + } + OutputTarget::Write(ref mut write, _) => { + write.write_all(bytes)?; + } + OutputTarget::Vec(ref mut vec) => { + assert!(self.buffer.pos_within_buf() == 0); + vec.extend(bytes); + self.buffer = OutputBuffer::new(vec.spare_capacity_mut()); + self.pos_of_buffer_start += bytes.len() as u64; + } + } + Ok(()) + } + + /// Write a tag + pub fn write_tag(&mut self, field_number: u32, wire_type: WireType) -> crate::Result<()> { + self.write_raw_varint32(wire_format::Tag::make(field_number, wire_type).value()) + } + + /// Write varint + pub fn write_raw_varint32(&mut self, value: u32) -> crate::Result<()> { + if self.buffer.unfilled_len() >= 5 { + // fast path + unsafe { + let len = encode_varint32(value, self.buffer.unfilled()); + self.buffer.advance(len); + }; + Ok(()) + } else { + // slow path + let buf = &mut [0u8; 5]; + let len = encode_varint32(value, unsafe { + slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut MaybeUninit<u8>, buf.len()) + }); + self.write_raw_bytes(&buf[..len]) + } + } + + /// Write varint + pub fn write_raw_varint64(&mut self, value: u64) -> crate::Result<()> { + if self.buffer.unfilled_len() >= MAX_VARINT_ENCODED_LEN { + // fast path + unsafe { + let len = encode_varint64(value, self.buffer.unfilled()); + self.buffer.advance(len); + }; + Ok(()) + } else { + // slow path + let buf = &mut [0u8; MAX_VARINT_ENCODED_LEN]; + let len = encode_varint64(value, unsafe { + slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut MaybeUninit<u8>, buf.len()) + }); + self.write_raw_bytes(&buf[..len]) + } + } + + /// Write 32-bit integer little endian + pub fn write_raw_little_endian32(&mut self, value: u32) -> crate::Result<()> { + self.write_raw_bytes(&value.to_le_bytes()) + } + + /// Write 64-bit integer little endian + pub fn write_raw_little_endian64(&mut self, value: u64) -> crate::Result<()> { + self.write_raw_bytes(&value.to_le_bytes()) + } + + /// Write `float` + pub fn write_float_no_tag(&mut self, value: f32) -> crate::Result<()> { + self.write_raw_little_endian32(value.to_bits()) + } + + /// Write `double` + pub fn write_double_no_tag(&mut self, value: f64) -> crate::Result<()> { + self.write_raw_little_endian64(value.to_bits()) + } + + /// Write `float` field + pub fn write_float(&mut self, field_number: u32, value: f32) -> crate::Result<()> { + self.write_tag(field_number, WireType::Fixed32)?; + self.write_float_no_tag(value)?; + Ok(()) + } + + /// Write `double` field + pub fn write_double(&mut self, field_number: u32, value: f64) -> crate::Result<()> { + self.write_tag(field_number, WireType::Fixed64)?; + self.write_double_no_tag(value)?; + Ok(()) + } + + /// Write varint + pub fn write_uint64_no_tag(&mut self, value: u64) -> crate::Result<()> { + self.write_raw_varint64(value) + } + + /// Write varint + pub fn write_uint32_no_tag(&mut self, value: u32) -> crate::Result<()> { + self.write_raw_varint32(value) + } + + /// Write varint + pub fn write_int64_no_tag(&mut self, value: i64) -> crate::Result<()> { + self.write_raw_varint64(value as u64) + } + + /// Write varint + pub fn write_int32_no_tag(&mut self, value: i32) -> crate::Result<()> { + self.write_raw_varint64(value as u64) + } + + /// Write zigzag varint + pub fn write_sint64_no_tag(&mut self, value: i64) -> crate::Result<()> { + self.write_uint64_no_tag(encode_zig_zag_64(value)) + } + + /// Write zigzag varint + pub fn write_sint32_no_tag(&mut self, value: i32) -> crate::Result<()> { + self.write_uint32_no_tag(encode_zig_zag_32(value)) + } + + /// Write `fixed64` + pub fn write_fixed64_no_tag(&mut self, value: u64) -> crate::Result<()> { + self.write_raw_little_endian64(value) + } + + /// Write `fixed32` + pub fn write_fixed32_no_tag(&mut self, value: u32) -> crate::Result<()> { + self.write_raw_little_endian32(value) + } + + /// Write `sfixed64` + pub fn write_sfixed64_no_tag(&mut self, value: i64) -> crate::Result<()> { + self.write_raw_little_endian64(value as u64) + } + + /// Write `sfixed32` + pub fn write_sfixed32_no_tag(&mut self, value: i32) -> crate::Result<()> { + self.write_raw_little_endian32(value as u32) + } + + /// Write `bool` + pub fn write_bool_no_tag(&mut self, value: bool) -> crate::Result<()> { + self.write_raw_varint32(if value { 1 } else { 0 }) + } + + /// Write `enum` + pub fn write_enum_no_tag(&mut self, value: i32) -> crate::Result<()> { + self.write_int32_no_tag(value) + } + + /// Write `enum` + pub fn write_enum_obj_no_tag<E>(&mut self, value: E) -> crate::Result<()> + where + E: Enum, + { + self.write_enum_no_tag(value.value()) + } + + /// Write `enum` + pub fn write_enum_or_unknown_no_tag<E>(&mut self, value: EnumOrUnknown<E>) -> crate::Result<()> + where + E: Enum, + { + self.write_enum_no_tag(value.value()) + } + + /// Write unknown value + pub fn write_unknown_no_tag(&mut self, unknown: UnknownValueRef) -> crate::Result<()> { + match unknown { + UnknownValueRef::Fixed64(fixed64) => self.write_raw_little_endian64(fixed64), + UnknownValueRef::Fixed32(fixed32) => self.write_raw_little_endian32(fixed32), + UnknownValueRef::Varint(varint) => self.write_raw_varint64(varint), + UnknownValueRef::LengthDelimited(bytes) => self.write_bytes_no_tag(bytes), + } + } + + /// Write `uint64` field + pub fn write_uint64(&mut self, field_number: u32, value: u64) -> crate::Result<()> { + self.write_tag(field_number, WireType::Varint)?; + self.write_uint64_no_tag(value)?; + Ok(()) + } + + /// Write `uint32` field + pub fn write_uint32(&mut self, field_number: u32, value: u32) -> crate::Result<()> { + self.write_tag(field_number, WireType::Varint)?; + self.write_uint32_no_tag(value)?; + Ok(()) + } + + /// Write `int64` field + pub fn write_int64(&mut self, field_number: u32, value: i64) -> crate::Result<()> { + self.write_tag(field_number, WireType::Varint)?; + self.write_int64_no_tag(value)?; + Ok(()) + } + + /// Write `int32` field + pub fn write_int32(&mut self, field_number: u32, value: i32) -> crate::Result<()> { + self.write_tag(field_number, WireType::Varint)?; + self.write_int32_no_tag(value)?; + Ok(()) + } + + /// Write `sint64` field + pub fn write_sint64(&mut self, field_number: u32, value: i64) -> crate::Result<()> { + self.write_tag(field_number, WireType::Varint)?; + self.write_sint64_no_tag(value)?; + Ok(()) + } + + /// Write `sint32` field + pub fn write_sint32(&mut self, field_number: u32, value: i32) -> crate::Result<()> { + self.write_tag(field_number, WireType::Varint)?; + self.write_sint32_no_tag(value)?; + Ok(()) + } + + /// Write `fixed64` field + pub fn write_fixed64(&mut self, field_number: u32, value: u64) -> crate::Result<()> { + self.write_tag(field_number, WireType::Fixed64)?; + self.write_fixed64_no_tag(value)?; + Ok(()) + } + + /// Write `fixed32` field + pub fn write_fixed32(&mut self, field_number: u32, value: u32) -> crate::Result<()> { + self.write_tag(field_number, WireType::Fixed32)?; + self.write_fixed32_no_tag(value)?; + Ok(()) + } + + /// Write `sfixed64` field + pub fn write_sfixed64(&mut self, field_number: u32, value: i64) -> crate::Result<()> { + self.write_tag(field_number, WireType::Fixed64)?; + self.write_sfixed64_no_tag(value)?; + Ok(()) + } + + /// Write `sfixed32` field + pub fn write_sfixed32(&mut self, field_number: u32, value: i32) -> crate::Result<()> { + self.write_tag(field_number, WireType::Fixed32)?; + self.write_sfixed32_no_tag(value)?; + Ok(()) + } + + /// Write `bool` field + pub fn write_bool(&mut self, field_number: u32, value: bool) -> crate::Result<()> { + self.write_tag(field_number, WireType::Varint)?; + self.write_bool_no_tag(value)?; + Ok(()) + } + + /// Write `enum` field + pub fn write_enum(&mut self, field_number: u32, value: i32) -> crate::Result<()> { + self.write_tag(field_number, WireType::Varint)?; + self.write_enum_no_tag(value)?; + Ok(()) + } + + /// Write `enum` field + pub fn write_enum_obj<E>(&mut self, field_number: u32, value: E) -> crate::Result<()> + where + E: Enum, + { + self.write_enum(field_number, value.value()) + } + + /// Write `enum` field + pub fn write_enum_or_unknown<E>( + &mut self, + field_number: u32, + value: EnumOrUnknown<E>, + ) -> crate::Result<()> + where + E: Enum, + { + self.write_enum(field_number, value.value()) + } + + /// Write unknown field + pub fn write_unknown( + &mut self, + field_number: u32, + value: UnknownValueRef, + ) -> crate::Result<()> { + self.write_tag(field_number, value.wire_type())?; + self.write_unknown_no_tag(value)?; + Ok(()) + } + + /// Write unknown fields + pub fn write_unknown_fields(&mut self, fields: &UnknownFields) -> crate::Result<()> { + for (number, value) in fields { + self.write_unknown(number, value)?; + } + Ok(()) + } + + /// Write unknown fields sorting them by name + // TODO: make unknown fields deterministic and remove this. + pub(crate) fn write_unknown_fields_sorted( + &mut self, + fields: &UnknownFields, + ) -> crate::Result<()> { + let mut fields: Vec<_> = fields.iter().collect(); + fields.sort_by_key(|(n, _)| *n); + for (number, value) in fields { + self.write_unknown(number, value)?; + } + Ok(()) + } + + /// Write bytes + pub fn write_bytes_no_tag(&mut self, bytes: &[u8]) -> crate::Result<()> { + self.write_raw_varint32(bytes.len() as u32)?; + self.write_raw_bytes(bytes)?; + Ok(()) + } + + /// Write string + pub fn write_string_no_tag(&mut self, s: &str) -> crate::Result<()> { + self.write_bytes_no_tag(s.as_bytes()) + } + + /// Write message + pub fn write_message_no_tag<M: Message>(&mut self, msg: &M) -> crate::Result<()> { + msg.write_length_delimited_to(self) + } + + /// Write dynamic message + pub fn write_message_no_tag_dyn(&mut self, msg: &dyn MessageDyn) -> crate::Result<()> { + let size = msg.compute_size_dyn(); + let size = check_message_size(size)?; + self.write_raw_varint32(size)?; + msg.write_to_dyn(self)?; + Ok(()) + } + + /// Write `bytes` field + pub fn write_bytes(&mut self, field_number: u32, bytes: &[u8]) -> crate::Result<()> { + self.write_tag(field_number, WireType::LengthDelimited)?; + self.write_bytes_no_tag(bytes)?; + Ok(()) + } + + /// Write `string` field + pub fn write_string(&mut self, field_number: u32, s: &str) -> crate::Result<()> { + self.write_tag(field_number, WireType::LengthDelimited)?; + self.write_string_no_tag(s)?; + Ok(()) + } + + /// Write repeated packed float values. + pub fn write_repeated_packed_float_no_tag(&mut self, values: &[f32]) -> crate::Result<()> { + if LITTLE_ENDIAN { + // SAFETY: it is safe to transmute floats to bytes. + let bytes = + unsafe { slice::from_raw_parts(values.as_ptr() as *const u8, values.len() * 4) }; + self.write_raw_bytes(bytes) + } else { + for v in values { + self.write_float_no_tag(*v)?; + } + Ok(()) + } + } + + /// Write field header and data for repeated packed float. + pub fn write_repeated_packed_float( + &mut self, + field_number: u32, + values: &[f32], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_fixed_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_float_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed double values. + pub fn write_repeated_packed_double_no_tag(&mut self, values: &[f64]) -> crate::Result<()> { + if LITTLE_ENDIAN { + // SAFETY: it is safe to transmute doubles to bytes. + let bytes = + unsafe { slice::from_raw_parts(values.as_ptr() as *const u8, values.len() * 8) }; + self.write_raw_bytes(bytes) + } else { + for v in values { + self.write_double_no_tag(*v)?; + } + Ok(()) + } + } + + /// Write field header and data for repeated packed double. + pub fn write_repeated_packed_double( + &mut self, + field_number: u32, + values: &[f64], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_fixed_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_double_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed fixed32 values. + pub fn write_repeated_packed_fixed32_no_tag(&mut self, values: &[u32]) -> crate::Result<()> { + if LITTLE_ENDIAN { + // SAFETY: it is safe to transmute integer to bytes. + let bytes = + unsafe { slice::from_raw_parts(values.as_ptr() as *const u8, values.len() * 4) }; + self.write_raw_bytes(bytes)?; + } else { + for v in values { + self.write_fixed32_no_tag(*v)?; + } + } + Ok(()) + } + + /// Write field header and data for repeated packed fixed32. + pub fn write_repeated_packed_fixed32( + &mut self, + field_number: u32, + values: &[u32], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_fixed_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_fixed32_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed fixed64 values. + pub fn write_repeated_packed_fixed64_no_tag(&mut self, values: &[u64]) -> crate::Result<()> { + if LITTLE_ENDIAN { + // SAFETY: it is safe to transmute integer to bytes. + let bytes = + unsafe { slice::from_raw_parts(values.as_ptr() as *const u8, values.len() * 8) }; + self.write_raw_bytes(bytes)?; + } else { + for v in values { + self.write_fixed64_no_tag(*v)?; + } + } + Ok(()) + } + + /// Write field header and data for repeated packed fixed64. + pub fn write_repeated_packed_fixed64( + &mut self, + field_number: u32, + values: &[u64], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_fixed_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_fixed64_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed sfixed32 values. + pub fn write_repeated_packed_sfixed32_no_tag(&mut self, values: &[i32]) -> crate::Result<()> { + if LITTLE_ENDIAN { + // SAFETY: it is safe to transmute integer to bytes. + let bytes = + unsafe { slice::from_raw_parts(values.as_ptr() as *const u8, values.len() * 4) }; + self.write_raw_bytes(bytes)?; + } else { + for v in values { + self.write_sfixed32_no_tag(*v)?; + } + } + Ok(()) + } + + /// Write field header and data for repeated packed sfixed32. + pub fn write_repeated_packed_sfixed32( + &mut self, + field_number: u32, + values: &[i32], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_fixed_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_sfixed32_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed sfixed64 values. + pub fn write_repeated_packed_sfixed64_no_tag(&mut self, values: &[i64]) -> crate::Result<()> { + if LITTLE_ENDIAN { + // SAFETY: it is safe to transmute integer to bytes. + let bytes = + unsafe { slice::from_raw_parts(values.as_ptr() as *const u8, values.len() * 8) }; + self.write_raw_bytes(bytes)?; + } else { + for v in values { + self.write_sfixed64_no_tag(*v)?; + } + } + Ok(()) + } + + /// Write field header and data for repeated packed sfixed64. + pub fn write_repeated_packed_sfixed64( + &mut self, + field_number: u32, + values: &[i64], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_fixed_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_sfixed64_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed int32 values. + pub fn write_repeated_packed_int32_no_tag(&mut self, values: &[i32]) -> crate::Result<()> { + for v in values { + self.write_int32_no_tag(*v)?; + } + Ok(()) + } + + /// Write field header and data for repeated packed int32. + pub fn write_repeated_packed_int32( + &mut self, + field_number: u32, + values: &[i32], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_varint_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_int32_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed int64 values. + pub fn write_repeated_packed_int64_no_tag(&mut self, values: &[i64]) -> crate::Result<()> { + for v in values { + self.write_int64_no_tag(*v)?; + } + Ok(()) + } + + /// Write field header and data for repeated packed int64. + pub fn write_repeated_packed_int64( + &mut self, + field_number: u32, + values: &[i64], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_varint_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_int64_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed uint32 values. + pub fn write_repeated_packed_uint32_no_tag(&mut self, values: &[u32]) -> crate::Result<()> { + for v in values { + self.write_uint32_no_tag(*v)?; + } + Ok(()) + } + + /// Write field header and data for repeated packed uint32. + pub fn write_repeated_packed_uint32( + &mut self, + field_number: u32, + values: &[u32], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_varint_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_uint32_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed uint64 values. + pub fn write_repeated_packed_uint64_no_tag(&mut self, values: &[u64]) -> crate::Result<()> { + for v in values { + self.write_uint64_no_tag(*v)?; + } + Ok(()) + } + + /// Write field header and data for repeated packed uint64. + pub fn write_repeated_packed_uint64( + &mut self, + field_number: u32, + values: &[u64], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_varint_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_uint64_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed sint32 values. + pub fn write_repeated_packed_sint32_no_tag(&mut self, values: &[i32]) -> crate::Result<()> { + for v in values { + self.write_sint32_no_tag(*v)?; + } + Ok(()) + } + + /// Write field header and data for repeated packed sint32. + pub fn write_repeated_packed_sint32( + &mut self, + field_number: u32, + values: &[i32], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_varint_zigzag_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_sint32_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed sint64 values. + pub fn write_repeated_packed_sint64_no_tag(&mut self, values: &[i64]) -> crate::Result<()> { + for v in values { + self.write_sint64_no_tag(*v)?; + } + Ok(()) + } + + /// Write field header and data for repeated packed sint64. + pub fn write_repeated_packed_sint64( + &mut self, + field_number: u32, + values: &[i64], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_varint_zigzag_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_sint64_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed bool values. + pub fn write_repeated_packed_bool_no_tag(&mut self, values: &[bool]) -> crate::Result<()> { + for v in values { + self.write_bool_no_tag(*v)?; + } + Ok(()) + } + + /// Write field header and data for repeated packed bool. + pub fn write_repeated_packed_bool( + &mut self, + field_number: u32, + values: &[bool], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_fixed_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_bool_no_tag(values)?; + Ok(()) + } + + /// Write repeated packed enum values. + #[inline] + pub fn write_repeated_packed_enum_or_unknown_no_tag<E: Enum>( + &mut self, + values: &[EnumOrUnknown<E>], + ) -> crate::Result<()> { + self.write_repeated_packed_int32_no_tag(EnumOrUnknown::cast_to_values(values)) + } + + /// Write field header and data for repeated packed enum. + pub fn write_repeated_packed_enum_or_unknown<E: Enum>( + &mut self, + field_number: u32, + values: &[EnumOrUnknown<E>], + ) -> crate::Result<()> { + if values.is_empty() { + return Ok(()); + } + self.write_tag(field_number, WireType::LengthDelimited)?; + let data_size = vec_packed_enum_or_unknown_data_size(values); + self.write_raw_varint32(data_size as u32)?; + self.write_repeated_packed_enum_or_unknown_no_tag(values)?; + Ok(()) + } + + /// Write `message` field + pub fn write_message<M: MessageFull>( + &mut self, + field_number: u32, + msg: &M, + ) -> crate::Result<()> { + self.write_tag(field_number, WireType::LengthDelimited)?; + self.write_message_no_tag(msg)?; + Ok(()) + } + + /// Write dynamic `message` field + pub fn write_message_dyn( + &mut self, + field_number: u32, + msg: &dyn MessageDyn, + ) -> crate::Result<()> { + self.write_tag(field_number, WireType::LengthDelimited)?; + self.write_message_no_tag_dyn(msg)?; + Ok(()) + } +} + +impl<'a> Write for CodedOutputStream<'a> { + fn write(&mut self, buf: &[u8]) -> io::Result<usize> { + self.write_raw_bytes(buf)?; + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + CodedOutputStream::flush(self).map_err(Into::into) + } +} + +impl<'a> Drop for CodedOutputStream<'a> { + fn drop(&mut self) { + // This may panic + CodedOutputStream::flush(self).expect("failed to flush"); + } +} + +#[cfg(test)] +mod test { + use std::iter; + + use super::*; + use crate::hex::decode_hex; + use crate::hex::encode_hex; + + fn test_write<F>(expected: &str, mut gen: F) + where + F: FnMut(&mut CodedOutputStream) -> crate::Result<()>, + { + let expected_bytes = decode_hex(expected); + + // write to Write + { + let mut v = Vec::new(); + { + let mut os = CodedOutputStream::new(&mut v as &mut dyn Write); + gen(&mut os).unwrap(); + os.flush().unwrap(); + } + assert_eq!(encode_hex(&expected_bytes), encode_hex(&v)); + } + + // write to &[u8] + { + let mut r = Vec::with_capacity(expected_bytes.len()); + r.resize(expected_bytes.len(), 0); + { + let mut os = CodedOutputStream::bytes(&mut r); + gen(&mut os).unwrap(); + os.check_eof(); + } + assert_eq!(encode_hex(&expected_bytes), encode_hex(&r)); + } + + // write to Vec<u8> + { + let mut r = Vec::new(); + r.extend(&[11, 22, 33, 44, 55, 66, 77]); + { + let mut os = CodedOutputStream::vec(&mut r); + gen(&mut os).unwrap(); + os.flush().unwrap(); + } + + r.drain(..7); + assert_eq!(encode_hex(&expected_bytes), encode_hex(&r)); + } + } + + #[test] + fn test_output_stream_write_raw_byte() { + test_write("a1", |os| os.write_raw_byte(0xa1)); + } + + #[test] + fn test_output_stream_write_tag() { + test_write("08", |os| os.write_tag(1, WireType::Varint)); + } + + #[test] + #[cfg_attr(miri, ignore)] // Miri is too slow for this test. + fn test_output_stream_write_raw_bytes() { + test_write("00 ab", |os| os.write_raw_bytes(&[0x00, 0xab])); + + let expected = iter::repeat("01 02 03 04") + .take(2048) + .collect::<Vec<_>>() + .join(" "); + test_write(&expected, |os| { + for _ in 0..2048 { + os.write_raw_bytes(&[0x01, 0x02, 0x03, 0x04])?; + } + + Ok(()) + }); + } + + #[test] + fn test_output_stream_write_raw_varint32() { + test_write("96 01", |os| os.write_raw_varint32(150)); + test_write("ff ff ff ff 0f", |os| os.write_raw_varint32(0xffffffff)); + } + + #[test] + fn test_output_stream_write_raw_varint64() { + test_write("96 01", |os| os.write_raw_varint64(150)); + test_write("ff ff ff ff ff ff ff ff ff 01", |os| { + os.write_raw_varint64(0xffffffffffffffff) + }); + } + + #[test] + fn test_output_stream_write_int32_no_tag() { + test_write("ff ff ff ff ff ff ff ff ff 01", |os| { + os.write_int32_no_tag(-1) + }); + } + + #[test] + fn test_output_stream_write_int64_no_tag() { + test_write("ff ff ff ff ff ff ff ff ff 01", |os| { + os.write_int64_no_tag(-1) + }); + } + + #[test] + fn test_output_stream_write_raw_little_endian32() { + test_write("f1 e2 d3 c4", |os| os.write_raw_little_endian32(0xc4d3e2f1)); + } + + #[test] + fn test_output_stream_write_float_no_tag() { + test_write("95 73 13 61", |os| os.write_float_no_tag(17e19)); + } + + #[test] + fn test_output_stream_write_double_no_tag() { + test_write("40 d5 ab 68 b3 07 3d 46", |os| { + os.write_double_no_tag(23e29) + }); + } + + #[test] + fn test_output_stream_write_raw_little_endian64() { + test_write("f1 e2 d3 c4 b5 a6 07 f8", |os| { + os.write_raw_little_endian64(0xf807a6b5c4d3e2f1) + }); + } + + #[test] + fn test_output_stream_io_write() { + let expected = [0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77]; + + // write to Write + { + let mut v = Vec::new(); + { + let mut os = CodedOutputStream::new(&mut v as &mut dyn Write); + Write::write(&mut os, &expected).expect("io::Write::write"); + Write::flush(&mut os).expect("io::Write::flush"); + } + assert_eq!(expected, *v); + } + + // write to &[u8] + { + let mut v = Vec::with_capacity(expected.len()); + v.resize(expected.len(), 0); + { + let mut os = CodedOutputStream::bytes(&mut v); + Write::write(&mut os, &expected).expect("io::Write::write"); + Write::flush(&mut os).expect("io::Write::flush"); + os.check_eof(); + } + assert_eq!(expected, *v); + } + + // write to Vec<u8> + { + let mut v = Vec::new(); + { + let mut os = CodedOutputStream::vec(&mut v); + Write::write(&mut os, &expected).expect("io::Write::write"); + Write::flush(&mut os).expect("io::Write::flush"); + } + assert_eq!(expected, *v); + } + } + + #[test] + fn flush_for_vec_does_not_allocate_more() { + let mut v = Vec::with_capacity(10); + { + let mut os = CodedOutputStream::vec(&mut v); + for i in 0..10 { + os.write_raw_byte(i as u8).unwrap(); + } + os.flush().unwrap(); + } + assert_eq!(10, v.len()); + // Previously, this allocated more data in buf. + assert_eq!(10, v.capacity()); + } + + #[test] + fn total_bytes_written_to_bytes() { + let mut buf = vec![0; 10]; + let mut stream = CodedOutputStream::bytes(&mut buf); + assert_eq!(0, stream.total_bytes_written()); + stream.write_raw_bytes(&[11, 22]).unwrap(); + assert_eq!(2, stream.total_bytes_written()); + stream.write_raw_bytes(&[33, 44, 55]).unwrap(); + assert_eq!(5, stream.total_bytes_written()); + } + + #[test] + fn total_bytes_written_to_vec() { + let mut buf = Vec::new(); + let mut stream = CodedOutputStream::vec(&mut buf); + for i in 0..100 { + stream.write_raw_bytes(&[0, 1, 2]).unwrap(); + assert_eq!((i + 1) * 3, stream.total_bytes_written()); + } + } +} |