use std::borrow::Cow; use encoding_rs::{ BIG5, CoderResult, EUC_KR, Encoding, GB18030, GBK, SHIFT_JIS, UTF_8, UTF_16BE, UTF_16LE, WINDOWS_1252, }; use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] #[serde(rename_all = "snake_case")] pub enum TextEncoding { #[default] Utf8, Gb18030, Gbk, Big5, ShiftJis, EucKr, Windows1252, Utf16Le, Utf16Be, } pub(crate) const TERMINAL_ENCODINGS: &[TextEncoding] = &[ TextEncoding::Utf8, TextEncoding::Gb18030, TextEncoding::Gbk, TextEncoding::Big5, TextEncoding::ShiftJis, TextEncoding::EucKr, TextEncoding::Windows1252, ]; pub(crate) const FILE_ENCODINGS: &[TextEncoding] = &[ TextEncoding::Utf8, TextEncoding::Utf16Le, TextEncoding::Utf16Be, TextEncoding::Gb18030, TextEncoding::Gbk, TextEncoding::Big5, TextEncoding::ShiftJis, TextEncoding::EucKr, TextEncoding::Windows1252, ]; impl TextEncoding { pub(crate) fn label(self) -> &'static str { match self { Self::Utf8 => "UTF-8", Self::Gb18030 => "GB18030", Self::Gbk => "GBK", Self::Big5 => "Big5", Self::ShiftJis => "Shift_JIS", Self::EucKr => "EUC-KR", Self::Windows1252 => "Windows-1252", Self::Utf16Le => "UTF-16 LE", Self::Utf16Be => "UTF-16 BE", } } pub(crate) fn encoding(self) -> &'static Encoding { match self { Self::Utf8 => UTF_8, Self::Gb18030 => GB18030, Self::Gbk => GBK, Self::Big5 => BIG5, Self::ShiftJis => SHIFT_JIS, Self::EucKr => EUC_KR, Self::Windows1252 => WINDOWS_1252, Self::Utf16Le => UTF_16LE, Self::Utf16Be => UTF_16BE, } } pub(crate) fn detect_bom(bytes: &[u8]) -> Option { if bytes.starts_with(&[0xEF, 0xBB, 0xBF]) { Some(Self::Utf8) } else if bytes.starts_with(&[0xFF, 0xFE]) { Some(Self::Utf16Le) } else if bytes.starts_with(&[0xFE, 0xFF]) { Some(Self::Utf16Be) } else { None } } pub(crate) fn decode_file(self, bytes: &[u8]) -> (String, bool, bool) { let bom_len = self.matching_bom_len(bytes); let (text, had_errors) = self .encoding() .decode_without_bom_handling(&bytes[bom_len..]); (text.into_owned(), had_errors, bom_len > 0) } pub(crate) fn encode_file(self, text: &str, with_bom: bool) -> (Vec, bool) { if matches!(self, Self::Utf16Le | Self::Utf16Be) { let mut bytes = Vec::with_capacity(text.len().saturating_mul(2).saturating_add(2)); if with_bom { bytes.extend_from_slice(if self == Self::Utf16Le { &[0xFF, 0xFE] } else { &[0xFE, 0xFF] }); } for code_unit in text.encode_utf16() { let encoded = if self == Self::Utf16Le { code_unit.to_le_bytes() } else { code_unit.to_be_bytes() }; bytes.extend_from_slice(&encoded); } return (bytes, false); } let (encoded, _, had_errors) = self.encoding().encode(text); let mut bytes = Vec::with_capacity(encoded.len() + 3); if with_bom { match self { Self::Utf8 => bytes.extend_from_slice(&[0xEF, 0xBB, 0xBF]), Self::Utf16Le => bytes.extend_from_slice(&[0xFF, 0xFE]), Self::Utf16Be => bytes.extend_from_slice(&[0xFE, 0xFF]), _ => {} } } bytes.extend_from_slice(&encoded); (bytes, had_errors) } pub(crate) fn default_bom(self) -> bool { matches!(self, Self::Utf16Le | Self::Utf16Be) } pub(crate) fn encode_terminal_input<'a>(self, bytes: &'a [u8]) -> Cow<'a, [u8]> { if self == Self::Utf8 { return Cow::Borrowed(bytes); } let Ok(text) = std::str::from_utf8(bytes) else { return Cow::Borrowed(bytes); }; let (encoded, _, _) = self.encoding().encode(text); encoded } fn matching_bom_len(self, bytes: &[u8]) -> usize { match self { Self::Utf8 if bytes.starts_with(&[0xEF, 0xBB, 0xBF]) => 3, Self::Utf16Le if bytes.starts_with(&[0xFF, 0xFE]) => 2, Self::Utf16Be if bytes.starts_with(&[0xFE, 0xFF]) => 2, _ => 0, } } } pub(crate) struct StreamingDecoder { decoder: encoding_rs::Decoder, } impl StreamingDecoder { pub(crate) fn new(encoding: TextEncoding) -> Self { Self { decoder: encoding.encoding().new_decoder_without_bom_handling(), } } pub(crate) fn decode(&mut self, bytes: &[u8]) -> Vec { let initial_capacity = self .decoder .max_utf8_buffer_length(bytes.len()) .unwrap_or_else(|| bytes.len().saturating_mul(4).saturating_add(16)); let mut output = String::with_capacity(initial_capacity); let mut remaining = bytes; loop { let (result, read, _) = self.decoder.decode_to_string(remaining, &mut output, false); remaining = &remaining[read..]; match result { CoderResult::InputEmpty => break, CoderResult::OutputFull => { let additional = self .decoder .max_utf8_buffer_length(remaining.len()) .unwrap_or_else(|| remaining.len().saturating_mul(4).saturating_add(16)) .max(16); output.reserve(additional); } } } output.into_bytes() } } #[cfg(test)] mod tests { use super::{StreamingDecoder, TextEncoding}; #[test] fn streaming_decoder_preserves_split_multibyte_characters() { let (encoded, _) = TextEncoding::Gbk.encode_file("中文", false); let mut decoder = StreamingDecoder::new(TextEncoding::Gbk); let mut decoded = decoder.decode(&encoded[..1]); decoded.extend(decoder.decode(&encoded[1..])); assert_eq!(String::from_utf8(decoded).unwrap(), "中文"); } #[test] fn file_encoding_preserves_matching_bom() { let original = [0xFF, 0xFE, b'A', 0x00]; let (text, had_errors, has_bom) = TextEncoding::Utf16Le.decode_file(&original); let (encoded, encode_errors) = TextEncoding::Utf16Le.encode_file(&text, has_bom); assert!(!had_errors); assert!(!encode_errors); assert_eq!(encoded, original); } }