tangle-test.rs

use std::{
    collections::{HashMap, VecDeque, hash_map::DefaultHasher},
    hash::{Hash, Hasher, BuildHasherDefault, BuildHasher},
};

#[derive(Debug)]
pub enum Error {
    NoSuchHash,
    DuplicateHash,
}

// TODO: Permit carrying something interesting in the transaction body
#[derive(Hash)]
pub struct TxBody(u32);

#[derive(Hash)]
pub struct Tx {
    parents: Option<[TxHash; 2]>,
    confirmed: bool,
    body: TxBody,
}

impl Tx {
    pub fn parents(&self) -> Option<[TxHash; 2]> {
        self.parents
    }

    pub fn body(&self) -> &TxBody {
        &self.body
    }

    pub fn confirmed(&self) -> bool {
        self.confirmed
    }

    pub fn compute_hash(&self) -> TxHash {
        // TODO: Replace with actual hash implementation later
        let mut hasher: DefaultHasher = BuildHasherDefault::default()
            .build_hasher();
        self.hash(&mut hasher);
        TxHash(hasher.finish())
    }
}

#[derive(Copy, Clone, Hash, PartialEq, Eq)]
pub struct TxHash(u64);

pub struct Tangle {
    txs: HashMap<TxHash, Tx>,
    txs_ages: VecDeque<TxHash>,
    max_txs: usize,
}

impl Tangle {
    fn maintain(&mut self) {
        while self.txs_ages.len() > self.max_txs {
            self.txs.remove(&self.txs_ages.pop_front().unwrap());
        }
    }

    fn insert_raw(&mut self, tx: Tx) -> Result<TxHash, Error> {
        let hash = tx.compute_hash();
        if !self.contains(hash) {
            self.txs.insert(hash, tx);
            self.txs_ages.push_back(hash);
            self.maintain();
            Ok(hash)
        } else {
            Err(Error::DuplicateHash)
        }
    }

    fn derive_confirmation(txs: &mut HashMap<TxHash, Tx>, hash: TxHash) -> bool {
        match txs.get(&hash) {
            Some(tx) if tx.confirmed() => true,
            None => false,
            Some(tx) => if let Some(parents) = tx.parents() {
                if parents
                    .iter()
                    .all(|parent| Self::derive_confirmation(txs, *parent))
                {
                    txs
                        .get_mut(&hash)
                        .unwrap()
                        .confirmed = true;
                    true
                } else {
                    false
                }
            } else {
                false
            },
        }
    }

    pub fn contains(&self, hash: TxHash) -> bool {
        self.txs.contains_key(&hash)
    }

    pub fn get(&self, hash: TxHash) -> Option<&Tx> {
        self.txs.get(&hash)
    }

    pub fn confirm(&mut self, hash: TxHash) {
        self.txs.get_mut(&hash).map(|tx| tx.confirmed = true);

        for hash in self.txs_ages.iter().rev() {
            if !self
                .get(*hash)
                .map(|tx| tx.confirmed())
                .unwrap_or(true)
            {
                Self::derive_confirmation(&mut self.txs, *hash);
            }
        }
    }

    pub fn insert_root(&mut self, body: TxBody) -> Result<TxHash, Error> {
        self.insert_raw(Tx {
            parents: None,
            confirmed: false,
            body,
        })
    }

    pub fn insert(&mut self, body: TxBody, parents: [TxHash; 2]) -> Result<TxHash, Error> {
        if parents.iter().all(|hash| self.contains(*hash)) {
            self.insert_raw(Tx {
                parents: Some(parents),
                confirmed: false,
                body,
            })
        } else {
            Err(Error::NoSuchHash)
        }
    }

    pub fn visit_parents(&self, hash: TxHash, mut f: impl FnMut(TxHash, &Tx)) {
        if let Some(tx) = self.get(hash) {
            f(hash, tx);
            tx
                .parents()
                .map(|parents| parents
                    .iter()
                    .for_each(|parent| self.visit_parents(*parent, &mut f)));
        }
    }
}

impl Default for Tangle {
    fn default() -> Self {
        Self {
            txs: HashMap::default(),
            txs_ages: VecDeque::default(),
            max_txs: 4096,
        }
    }
}

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

    #[test]
    fn simple() {
        let mut t = Tangle::default();

        let root0 = t.insert_root(TxBody(0)).unwrap();
        let root1 = t.insert_root(TxBody(1)).unwrap();

        let tx0 = t.insert(TxBody(0), [root0, root1]).unwrap();
        let tx1 = t.insert(TxBody(0), [tx0, root0]).unwrap();
        let tx2 = t.insert(TxBody(0), [tx0, tx1]).unwrap();
        let tx3 = t.insert(TxBody(0), [tx2, tx0]).unwrap();

        assert_eq!(t.get(tx3).unwrap().confirmed(), false);

        t.confirm(root0);

        assert_eq!(t.get(tx3).unwrap().confirmed(), false);

        t.confirm(root1);

        assert_eq!(t.get(tx3).unwrap().confirmed(), true);
    }
}