1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170

use super::Range;
use encoding_rs::{Encoding, WINDOWS_1252};
use memchr::{memchr, memchr3};
use std::borrow::Cow;
use std::fmt::{self, Debug};
use std::ops::Deref;
use std::str;

/// An error used to indicate that an encoded string has replacements and can't be converted losslessly.
#[derive(Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct HasReplacementsError;

/// A thin wrapper around either byte slice or owned bytes with some handy APIs attached
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct Bytes<'b>(Cow<'b, [u8]>);

impl<'b> Bytes<'b> {
    #[inline]
    pub fn from_str(string: &'b str, encoding: &'static Encoding) -> Self {
        encoding.encode(string).0.into()
    }

    #[inline]
    pub fn from_str_without_replacements(
        string: &'b str,
        encoding: &'static Encoding,
    ) -> Result<Self, HasReplacementsError> {
        let (res, _, has_replacements) = encoding.encode(string);

        if has_replacements {
            Err(HasReplacementsError)
        } else {
            Ok(res.into())
        }
    }

    #[inline]
    pub fn as_string(&self, encoding: &'static Encoding) -> String {
        encoding.decode(self).0.into_owned()
    }

    #[inline]
    pub fn as_lowercase_string(&self, encoding: &'static Encoding) -> String {
        encoding.decode(self).0.to_ascii_lowercase()
    }

    #[inline]
    pub fn into_owned(self) -> Bytes<'static> {
        Bytes(Cow::Owned(self.0.into_owned()))
    }

    #[inline]
    pub fn slice(&self, range: Range) -> Bytes {
        self.0[range.start..range.end].into()
    }

    #[inline]
    pub fn opt_slice(&self, range: Option<Range>) -> Option<Bytes> {
        range.map(|range| self.slice(range))
    }

    pub(crate) fn as_debug_string(&self) -> String {
        // NOTE: use WINDOWS_1252 (superset of ASCII) encoding here as
        // the most safe variant since we don't know which actual encoding
        // has been used for bytes.
        self.as_string(WINDOWS_1252)
    }
}

macro_rules! impl_replace_byte {
    ($self:tt, $output_handler:ident, $impls:ident) => {
        let mut tail: &[u8] = $self;

        loop {
            match $impls!(@find tail) {
                Some(pos) => {
                    let replacement = $impls!(@get_replacement tail, pos);
                    let chunk = &tail[..pos];

                    if !chunk.is_empty() {
                        $output_handler(chunk);
                    }

                    $output_handler(&replacement);
                    tail = &tail[pos + 1..];
                }
                None => {
                    if !tail.is_empty() {
                        $output_handler(&tail);
                    }
                    break;
                }
            }
        }
    };
}

impl<'b> Bytes<'b> {
    #[inline]
    pub fn replace_byte(&self, (needle, repl): (u8, &[u8]), output_handler: &mut dyn FnMut(&[u8])) {
        macro_rules! impls {
            (@find $tail:ident) => {
                memchr(needle, $tail)
            };

            (@get_replacement $tail:ident, $pos:ident) => {
                repl
            };
        }

        impl_replace_byte!(self, output_handler, impls);
    }

    #[inline]
    pub fn replace_byte3(
        &self,
        (needle1, repl1): (u8, &[u8]),
        (needle2, repl2): (u8, &[u8]),
        (needle3, repl3): (u8, &[u8]),
        output_handler: &mut dyn FnMut(&[u8]),
    ) {
        macro_rules! impls {
            (@find $tail:ident) => {
                memchr3(needle1, needle2, needle3, $tail)
            };

            (@get_replacement $tail:ident, $pos:ident) => {{
                let matched = $tail[$pos];

                if matched == needle1 {
                    repl1
                } else if matched == needle2 {
                    repl2
                } else {
                    repl3
                }
            }};
        }

        impl_replace_byte!(self, output_handler, impls);
    }
}

impl<'b> From<Cow<'b, [u8]>> for Bytes<'b> {
    #[inline]
    fn from(bytes: Cow<'b, [u8]>) -> Self {
        Bytes(bytes)
    }
}

impl<'b> From<&'b [u8]> for Bytes<'b> {
    #[inline]
    fn from(bytes: &'b [u8]) -> Self {
        Bytes(bytes.into())
    }
}

impl Debug for Bytes<'_> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "`{}`", self.as_debug_string())
    }
}

impl Deref for Bytes<'_> {
    type Target = [u8];

    fn deref(&self) -> &[u8] {
        &*self.0
    }
}