Fix varint::encode allocating for small inputs (#103)

* Fix `encode` allocating for small inputs

* Fix protocol listener expecting `Vec`
This commit is contained in:
Jef
2018-01-24 11:14:17 +01:00
committed by Pierre Krieger
parent 71dae91d15
commit cace5bf577
2 changed files with 14 additions and 12 deletions

View File

@ -113,11 +113,13 @@ impl<R> Sink for Listener<R>
} }
ListenerToDialerMessage::ProtocolsListResponse { list } => { ListenerToDialerMessage::ProtocolsListResponse { list } => {
use std::iter;
let mut out_msg = varint::encode(list.len()); let mut out_msg = varint::encode(list.len());
for elem in list.iter() { for elem in list.iter() {
out_msg.push(b'\r'); out_msg.extend(iter::once(b'\r'));
out_msg.extend_from_slice(elem); out_msg.extend_from_slice(elem);
out_msg.push(b'\n'); out_msg.extend(iter::once(b'\n'));
} }
match self.inner.start_send(BytesMut::from(out_msg)) { match self.inner.start_send(BytesMut::from(out_msg)) {

View File

@ -31,7 +31,7 @@ extern crate futures;
#[macro_use] #[macro_use]
extern crate error_chain; extern crate error_chain;
use bytes::{BufMut, BytesMut, IntoBuf}; use bytes::{BufMut, Bytes, BytesMut, IntoBuf};
use futures::{Poll, Async}; use futures::{Poll, Async};
use num_bigint::BigUint; use num_bigint::BigUint;
use num_traits::ToPrimitive; use num_traits::ToPrimitive;
@ -421,7 +421,7 @@ impl<T> Default for VarintCodec<T> {
enum VarintCodecInner { enum VarintCodecInner {
WaitingForLen(VarintDecoder<usize>), WaitingForLen(VarintDecoder<usize>),
WaitingForData(usize), WaitingForData(usize),
Poisonned, Poisoned,
} }
impl<T> Decoder for VarintCodec<T> { impl<T> Decoder for VarintCodec<T> {
@ -430,7 +430,7 @@ impl<T> Decoder for VarintCodec<T> {
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> { fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
loop { loop {
match mem::replace(&mut self.inner, VarintCodecInner::Poisonned) { match mem::replace(&mut self.inner, VarintCodecInner::Poisoned) {
VarintCodecInner::WaitingForData(len) => { VarintCodecInner::WaitingForData(len) => {
if src.len() >= len { if src.len() >= len {
self.inner = VarintCodecInner::WaitingForLen(VarintDecoder::default()); self.inner = VarintCodecInner::WaitingForLen(VarintDecoder::default());
@ -451,9 +451,7 @@ impl<T> Decoder for VarintCodec<T> {
}, },
} }
}, },
VarintCodecInner::Poisonned => { VarintCodecInner::Poisoned => panic!("varint codec was poisoned"),
panic!("varint codec was poisoned")
},
} }
} }
} }
@ -466,7 +464,7 @@ impl<D> Encoder for VarintCodec<D>
type Error = io::Error; type Error = io::Error;
fn encode(&mut self, item: D, dst: &mut BytesMut) -> Result<(), io::Error> { fn encode(&mut self, item: D, dst: &mut BytesMut) -> Result<(), io::Error> {
let encoded_len = encode(item.as_ref().len()); // TODO: can be optimized by not allocating? let encoded_len = encode(item.as_ref().len());
dst.put(encoded_len); dst.put(encoded_len);
dst.put(item); dst.put(item);
Ok(()) Ok(())
@ -495,12 +493,14 @@ pub fn decode<R: Read, T: Default + DecoderHelper>(mut input: R) -> errors::Resu
} }
/// Syncronously decode a number from a `Read` /// Syncronously decode a number from a `Read`
pub fn encode<T: EncoderHelper + Bits>(input: T) -> Vec<u8> { pub fn encode<T: EncoderHelper + Bits>(input: T) -> Bytes {
use tokio_io::io::AllowStdIo;
let mut encoder = EncoderState::new(input); let mut encoder = EncoderState::new(input);
let mut out = io::Cursor::new(Vec::with_capacity(1)); let mut out = AllowStdIo::new(BytesMut::new().writer());
match T::write(&mut encoder, &mut out).expect("Writing to a vec should never fail, Q.E.D") { match T::write(&mut encoder, &mut out).expect("Writing to a vec should never fail, Q.E.D") {
Async::Ready(_) => out.into_inner(), Async::Ready(_) => out.into_inner().into_inner().freeze(),
Async::NotReady => unreachable!(), Async::NotReady => unreachable!(),
} }
} }