1use core::fmt;
4
5#[cfg(feature = "alloc")]
6use alloc::{string::String, vec::Vec};
7
8use crate::Check;
9#[cfg(any(feature = "check", feature = "cb58"))]
10use crate::CHECKSUM_LEN;
11
12use crate::Alphabet;
13
14#[allow(missing_debug_implementations)]
16pub struct EncodeBuilder<'a, I: AsRef<[u8]>> {
17 input: I,
18 alpha: &'a Alphabet,
19 check: Check,
20}
21
22pub type Result<T> = core::result::Result<T, Error>;
24
25#[derive(Copy, Clone, Debug, Eq, PartialEq)]
27#[non_exhaustive]
28pub enum Error {
29 BufferTooSmall,
31}
32
33pub trait EncodeTarget {
36 fn encode_with(
40 &mut self,
41 max_len: usize,
42 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
43 ) -> Result<usize>;
44}
45
46impl<T: EncodeTarget + ?Sized> EncodeTarget for &mut T {
47 fn encode_with(
48 &mut self,
49 max_len: usize,
50 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
51 ) -> Result<usize> {
52 T::encode_with(self, max_len, f)
53 }
54}
55
56#[cfg(feature = "alloc")]
57impl EncodeTarget for Vec<u8> {
58 fn encode_with(
59 &mut self,
60 max_len: usize,
61 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
62 ) -> Result<usize> {
63 let original = self.len();
64 self.resize(original + max_len, 0);
65 let len = f(&mut self[original..])?;
66 self.truncate(original + len);
67 Ok(len)
68 }
69}
70
71#[cfg(feature = "smallvec")]
72impl<A: smallvec::Array<Item = u8>> EncodeTarget for smallvec::SmallVec<A> {
73 fn encode_with(
81 &mut self,
82 max_len: usize,
83 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
84 ) -> Result<usize> {
85 let original = self.len();
86 self.resize(original + max_len, 0);
87 let len = f(&mut self[original..])?;
88 self.truncate(original + len);
89 Ok(len)
90 }
91}
92
93#[cfg(feature = "tinyvec")]
94impl<A: tinyvec::Array<Item = u8>> EncodeTarget for tinyvec::ArrayVec<A> {
95 fn encode_with(
96 &mut self,
97 max_len: usize,
98 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
99 ) -> Result<usize> {
100 let _ = max_len;
101 let original = self.len();
102 let len = f(self.grab_spare_slice_mut())?;
103 self.set_len(original + len);
104 Ok(len)
105 }
106}
107
108#[cfg(feature = "tinyvec")]
109impl EncodeTarget for tinyvec::SliceVec<'_, u8> {
110 fn encode_with(
111 &mut self,
112 max_len: usize,
113 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
114 ) -> Result<usize> {
115 let _ = max_len;
116 let original = self.len();
117 let len = f(self.grab_spare_slice_mut())?;
118 self.set_len(original + len);
119 Ok(len)
120 }
121}
122
123#[cfg(all(feature = "tinyvec", feature = "alloc"))]
124impl<A: tinyvec::Array<Item = u8>> EncodeTarget for tinyvec::TinyVec<A> {
125 fn encode_with(
126 &mut self,
127 max_len: usize,
128 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
129 ) -> Result<usize> {
130 let original = self.len();
131 self.resize(original + max_len, 0);
132 let len = f(&mut self[original..])?;
133 self.truncate(original + len);
134 Ok(len)
135 }
136}
137
138#[cfg(feature = "alloc")]
139impl EncodeTarget for String {
140 fn encode_with(
141 &mut self,
142 max_len: usize,
143 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
144 ) -> Result<usize> {
145 let mut output = core::mem::take(self).into_bytes();
146 let len = output.encode_with(max_len, f)?;
147 *self = String::from_utf8(output).unwrap();
148 Ok(len)
149 }
150}
151
152impl EncodeTarget for [u8] {
153 fn encode_with(
154 &mut self,
155 max_len: usize,
156 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
157 ) -> Result<usize> {
158 let _ = max_len;
159 f(&mut *self)
160 }
161}
162
163impl EncodeTarget for str {
164 fn encode_with(
165 &mut self,
166 max_len: usize,
167 f: impl for<'a> FnOnce(&'a mut [u8]) -> Result<usize>,
168 ) -> Result<usize> {
169 struct Guard<'a>(&'a mut [u8]);
170
171 impl Drop for Guard<'_> {
172 fn drop(&mut self) {
173 let mut index = 0;
174 loop {
175 match core::str::from_utf8(&self.0[index..]) {
176 Ok(_) => return,
177 Err(e) => {
178 index += e.valid_up_to();
179 if let Some(len) = e.error_len() {
180 for i in &mut self.0[index..index + len] {
181 *i = 0;
182 }
183 index += len;
184 } else {
185 for i in &mut self.0[index..] {
186 *i = 0;
187 }
188 index += self.0[index..].len();
189 }
190 }
191 }
192 }
193 }
194 }
195
196 let _ = max_len;
197
198 #[allow(unsafe_code)]
199 let guard = Guard(unsafe { self.as_bytes_mut() });
202 f(&mut *guard.0)
203 }
204}
205
206impl<'a, I: AsRef<[u8]>> EncodeBuilder<'a, I> {
207 pub fn new(input: I, alpha: &'a Alphabet) -> EncodeBuilder<'a, I> {
211 EncodeBuilder {
212 input,
213 alpha,
214 check: Check::Disabled,
215 }
216 }
217
218 pub(crate) fn from_input(input: I) -> EncodeBuilder<'static, I> {
220 EncodeBuilder {
221 input,
222 alpha: Alphabet::DEFAULT,
223 check: Check::Disabled,
224 }
225 }
226
227 pub fn with_alphabet(self, alpha: &'a Alphabet) -> EncodeBuilder<'a, I> {
240 EncodeBuilder { alpha, ..self }
241 }
242
243 #[cfg(feature = "check")]
259 pub fn with_check(self) -> EncodeBuilder<'a, I> {
260 let check = Check::Enabled(None);
261 EncodeBuilder { check, ..self }
262 }
263
264 #[cfg(feature = "check")]
280 pub fn with_check_version(self, expected_ver: u8) -> EncodeBuilder<'a, I> {
281 let check = Check::Enabled(Some(expected_ver));
282 EncodeBuilder { check, ..self }
283 }
284
285 #[cfg(feature = "cb58")]
301 pub fn as_cb58(self, expected_ver: Option<u8>) -> EncodeBuilder<'a, I> {
302 let check = Check::CB58(expected_ver);
303 EncodeBuilder { check, ..self }
304 }
305
306 #[cfg(feature = "alloc")]
315 pub fn into_string(self) -> String {
316 let mut output = String::new();
317 self.onto(&mut output).unwrap();
318 output
319 }
320
321 #[cfg(feature = "alloc")]
330 pub fn into_vec(self) -> Vec<u8> {
331 let mut output = Vec::new();
332 self.onto(&mut output).unwrap();
333 output
334 }
335
336 pub fn onto(self, mut output: impl EncodeTarget) -> Result<usize> {
402 let input = self.input.as_ref();
403 match self.check {
404 Check::Disabled => output.encode_with(max_encoded_len(input.len()), |output| {
405 encode_into(input, output, self.alpha)
406 }),
407 #[cfg(feature = "check")]
408 Check::Enabled(version) => {
409 let input_len = input.len() + CHECKSUM_LEN + version.map_or(0, |_| 1);
410 output.encode_with(max_encoded_len(input_len), |output| {
411 encode_check_into(self.input.as_ref(), output, self.alpha, version)
412 })
413 }
414 #[cfg(feature = "cb58")]
415 Check::CB58(version) => {
416 let input_len = input.len() + CHECKSUM_LEN + version.map_or(0, |_| 1);
417 output.encode_with(max_encoded_len(input_len), |output| {
418 encode_cb58_into(self.input.as_ref(), output, self.alpha, version)
419 })
420 }
421 }
422 }
423}
424
425fn max_encoded_len(len: usize) -> usize {
430 len + (len + 1) / 2
432}
433
434fn encode_into<'a, I>(input: I, output: &mut [u8], alpha: &Alphabet) -> Result<usize>
435where
436 I: Clone + IntoIterator<Item = &'a u8>,
437{
438 let mut index = 0;
439 for &val in input.clone() {
440 let mut carry = val as usize;
441 for byte in &mut output[..index] {
442 carry += (*byte as usize) << 8;
443 *byte = (carry % 58) as u8;
444 carry /= 58;
445 }
446 while carry > 0 {
447 if index == output.len() {
448 return Err(Error::BufferTooSmall);
449 }
450 output[index] = (carry % 58) as u8;
451 index += 1;
452 carry /= 58;
453 }
454 }
455
456 for _ in input.into_iter().take_while(|v| **v == 0) {
457 if index == output.len() {
458 return Err(Error::BufferTooSmall);
459 }
460 output[index] = 0;
461 index += 1;
462 }
463
464 for val in &mut output[..index] {
465 *val = alpha.encode[*val as usize];
466 }
467
468 output[..index].reverse();
469 Ok(index)
470}
471
472#[cfg(feature = "check")]
473fn encode_check_into(
474 input: &[u8],
475 output: &mut [u8],
476 alpha: &Alphabet,
477 version: Option<u8>,
478) -> Result<usize> {
479 use sha2::{Digest, Sha256};
480
481 let mut first_hash = Sha256::new();
482 if let Some(version) = version {
483 first_hash.update([version; 1]);
484 }
485 let first_hash = first_hash.chain_update(input).finalize();
486 let second_hash = Sha256::digest(first_hash);
487
488 let checksum = &second_hash[0..CHECKSUM_LEN];
489
490 encode_into(
491 version.iter().chain(input.iter()).chain(checksum.iter()),
492 output,
493 alpha,
494 )
495}
496
497#[cfg(feature = "cb58")]
498fn encode_cb58_into(
499 input: &[u8],
500 output: &mut [u8],
501 alpha: &Alphabet,
502 version: Option<u8>,
503) -> Result<usize> {
504 use sha2::{Digest, Sha256};
505
506 let mut hash = Sha256::new();
507 if let Some(version) = version {
508 hash.update([version; 1]);
509 }
510 let hash = hash.chain_update(input).finalize();
511
512 let checksum = &hash[hash.len() - CHECKSUM_LEN..];
513
514 encode_into(
515 version.iter().chain(input.iter()).chain(checksum.iter()),
516 output,
517 alpha,
518 )
519}
520
521#[cfg(feature = "std")]
522impl std::error::Error for Error {}
523
524impl fmt::Display for Error {
525 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
526 match *self {
527 Error::BufferTooSmall => write!(
528 f,
529 "buffer provided to encode base58 string into was too small"
530 ),
531 }
532 }
533}