ledger-prototype/utxo/src/id.rs

use serde::{Deserialize, Deserializer, Serialize};
use sha2::{Digest, Sha256};
use std::fmt::Display;
use std::str::FromStr;

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Invalid length for {0}: {1} (expected: {2})")]
    InvalidLength(String, usize, usize),

    #[error("Unknown prefix {0}")]
    UnknownPrefix(String),
}

macro_rules! Id {
    ($id:ident, $suffix:expr) => {
        #[derive(Clone, Debug, Eq, PartialOrd, Ord, Hash, PartialEq)]
        pub struct $id {
            bytes: [u8; 32],
        }

        impl $id {
            pub fn new(bytes: [u8; 32]) -> Self {
                Self { bytes }
            }
        }

        impl FromStr for $id {
            type Err = Error;
            fn from_str(value: &str) -> Result<Self, Self::Err> {
                Ok(Self::try_from(value).unwrap_or_else(|_| {
                    let mut hasher = Sha256::new();
                    hasher.update(&value);
                    Self {
                        bytes: hasher.finalize().into(),
                    }
                }))
            }
        }

        impl Serialize for $id {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: serde::Serializer,
            {
                serializer.collect_str(&self)
            }
        }

        impl<'de> Deserialize<'de> for $id {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: Deserializer<'de>,
            {
                let s = String::deserialize(deserializer)?;
                // Use FromStr to parse the string and construct the struct
                $id::from_str(&s).map_err(serde::de::Error::custom)
            }
        }

        impl TryFrom<&str> for $id {
            type Error = Error;
            fn try_from(value: &str) -> Result<Self, Self::Error> {
                if $suffix.len() + 64 != value.len() {
                    return Err(Error::InvalidLength(
                        stringify!($id).to_owned(),
                        value.len(),
                        $suffix.len() + 64,
                    ));
                }

                if !value.starts_with($suffix) {
                    return Err(Error::InvalidLength(
                        stringify!($id).to_owned(),
                        value.len(),
                        $suffix.len() + 64,
                    ));
                }

                let bytes = hex::decode(&value[$suffix.len()..]).map_err(|_| {
                    Error::InvalidLength(stringify!($id).to_owned(), value.len(), 32)
                })?;
                bytes.try_into()
            }
        }

        impl TryFrom<&[u8]> for $id {
            type Error = Error;

            fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
                if value.len() != 32 {
                    return Err(Error::InvalidLength(
                        stringify!($id).to_owned(),
                        value.len(),
                        32,
                    ));
                }
                let mut bytes = [0u8; 32];
                bytes.copy_from_slice(&value);
                Ok(Self { bytes })
            }
        }

        impl TryFrom<Vec<u8>> for $id {
            type Error = Error;

            fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
                if value.len() != 32 {
                    return Err(Error::InvalidLength(
                        stringify!($id).to_owned(),
                        value.len(),
                        32,
                    ));
                }
                let mut bytes = [0u8; 32];
                bytes.copy_from_slice(&value);
                Ok(Self { bytes })
            }
        }

        impl Display for $id {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "{}{}", $suffix, hex::encode(self.bytes))
            }
        }

        impl AsRef<[u8]> for $id {
            fn as_ref(&self) -> &[u8] {
                &self.bytes
            }
        }

        impl AsRef<[u8; 32]> for $id {
            fn as_ref(&self) -> &[u8; 32] {
                &self.bytes
            }
        }

        impl std::ops::Deref for $id {
            type Target = [u8; 32];
            fn deref(&self) -> &Self::Target {
                &self.bytes
            }
        }
    };
}

Id!(AccountId, "account");
Id!(TransactionId, "tx");

#[derive(Debug)]
pub enum AnyId {
    Account(AccountId),
    Transaction(TransactionId),
}

impl FromStr for AnyId {
    type Err = Error;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        if value.starts_with("account") {
            Ok(Self::Account(value.parse()?))
        } else if value.starts_with("tx") {
            Ok(Self::Transaction(value.parse()?))
        } else {
            Err(Error::UnknownPrefix(value.to_owned()))
        }
    }
}

impl<'de> Deserialize<'de> for AnyId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        AnyId::from_str(&s).map_err(serde::de::Error::custom)
    }
}

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

    #[test]
    fn from_random_data() {
        let id = "something".parse::<AccountId>().expect("hashed value");
        let id_str = id.to_string();
        let id_bin: &[u8] = id.as_ref();
        assert_eq!(71, id_str.len());
        assert_eq!(
            <&str as TryInto<AccountId>>::try_into(id_str.as_str()).expect("valid"),
            id
        );
        assert_eq!(
            <Vec<u8> as TryInto<AccountId>>::try_into(id_bin.to_owned()).expect("valid"),
            id
        );
        assert_eq!(
            <&[u8] as TryInto<AccountId>>::try_into(id_bin).expect("valid"),
            id
        );
    }
}