kaspa-mining 2.1.0

use itertools::Itertools;
use kaspa_consensus_core::tx::Transaction;
use std::{
    collections::{HashMap, HashSet, VecDeque},
    iter::{FusedIterator, Map},
};

type IndexSet = HashSet<usize>;

pub trait IterTopologically<T>
where
    T: AsRef<Transaction>,
{
    fn topological_iter(&self) -> TopologicalIter<'_, T>;
}

impl<T: AsRef<Transaction>> IterTopologically<T> for &[T] {
    fn topological_iter(&self) -> TopologicalIter<'_, T> {
        TopologicalIter::new(self)
    }
}

impl<T: AsRef<Transaction>> IterTopologically<T> for Vec<T> {
    fn topological_iter(&self) -> TopologicalIter<'_, T> {
        TopologicalIter::new(self)
    }
}

pub struct TopologicalIter<'a, T: AsRef<Transaction>> {
    transactions: &'a [T],
    in_degree: Vec<isize>,
    edges: Vec<Option<IndexSet>>,
    queue: VecDeque<usize>,
    yields_count: usize,
}

impl<'a, T: AsRef<Transaction>> TopologicalIter<'a, T> {
    pub fn new(transactions: &'a [T]) -> Self {
        let mut in_degree: Vec<isize> = vec![0; transactions.len()];

        // Index on transaction ids
        let mut index = HashMap::with_capacity(transactions.len());
        transactions.iter().enumerate().for_each(|(idx, tx)| {
            let _ = index.insert(tx.as_ref().id(), idx);
        });

        // Transaction edges
        let mut edges: Vec<Option<IndexSet>> = vec![None; transactions.len()];
        transactions.iter().enumerate().for_each(|(destination_idx, tx)| {
            tx.as_ref().inputs.iter().for_each(|input| {
                if let Some(origin_idx) = index.get(&input.previous_outpoint.transaction_id) {
                    edges[*origin_idx].get_or_insert_with(IndexSet::new).insert(destination_idx);
                }
            })
        });

        // Degrees
        (0..transactions.len()).for_each(|origin_idx| {
            if let Some(ref edges) = edges[origin_idx] {
                edges.iter().for_each(|destination_idx| {
                    #[allow(clippy::arithmetic_side_effects, reason = "ARITH-SAFETY(COUNTER)")]
                    {
                        in_degree[*destination_idx] += 1;
                    }
                });
            }
        });

        // Degree 0
        let mut queue = VecDeque::with_capacity(transactions.len());
        (0..transactions.len()).for_each(|destination_idx| {
            if in_degree[destination_idx] == 0 {
                queue.push_back(destination_idx);
            }
        });
        Self { transactions, in_degree, edges, queue, yields_count: 0 }
    }
}

impl<'a, T: AsRef<Transaction>> Iterator for TopologicalIter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        match self.queue.pop_front() {
            Some(current) => {
                if let Some(ref edges) = self.edges[current] {
                    edges.iter().for_each(|destination_idx| {
                        let degree = self.in_degree.get_mut(*destination_idx).unwrap();
                        #[allow(
                            clippy::arithmetic_side_effects,
                            reason = "Since the edge goes into `destination_idx`, the degree must be at least 1."
                        )]
                        {
                            *degree -= 1;
                        }
                        if *degree == 0 {
                            self.queue.push_back(*destination_idx);
                        }
                    });
                }
                #[allow(clippy::arithmetic_side_effects, reason = "ARITH-SAFETY(COUNTER)")]
                {
                    self.yields_count += 1;
                }
                Some(&self.transactions[current])
            }
            None => None,
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        #[allow(
            clippy::arithmetic_side_effects,
            reason = "the subtraction is safe since we subtract at most `self.transactions.len()`."
        )]
        let items_remaining = self.transactions.len() - self.yields_count.min(self.transactions.len());
        (self.yields_count, Some(items_remaining))
    }
}

impl<T: AsRef<Transaction>> FusedIterator for TopologicalIter<'_, T> {}
impl<T: AsRef<Transaction>> ExactSizeIterator for TopologicalIter<'_, T> {
    fn len(&self) -> usize {
        self.transactions.len()
    }
}

pub trait IntoIterTopologically<T>
where
    T: AsRef<Transaction>,
{
    fn topological_into_iter(self) -> TopologicalIntoIter<T>;
}

impl<T: AsRef<Transaction>> IntoIterTopologically<T> for Vec<T> {
    fn topological_into_iter(self) -> TopologicalIntoIter<T> {
        TopologicalIntoIter::new(self)
    }
}

impl<T, I, F> IntoIterTopologically<T> for Map<I, F>
where
    T: AsRef<Transaction>,
    I: Iterator,
    F: FnMut(<I as Iterator>::Item) -> T,
{
    fn topological_into_iter(self) -> TopologicalIntoIter<T> {
        TopologicalIntoIter::new(self)
    }
}

pub struct TopologicalIntoIter<T: AsRef<Transaction>> {
    transactions: Vec<Option<T>>,
    in_degree: Vec<isize>,
    edges: Vec<Option<IndexSet>>,
    queue: VecDeque<usize>,
    yields_count: usize,
}

impl<T: AsRef<Transaction>> TopologicalIntoIter<T> {
    pub fn new(transactions: impl IntoIterator<Item = T>) -> Self {
        // Collect all transactions
        let transactions = transactions.into_iter().map(|tx| Some(tx)).collect_vec();

        let mut in_degree: Vec<isize> = vec![0; transactions.len()];

        // Index on transaction ids
        let mut index = HashMap::with_capacity(transactions.len());
        transactions.iter().enumerate().for_each(|(idx, tx)| {
            let _ = index.insert(tx.as_ref().unwrap().as_ref().id(), idx);
        });

        // Transaction edges
        let mut edges: Vec<Option<IndexSet>> = vec![None; transactions.len()];
        transactions.iter().enumerate().for_each(|(destination_idx, tx)| {
            tx.as_ref().unwrap().as_ref().inputs.iter().for_each(|input| {
                if let Some(origin_idx) = index.get(&input.previous_outpoint.transaction_id) {
                    edges[*origin_idx].get_or_insert_with(IndexSet::new).insert(destination_idx);
                }
            })
        });

        // Degrees
        (0..transactions.len()).for_each(|origin_idx| {
            if let Some(ref edges) = edges[origin_idx] {
                edges.iter().for_each(|destination_idx| {
                    #[allow(clippy::arithmetic_side_effects, reason = "ARITH-SAFETY(COUNTER)")]
                    {
                        in_degree[*destination_idx] += 1;
                    }
                });
            }
        });

        // Degree 0
        let mut queue = VecDeque::with_capacity(transactions.len());
        (0..transactions.len()).for_each(|destination_idx| {
            if in_degree[destination_idx] == 0 {
                queue.push_back(destination_idx);
            }
        });
        Self { transactions, in_degree, edges, queue, yields_count: 0 }
    }
}

impl<T: AsRef<Transaction>> Iterator for TopologicalIntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        match self.queue.pop_front() {
            Some(current) => {
                if let Some(ref edges) = self.edges[current] {
                    edges.iter().for_each(|destination_idx| {
                        let degree = self.in_degree.get_mut(*destination_idx).unwrap();
                        #[allow(
                            clippy::arithmetic_side_effects,
                            reason = "Since the edge goes into `destination_idx`, the degree must be at least 1."
                        )]
                        {
                            *degree -= 1;
                        }
                        if *degree == 0 {
                            self.queue.push_back(*destination_idx);
                        }
                    });
                }
                #[allow(clippy::arithmetic_side_effects, reason = "ARITH-SAFETY(COUNTER)")]
                {
                    self.yields_count += 1;
                }

                self.transactions[current].take()
            }
            None => None,
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        #[allow(
            clippy::arithmetic_side_effects,
            reason = "the subtractin is safe since we subtract at most `self.transactions.len()`."
        )]
        let items_remaining = self.transactions.len() - self.yields_count.min(self.transactions.len());
        (self.yields_count, Some(items_remaining))
    }
}

impl<T: AsRef<Transaction>> FusedIterator for TopologicalIntoIter<T> {}
impl<T: AsRef<Transaction>> ExactSizeIterator for TopologicalIntoIter<T> {
    fn len(&self) -> usize {
        self.transactions.len()
    }
}