smol_base_x/
base.rs

1use super::DecodeError;
2use crate::util::*;
3
4/// ## Base-x for Ascii alphabets (which is most)
5pub trait Base<const BASE: usize> {
6    const ALPHABET: [u8; BASE];
7
8    const LUT: [i8; 256] = gen_lut(&Self::ALPHABET);
9
10    const BASE: usize = Self::ALPHABET.len();
11
12    /// decode input base encoding into buffer
13    /// please do not pass in a non-empty buffer (this will output garbled data).
14    ///
15    /// ```rust
16    /// use smol_base_x::*;
17    ///
18    /// let mut buf = [0u8; 16];
19    /// let bytes_written = Base58Btc::decode_mut("ZiCa", &mut buf).unwrap();
20    ///
21    /// let expected = b"abc";
22    /// assert_eq!(&buf[..bytes_written], expected.as_slice());
23    /// ```
24    fn decode_mut<I: AsRef<[u8]>>(input: I, buf: &mut [u8]) -> Result<usize, DecodeError> {
25        let input = input.as_ref();
26
27        if !input.is_ascii() {
28            return Err(DecodeError::InvalidChar);
29        }
30        // thanks to https://sts10.github.io/2020/10/06/peeking-the-pivot.html for the great notes on iterators with look ahead
31        let mut iter = input.iter().peekable();
32
33        // Skip leading spaces.
34        while let Some(&&ch) = iter.peek() {
35            if ch == b' ' {
36                iter.next();
37            } else {
38                break;
39            }
40        }
41
42        // skip & count leading '1's.
43        let mut ones = 0;
44        while let Some(&&ch) = iter.peek() {
45            if ch == b'1' {
46                ones += 1;
47                iter.next();
48            } else {
49                break;
50            }
51        }
52
53        // in C++ this is used to allocate a vec, but this will overestimate if there are trailing zeroes in the input
54        // let size = decoded_size(Self::BASE, iter.len());
55        // if size + ones > buf.len() {
56        //     return Err(DecodeError::InvalidLength(size + ones));
57        // }
58
59        let mut length = 0;
60
61        // Process the characters.
62        while let Some(&&ch) = iter.peek() {
63            // move forward only if next char is not a space
64            if ch == b' ' {
65                break;
66            }
67            iter.next();
68
69            // Decode base-x character
70            let mut carry = match Self::lookup_ascii(ch) {
71                Some(x) => x,
72                None => return Err(DecodeError::InvalidChar),
73            };
74
75            let mut rev = buf.iter_mut().rev();
76            let mut i = 0;
77
78            while i < length || carry != 0 {
79                match rev.next() {
80                    Some(it) => {
81                        carry += BASE * (*it as usize);
82                        *it = (carry % 256) as u8;
83                        carry /= 256;
84                        i += 1;
85                    }
86                    None => break,
87                }
88            }
89
90            length = i;
91            if length + ones > buf.len() {
92                return Err(DecodeError::InvalidLength(length));
93            }
94
95            // generally we dont want to panic in release if we can avoid it, consider using debug_assert
96            assert!(carry == 0);
97        }
98        length += ones;
99
100        // Skip trailing spaces.
101        while let Some(&&ch) = iter.peek() {
102            if ch == b' ' {
103                iter.next();
104            } else {
105                return Err(DecodeError::CharAfterTrailingSpaces);
106            }
107        }
108
109        buf.rotate_left(buf.len() - length);
110        Ok(length)
111    }
112
113    /// output buff is intentionally a slice since `&mut str` is essentially useless
114    /// users will have to convert output bytes into a str.
115    /// it is assumed buf is zeroed, passing one in that isn't will produce garbage (and extremely likely not valid UTF-8) data.
116    ///
117    /// ```rust
118    /// use smol_base_x::*;
119    ///
120    /// let mut buf = [0u8; 16];
121    /// let bytes_written = Base58Btc::encode_mut("abc", &mut buf).unwrap();
122    ///
123    /// // Here
124    /// let output = core::str::from_utf8(&buf[..bytes_written]).unwrap();
125    ///
126    /// let expected = "ZiCa";
127    /// assert_eq!(output, expected);
128    /// ```
129    fn encode_mut<I: AsRef<[u8]>>(input: I, buf: &mut [u8]) -> Result<usize, DecodeError> {
130        let input = input.as_ref();
131
132        // thanks to https://sts10.github.io/2020/10/06/peeking-the-pivot.html for the great notes on iterators with look ahead
133        let mut iter = input.iter().peekable();
134
135        let mut zeroes = 0;
136        // skip & count leading zeros
137        while let Some(&&ch) = iter.peek() {
138            if ch == 0 {
139                zeroes += 1;
140                iter.next();
141            } else {
142                break;
143            }
144        }
145
146        let size = encoded_size(BASE, iter.len());
147
148        // buf is too small to fit string
149        if size > buf.len() {
150            return Err(DecodeError::InvalidLength(size));
151        }
152
153        let mut length = 0;
154
155        for &ch in iter {
156            let mut carry = ch as usize;
157
158            let mut i = 0;
159            let mut rev = buf.iter_mut().rev();
160            while i < length || carry != 0 {
161                match rev.next() {
162                    Some(it) => {
163                        carry += 256 * (*it as usize);
164                        *it = (carry % BASE) as u8;
165                        carry /= BASE;
166                        i += 1;
167                    }
168                    None => break,
169                }
170            }
171
172            length = i;
173            if length > buf.len() {
174                return Err(DecodeError::InvalidLength(length));
175            }
176
177            assert!(carry == 0);
178        }
179
180        length += zeroes;
181
182        buf.rotate_left(buf.len() - length);
183
184        // translate index into alphabet letter
185        for i in buf.iter_mut().take(length) {
186            *i = Self::ALPHABET[*i as usize];
187        }
188
189        Ok(length)
190    }
191
192    /// Lookup the value for the current char index
193    fn lookup_ascii(ch: u8) -> Option<usize> {
194        match Self::LUT[ch as usize] {
195            -1 => None,
196            i => Some(i as usize),
197        }
198    }
199
200    /// output is `(decoded bytes, bytes written)`
201    #[cfg(feature = "unstable")]
202    fn decode_arr<const LEN: usize, I: Into<[u8; LEN]>>(
203        input: I,
204    ) -> Result<([u8; decoded_arr_size(BASE, LEN)], usize), DecodeError> {
205        let input = input.into();
206        let mut arr = [0u8; decoded_arr_size(BASE, LEN)];
207
208        if !input.is_ascii() {
209            return Err(DecodeError::InvalidChar);
210        }
211        // thanks to https://sts10.github.io/2020/10/06/peeking-the-pivot.html for the great notes on iterators with look ahead
212        let mut iter = input.iter().peekable();
213
214        // Skip leading spaces.
215        while let Some(&&ch) = iter.peek() {
216            if ch == b' ' {
217                iter.next();
218            } else {
219                break;
220            }
221        }
222
223        // skip & count leading '1's.
224        let mut ones = 0;
225        while let Some(&&ch) = iter.peek() {
226            if ch == b'1' {
227                ones += 1;
228                iter.next();
229            } else {
230                break;
231            }
232        }
233
234        let mut length = 0;
235
236        // Process the characters.
237        while let Some(&&ch) = iter.peek() {
238            // move forward only if next char is not a space
239            if ch == b' ' {
240                break;
241            }
242            iter.next();
243
244            // Decode base-x character
245            let mut carry = match Self::lookup_ascii(ch) {
246                Some(x) => x,
247                None => return Err(DecodeError::InvalidChar),
248            };
249
250            let mut rev = arr.iter_mut().rev();
251            let mut i = 0;
252
253            while i < length || carry != 0 {
254                match rev.next() {
255                    Some(it) => {
256                        carry += BASE * (*it as usize);
257                        *it = (carry % 256) as u8;
258                        carry /= 256;
259                        i += 1;
260                    }
261                    None => break,
262                }
263            }
264
265            length = i;
266            // size is already known to be correct
267            // if length > arr.len() {
268            //     return Err(DecodeError::InvalidLength(length));
269            // }
270
271            // generally we dont want to panic in release if we can avoid it, consider using debug_assert
272            assert!(carry == 0);
273        }
274
275        length += ones;
276
277        // Skip trailing spaces.
278        while let Some(&&ch) = iter.peek() {
279            if ch == b' ' {
280                iter.next();
281            } else {
282                return Err(DecodeError::CharAfterTrailingSpaces);
283            }
284        }
285
286        let mid = arr.len() - length;
287        arr.rotate_left(mid);
288
289        Ok((arr, length))
290    }
291
292    /// output is `(encoded chars, chars written)`
293    #[cfg(feature = "unstable")]
294    fn encode_arr<const BYTES: usize, I: Into<[u8; BYTES]>>(
295        input: I,
296    ) -> Result<([u8; encoded_arr_size(Self::BASE, BYTES)], usize), DecodeError> {
297        let input = input.into();
298
299        let mut arr = [0u8; encoded_arr_size(Self::BASE, BYTES)];
300
301        // thanks to https://sts10.github.io/2020/10/06/peeking-the-pivot.html for the great notes on iterators with look ahead
302        let mut iter = input.iter().peekable();
303
304        let mut zeroes = 0;
305        // skip & count leading zeros
306        while let Some(&&ch) = iter.peek() {
307            if ch == 0 {
308                zeroes += 1;
309                iter.next();
310            } else {
311                break;
312            }
313        }
314
315        // let size = encoded_size(BASE, iter.len());
316
317        // array is properly sized at compile time
318        // // buf is too small to fit string
319        // if size > buf.len() {
320        //     return Err(DecodeError::InvalidLength(size));
321        // }
322
323        let mut length = 0;
324
325        for &ch in iter {
326            let mut carry = ch as usize;
327
328            let mut i = 0;
329            let mut rev = arr.iter_mut().rev();
330            while i < length || carry != 0 {
331                match rev.next() {
332                    Some(it) => {
333                        carry += 256 * (*it as usize);
334                        *it = (carry % BASE) as u8;
335                        carry /= BASE;
336                        i += 1;
337                    }
338                    None => break,
339                }
340            }
341
342            length = i;
343            // length known at compile time
344            // if length > arr.len() {
345            //     return Err(DecodeError::InvalidLength(length));
346            // }
347
348            assert!(carry == 0);
349        }
350
351        length += zeroes;
352
353        let mid = arr.len() - length;
354        arr.rotate_left(mid);
355
356        // translate index into alphabet letter
357        for i in arr.iter_mut().take(length) {
358            *i = Self::ALPHABET[*i as usize];
359        }
360
361        Ok((arr, length))
362    }
363}