byte_buffer: use slices instead of Vec

This commit is contained in:
Matt Bilker
2018-07-11 22:10:53 -04:00
parent 3df2b1eda1
commit f4e34082bc
4 changed files with 59 additions and 39 deletions

View File

@@ -1,5 +1,5 @@
use std::cmp::max;
use std::io::{Cursor, Read, Seek, SeekFrom, Write};
use std::io::{Cursor, Seek, SeekFrom, Write};
use std::ops::{Deref, DerefMut};
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
@@ -10,8 +10,8 @@ use node_types::KbinType;
pub use encoding_type::EncodingType;
pub use error::{KbinError, KbinErrorKind, Result};
pub struct ByteBufferRead<R: AsRef<[u8]>> {
buffer: Cursor<R>,
pub struct ByteBufferRead<'buf> {
buffer: Cursor<&'buf [u8]>,
offset_1: u64,
offset_2: u64,
}
@@ -22,8 +22,8 @@ pub struct ByteBufferWrite {
offset_2: u64,
}
impl<R: AsRef<[u8]>> ByteBufferRead<R> {
pub fn new(buffer: R) -> Self {
impl<'buf> ByteBufferRead<'buf> {
pub fn new(buffer: &'buf [u8]) -> Self {
Self {
buffer: Cursor::new(buffer),
offset_1: 0,
@@ -41,21 +41,34 @@ impl<R: AsRef<[u8]>> ByteBufferRead<R> {
self.buffer.position()
}
pub fn buf_read(&mut self) -> Result<Vec<u8>> {
fn buf_read_size(&mut self, size: usize) -> Result<&'buf [u8]> {
// To avoid an allocation of a `Vec` here, the raw input byte array is used
let start = self.buffer.position() as usize;
let end = start + size as usize;
if end > self.buffer.get_ref().len() {
return Err(KbinErrorKind::DataRead(size as usize).into());
}
let data = &self.buffer.get_ref()[start..end];
trace!("buf_read_size => index: {}, size: {}, data: 0x{:02x?}", self.buffer.position(), data.len(), data);
self.buffer.seek(SeekFrom::Current(size as i64)).context(KbinErrorKind::DataRead(size as usize))?;
Ok(data)
}
pub fn buf_read(&mut self) -> Result<&'buf [u8]> {
let size = self.buffer.read_u32::<BigEndian>().context(KbinErrorKind::DataReadSize)?;
debug!("buf_read => index: {}, size: {}", self.buffer.position(), size);
let mut data = vec![0; size as usize];
self.buffer.read_exact(&mut data).context(KbinErrorKind::DataRead(size as usize))?;
trace!("buf_read => index: {}, size: {}, data: 0x{:02x?}", self.buffer.position(), data.len(), data);
let data = self.buf_read_size(size as usize)?;
self.realign_reads(None)?;
Ok(data)
}
pub fn read_str(&mut self, encoding: EncodingType) -> Result<String> {
let mut data = self.buf_read()?;
let data = self.buf_read()?;
// Remove trailing null bytes
let mut index = data.len() - 1;
@@ -63,20 +76,20 @@ impl<R: AsRef<[u8]>> ByteBufferRead<R> {
while index > 0 && index < len && data[index] == 0x00 {
index -= 1;
}
data.truncate(index + 1);
let data = &data[..=index];
trace!("read_str => size: {}, data: 0x{:02x?}", data.len(), data);
encoding.decode_bytes(data)
}
pub fn get(&mut self, size: u32) -> Result<Vec<u8>> {
let mut data = vec![0; size as usize];
self.buffer.read_exact(&mut data).context(KbinErrorKind::DataRead(size as usize))?;
pub fn get(&mut self, size: u32) -> Result<&'buf [u8]> {
let data = self.buf_read_size(size as usize)?;
trace!("get => size: {}, data: 0x{:02x?}", size, data);
Ok(data)
}
pub fn get_aligned(&mut self, data_type: KbinType) -> Result<Vec<u8>> {
pub fn get_aligned(&mut self, data_type: KbinType) -> Result<&'buf [u8]> {
if self.offset_1 % 4 == 0 {
self.offset_1 = self.data_buf_offset();
}
@@ -92,23 +105,21 @@ impl<R: AsRef<[u8]>> ByteBufferRead<R> {
1 => {
self.buffer.seek(SeekFrom::Start(self.offset_1)).context(KbinErrorKind::Seek)?;
let data = self.buffer.read_u8().context(KbinErrorKind::DataReadOneByte)?;
let data = self.buf_read_size(1).context(KbinErrorKind::DataReadOneByte)?;
self.offset_1 += 1;
(true, vec![data])
(true, data)
},
2 => {
self.buffer.seek(SeekFrom::Start(self.offset_2)).context(KbinErrorKind::Seek)?;
let mut data = vec![0; 2];
self.buffer.read_exact(&mut data).context(KbinErrorKind::DataReadTwoByte)?;
let data = self.buf_read_size(2).context(KbinErrorKind::DataReadTwoByte)?;
self.offset_2 += 2;
(true, data)
},
size => {
let mut data = vec![0; size as usize];
self.buffer.read_exact(&mut data).context(KbinErrorKind::DataReadAligned)?;
let data = self.buf_read_size(size as usize).context(KbinErrorKind::DataReadAligned)?;
self.realign_reads(None)?;
(false, data)
@@ -264,15 +275,15 @@ impl ByteBufferWrite {
}
}
impl<R> Deref for ByteBufferRead<R> where R: AsRef<[u8]> {
type Target = Cursor<R>;
impl<'buf> Deref for ByteBufferRead<'buf> {
type Target = Cursor<&'buf [u8]>;
fn deref(&self) -> &Self::Target {
&self.buffer
}
}
impl<R> DerefMut for ByteBufferRead<R> where R: AsRef<[u8]> {
impl<'buf> DerefMut for ByteBufferRead<'buf> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.buffer
}

View File

@@ -50,19 +50,19 @@ impl EncodingType {
///
/// A `Some` value indicates an encoding should be used from the `encoding`
/// crate. A `None` value indicates Rust's own UTF-8 handling should be used.
pub fn decode_bytes(&self, input: Vec<u8>) -> Result<String, KbinError> {
pub fn decode_bytes(&self, input: &[u8]) -> Result<String, KbinError> {
let decoder_fail = |e| {
format_err!("{}", e).context(KbinErrorKind::Encoding)
};
let result = match *self {
EncodingType::None |
EncodingType::UTF_8 => String::from_utf8(input)?,
EncodingType::UTF_8 => String::from_utf8(input.to_vec())?,
EncodingType::ASCII => ASCII.decode(&input, DecoderTrap::Strict).map_err(decoder_fail)?,
EncodingType::ISO_8859_1 => ISO_8859_1.decode(&input, DecoderTrap::Strict).map_err(decoder_fail)?,
EncodingType::EUC_JP => EUC_JP.decode(&input, DecoderTrap::Strict).map_err(decoder_fail)?,
EncodingType::SHIFT_JIS => WINDOWS_31J.decode(&input, DecoderTrap::Strict).map_err(decoder_fail)?,
EncodingType::ASCII => ASCII.decode(input, DecoderTrap::Strict).map_err(decoder_fail)?,
EncodingType::ISO_8859_1 => ISO_8859_1.decode(input, DecoderTrap::Strict).map_err(decoder_fail)?,
EncodingType::EUC_JP => EUC_JP.decode(input, DecoderTrap::Strict).map_err(decoder_fail)?,
EncodingType::SHIFT_JIS => WINDOWS_31J.decode(input, DecoderTrap::Strict).map_err(decoder_fail)?,
};
Ok(result)

View File

@@ -17,7 +17,7 @@ extern crate serde;
use std::fmt::Write as FmtWrite;
use std::io::{Cursor, Write};
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use byteorder::{BigEndian, WriteBytesExt};
use failure::ResultExt;
use minidom::Element;
use rustc_hex::FromHex;

View File

@@ -12,8 +12,8 @@ use super::{ARRAY_MASK, SIGNATURE, SIG_COMPRESSED};
pub struct Reader<'buf> {
encoding: EncodingType,
pub(crate) node_buf: ByteBufferRead<&'buf [u8]>,
pub(crate) data_buf: ByteBufferRead<&'buf [u8]>,
pub(crate) node_buf: ByteBufferRead<'buf>,
pub(crate) data_buf: ByteBufferRead<'buf>,
data_buf_start: u64,
}
@@ -81,13 +81,12 @@ impl<'buf> Reader<'buf> {
self.data_buf_start
}
pub fn read_node_type(&mut self) -> Result<(StandardType, bool)> {
let raw_node_type = self.node_buf.read_u8().context(KbinErrorKind::NodeTypeRead)?;
fn parse_node_type(&self, raw_node_type: u8) -> Result<(StandardType, bool)> {
let is_array = raw_node_type & ARRAY_MASK == ARRAY_MASK;
let node_type = raw_node_type & !ARRAY_MASK;
let xml_type = StandardType::from_u8(node_type);
debug!("Reader::read_node_type() => raw_node_type: {}, node_type: {:?} ({}), is_array: {}",
debug!("Reader::parse_node_type() => raw_node_type: {}, node_type: {:?} ({}), is_array: {}",
raw_node_type,
xml_type,
node_type,
@@ -96,6 +95,17 @@ impl<'buf> Reader<'buf> {
Ok((xml_type, is_array))
}
pub fn peek_node_type(&self) -> Result<(StandardType, bool)> {
let pos = self.node_buf.position();
let raw_node_type = self.node_buf.get_ref()[pos as usize];
self.parse_node_type(raw_node_type)
}
pub fn read_node_type(&mut self) -> Result<(StandardType, bool)> {
let raw_node_type = self.node_buf.read_u8().context(KbinErrorKind::NodeTypeRead)?;
self.parse_node_type(raw_node_type)
}
pub fn read_node_identifier(&mut self) -> Result<String> {
let value = unpack_sixbit(&mut *self.node_buf)?;
debug!("Reader::read_node_identifier() => value: {:?}", value);
@@ -124,9 +134,8 @@ impl<'buf> Reader<'buf> {
Ok(value)
}
// TODO: make a more intelligent reader to avoid allocating the Vec
#[inline]
pub fn read_bytes(&mut self) -> Result<Vec<u8>> {
pub fn read_bytes(&mut self) -> Result<&'buf [u8]> {
self.data_buf.buf_read()
}
}