base58.rs

use crate::hash;
use std::convert::TryFrom;

const ALPHABET: [char; 58] = [
    '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K',
    'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e',
    'f', 'g', 'h', 'i', 'j', 'k', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y',
    'z',
];

#[derive(Debug, PartialEq)]
pub struct Base58String(String);

#[derive(Debug, PartialEq)]
pub struct Base58CheckString(String);

#[derive(Debug)]
pub enum Base58Error {
    InvalidBase58Character { position: usize, character: char },
    InvalidChecksum,
}

impl Base58String {
    pub fn encode(data: Vec<u8>) -> Base58String {
        // log 256 ÷ log 58 ≈ 138 ÷ 100
        let mut buf = vec![0; data.len() * 138 / 100 + 1];
        let mut end = 0;

        for byte in data.iter().skip_while(|x| **x == 0) {
            let mut carry = *byte as u32;
            let mut cursor = 0;

            while carry != 0 || cursor < end {
                carry += 256 * buf[cursor] as u32;
                buf[cursor] = (carry % 58) as u8;
                carry /= 58;
                cursor += 1;
            }

            end = cursor;
        }

        buf.truncate(end);

        let head = data.iter().take_while(|x| **x == 0).map(|_| '1');
        let tail = buf.iter().rev().map(|x| ALPHABET[*x as usize]);

        Base58String(head.chain(tail).collect())
    }

    pub fn decode(self) -> Vec<u8> {
        let data = self.0;

        let decoded_tail = data
            .chars()
            .skip_while(|x| *x == '1')
            .map(|x| {
                ALPHABET
                    .iter()
                    .position(|y| x == *y)
                    .map(|x| x as u8)
                    .unwrap()
            })
            .collect::<Vec<u8>>();

        // log 58 ÷ log 256 ≈ 733 ÷ 1000
        let mut buf = vec![0u8; decoded_tail.len() * 733 / 1000 + 1];
        let mut end = 0;

        for base85_byte in decoded_tail.iter() {
            let mut carry = *base85_byte as u32;
            let mut cursor = 0;

            while carry != 0 || cursor < end {
                carry += 58 * buf[cursor] as u32;
                buf[cursor] = (carry % 256) as u8;
                carry /= 256;
                cursor += 1;
            }

            end = cursor;
        }

        buf.truncate(end);

        let head = data.chars().take_while(|x| *x == '1').map(|_| 0);
        let tail = buf.into_iter().rev();

        head.chain(tail).collect()
    }
}

impl Base58CheckString {
    pub fn encode(mut data: Vec<u8>) -> Base58CheckString {
        let mut checksum = hash::double_sha256(&data);
        checksum.truncate(4);
        data.append(&mut checksum);
        Base58CheckString(Base58String::encode(data).into())
    }

    pub fn decode(self) -> Vec<u8> {
        let mut data = Base58String(self.0).decode();
        data.truncate(data.len() - 4);
        data
    }
}

impl From<Base58CheckString> for Base58String {
    fn from(string: Base58CheckString) -> Base58String {
        Base58String::encode(string.decode())
    }
}

impl From<Base58String> for Base58CheckString {
    fn from(string: Base58String) -> Base58CheckString {
        Base58CheckString::encode(string.decode())
    }
}

impl TryFrom<String> for Base58String {
    type Error = Base58Error;

    fn try_from(string: String) -> Result<Base58String, Base58Error> {
        for (position, character) in string.chars().enumerate() {
            if !ALPHABET.contains(&character) {
                return Err(Base58Error::InvalidBase58Character {
                    position,
                    character,
                });
            }
        }

        Ok(Base58String(string))
    }
}

impl TryFrom<String> for Base58CheckString {
    type Error = Base58Error;

    fn try_from(string: String) -> Result<Base58CheckString, Base58Error> {
        let mut data = Base58String::try_from(string.clone())?.decode();
        let given_checksum = data.drain(data.len() - 4..).collect::<Vec<u8>>();
        let computed_checksum = &hash::double_sha256(&data)[..4];

        if given_checksum == computed_checksum {
            Ok(Base58CheckString(string))
        } else {
            Err(Base58Error::InvalidChecksum)
        }
    }
}

impl Into<String> for Base58String {
    fn into(self) -> String {
        self.0
    }
}

impl Into<String> for Base58CheckString {
    fn into(self) -> String {
        self.0
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_base58string_encode() {
        let table = vec![
            (vec![], ""),
            (vec![45], "n"),
            (vec![48], "q"),
            (vec![49], "r"),
            (vec![57], "z"),
            (vec![45, 49], "4SU"),
            (vec![49, 49], "4k8"),
            (b"abc".to_vec(), "ZiCa"),
            (b"1234598760".to_vec(), "3mJr7AoUXx2Wqd"),
            (
                b"abcdefghijklmnopqrstuvwxyz".to_vec(),
                "3yxU3u1igY8WkgtjK92fbJQCd4BZiiT1v25f",
            ),
        ];

        for (data, base58) in table {
            assert_eq!(Base58String::encode(data), Base58String(base58.into()));
        }
    }

    #[test]
    fn test_base58string_encode_initial_zeroes() {
        let table = vec![
            (b"\0abc".to_vec(), "1ZiCa"),
            (b"\0\0abc".to_vec(), "11ZiCa"),
            (b"\0\0\0abc".to_vec(), "111ZiCa"),
            (b"\0\0\0\0abc".to_vec(), "1111ZiCa"),
        ];

        for (data, base58) in table {
            assert_eq!(Base58String::encode(data), Base58String(base58.into()));
        }
    }

    #[test]
    fn test_base58string_decode() {
        let table = vec![
            ("", vec![]),
            ("Z", vec![32]),
            ("n", vec![45]),
            ("q", vec![48]),
            ("r", vec![49]),
            ("z", vec![57]),
            ("4SU", vec![45, 49]),
            ("4k8", vec![49, 49]),
            ("ZiCa", vec![97, 98, 99]),
            ("3mJr7AoUXx2Wqd", b"1234598760".to_vec()),
            (
                "3yxU3u1igY8WkgtjK92fbJQCd4BZiiT1v25f",
                b"abcdefghijklmnopqrstuvwxyz".to_vec(),
            ),
        ];

        for (base58, data) in table {
            let string = Base58String(base58.into());
            assert_eq!(string.decode(), data);
        }
    }

    #[test]
    fn test_base58string_decode_initial_zeroes() {
        let table = vec![
            ("1ZiCa", b"\0abc".to_vec()),
            ("11ZiCa", b"\0\0abc".to_vec()),
            ("111ZiCa", b"\0\0\0abc".to_vec()),
            ("1111ZiCa", b"\0\0\0\0abc".to_vec()),
        ];

        for (base58, data) in table {
            let string = Base58String(base58.into());
            assert_eq!(string.decode(), data);
        }
    }

    #[test]
    fn test_base58string_try_from_string() {
        let table = vec!["", "Z", "n", "q", "r", "z", "4SU", "4k8", "ZiCa"];

        for base58 in table {
            assert!(Base58String::try_from(String::from(base58)).is_ok());
        }
    }

    #[test]
    fn test_base58string_try_from_string_invalid_char() {
        let table = vec![
            "0", "O", "I", "l", "3mJr0", "O3yxU", "3sNI", "4kl8", "s!5<", "t$@mX<*",
        ];

        for invalid_base58 in table {
            assert!(Base58String::try_from(String::from(invalid_base58)).is_err());
        }
    }

    #[test]
    fn test_base58checkstring_encode() {
        let table = vec![
            (b"abc".to_vec(), "4h3c6RH52R"),
            (b"\0hello".to_vec(), "12L5B5yqsf7vwb"),
        ];

        for (data, base58) in table {
            assert_eq!(
                Base58CheckString::encode(data),
                Base58CheckString(base58.into())
            );
        }
    }

    #[test]
    fn test_base58checkstring_try_from_string() {
        assert!(Base58CheckString::try_from(String::from("12L5B5yqsf7vwb")).is_ok());
    }

    #[test]
    fn test_base58checkstring_try_from_string_invalid_char() {
        assert!(Base58CheckString::try_from(String::from("02L5B5yqsf7vwb")).is_err());
    }

    #[test]
    fn test_base58checkstring_try_from_string_invalid_checksum() {
        assert!(Base58CheckString::try_from(String::from("12L5B5yqsf7vwc")).is_err());
    }

    #[test]
    fn test_base58checkstring_from_base58string() {
        assert_eq!(
            Base58CheckString::from(Base58String("Ldp".into())),
            Base58CheckString("3DUz7ncyT".into())
        );
    }

    #[test]
    fn test_base58string_from_base58checkstring() {
        assert_eq!(
            Base58String::from(Base58CheckString("3DUz7ncyT".into())),
            Base58String("Ldp".into())
        );
    }
}