logos-execution-zone/key_protocol/src/key_management/key_tree/chain_index.rs

use std::{fmt::Display, str::FromStr};

use serde::{Deserialize, Serialize};

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize)]
pub struct ChainIndex(Vec<u32>);

#[derive(thiserror::Error, Debug)]
pub enum ChainIndexError {
    #[error("No root found")]
    NoRootFound,
    #[error("Failed to parse segment into a number")]
    ParseIntError(#[from] std::num::ParseIntError),
}

impl FromStr for ChainIndex {
    type Err = ChainIndexError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if !s.starts_with('/') {
            return  Err(ChainIndexError:NoRootFound);
        }
        
        s
            .split("/")
            .map(u32::from_str)
            .collect()
            .map_err(Into::into)
    }
}

impl Display for ChainIndex {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "/")?;
        for cci in &self.0[..(self.0.len() - 1)] {
            write!(f, "{cci}/")?;
        }
        write!(f, "{}", self.0.last().unwrap())
    }
}

impl ChainIndex {
    pub fn root() -> Self {
        ChainIndex::from_str("/").unwrap()
    }

    pub fn chain(&self) -> &[u32] {
        &self.0
    }

    pub fn next_in_line(&self) -> ChainIndex {
        let mut chain = self.0.clone();
        //ToDo: Add overflow check
        if let Some(last_p) = chain.last_mut() {
            *last_p += 1
        }

        ChainIndex(chain)
    }

    pub fn nth_child(&self, child_id: u32) -> ChainIndex {
        let mut chain = self.0.clone();
        chain.push(child_id);

        ChainIndex(chain)
    }
}

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

    #[test]
    fn test_chain_id_root_correct() {
        let chain_id = ChainIndex::root();
        let chain_id_2 = ChainIndex::from_str("/").unwrap();

        assert_eq!(chain_id, chain_id_2);
    }

    #[test]
    fn test_chain_id_deser_correct() {
        let chain_id = ChainIndex::from_str("/257").unwrap();

        assert_eq!(chain_id.chain(), &[257]);
    }

    #[test]
    fn test_chain_id_next_in_line_correct() {
        let chain_id = ChainIndex::from_str("/257").unwrap();
        let next_in_line = chain_id.next_in_line();

        assert_eq!(next_in_line, ChainIndex::from_str("/258").unwrap());
    }

    #[test]
    fn test_chain_id_child_correct() {
        let chain_id = ChainIndex::from_str("/257").unwrap();
        let child = chain_id.n_th_child(3);

        assert_eq!(child, ChainIndex::from_str("/257/3").unwrap());
    }

    #[test]
    fn test_correct_display() {
        let chainid = ChainIndex(vec![5, 7, 8]);

        let string_index = format!("{chainid}");

        assert_eq!(string_index, "/5/7/8".to_string());
    }
}