use crate::AsyncBufRead;
use encoding_rs::{Decoder, Encoder, Encoding};
use futures_util::ready;
use std::fmt;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
pub struct CharCodec {
dec: Decoder,
enc: Option<Encoder>,
decoded: String,
is_end: bool,
}
#[allow(missing_docs)]
impl CharCodec {
pub fn new(from: &'static Encoding, to: &'static Encoding) -> CharCodec {
CharCodec {
dec: from.new_decoder(),
enc: if to == encoding_rs::UTF_8 {
None
} else {
Some(to.new_encoder())
},
decoded: String::new(),
is_end: false,
}
}
pub fn remove_encoder(&mut self) {
self.enc = None;
}
pub fn poll_codec<R: AsyncBufRead + Unpin>(
&mut self,
cx: &mut Context,
from: &mut R,
dst: &mut [u8],
) -> Poll<Result<usize, io::Error>> {
let src = ready!(Pin::new(&mut *from).poll_fill_buf(cx))?;
let mut consumed = 0;
let ret = self.decode_from_buf(src, dst, &mut consumed);
Pin::new(&mut *from).consume(consumed);
Poll::Ready(ret)
}
pub fn decode_from_buf(
&mut self,
src: &[u8],
dst: &mut [u8],
consumed: &mut usize,
) -> Result<usize, io::Error> {
let mut became_end = false;
if !self.is_end && src.is_empty() {
became_end = true;
self.is_end = true;
}
if (!self.is_end && self.decoded.len() < 128) || became_end {
let mut decode_to = [0_u8; 8_192];
let (_, decode_read, decode_written, decode_had_errors) =
self.dec.decode_to_utf8(src, &mut decode_to[..], became_end);
if decode_had_errors {
debug!("Character decoder had errors");
}
*consumed = decode_read;
let decoded = unsafe { std::str::from_utf8_unchecked(&decode_to[0..decode_written]) };
self.decoded.push_str(decoded);
}
if let Some(enc) = &mut self.enc {
let (_, encode_read, encode_written, encode_had_errors) =
enc.encode_from_utf8(&self.decoded[..], dst, self.is_end);
if encode_had_errors {
debug!("Character encoder had errors");
}
let rest = self.decoded.split_off(encode_read);
self.decoded = rest;
Ok(encode_written)
} else {
let decoded_bytes = self.decoded.as_bytes();
let mut max = decoded_bytes.len().min(dst.len());
while max > 0 && !self.decoded.is_char_boundary(max) {
max -= 1;
}
(&mut dst[0..max]).copy_from_slice(&decoded_bytes[0..max]);
let vec = unsafe { self.decoded.as_mut_vec() };
let rest = vec.split_off(max);
*vec = rest;
Ok(max)
}
}
}
impl fmt::Debug for CharCodec {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"CharCodec {{ from: {}, to: {} }}",
self.dec.encoding().name(),
self.enc
.as_ref()
.map(|e| e.encoding())
.unwrap_or(encoding_rs::UTF_8)
.name()
)
}
}