1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

use alloc::vec::Vec;

use necsim_core::lineage::{GlobalLineageReference, Lineage};

#[allow(clippy::module_name_repetitions)]
#[derive(Debug)]
pub enum ExceptionalLineage {
    Coalescence {
        child: Lineage,
        parent: GlobalLineageReference,
    },
    OutOfDeme(Lineage),
    OutOfHabitat(Lineage),
}

pub struct SplitExceptionalLineages {
    pub coalescence: Vec<(Lineage, GlobalLineageReference)>,
    pub out_of_deme: Vec<Lineage>,
    pub out_of_habitat: Vec<Lineage>,
}

impl ExceptionalLineage {
    #[must_use]
    pub fn split_vec(exceptional_lineages: Vec<ExceptionalLineage>) -> SplitExceptionalLineages {
        let mut coalescence_lineages = Vec::new();
        let mut out_of_deme_lineages = Vec::new();
        let mut out_of_habitat_lineages = Vec::new();

        for lineage in exceptional_lineages {
            match lineage {
                ExceptionalLineage::Coalescence { child, parent } => {
                    coalescence_lineages.push((child, parent));
                },
                ExceptionalLineage::OutOfDeme(lineage) => out_of_deme_lineages.push(lineage),
                ExceptionalLineage::OutOfHabitat(lineage) => out_of_habitat_lineages.push(lineage),
            }
        }

        SplitExceptionalLineages {
            coalescence: coalescence_lineages,
            out_of_deme: out_of_deme_lineages,
            out_of_habitat: out_of_habitat_lineages,
        }
    }

    pub fn drain_coalescing_lineages(
        exceptional_lineages: &mut Vec<ExceptionalLineage>,
    ) -> impl Iterator<Item = Lineage> + '_ {
        exceptional_lineages
            .extract_if(|exceptional_lineage| {
                matches!(exceptional_lineage, ExceptionalLineage::Coalescence { .. })
            })
            .map(|exceptional_lineage| match exceptional_lineage {
                ExceptionalLineage::Coalescence { child: lineage, .. }
                | ExceptionalLineage::OutOfDeme(lineage)
                | ExceptionalLineage::OutOfHabitat(lineage) => lineage,
            })
    }
}