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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use core::marker::PhantomData;

use necsim_core::{
    cogs::{Habitat, MathsCore},
    lineage::Lineage,
};

use crate::{
    cogs::origin_sampler::{
        pre_sampler::OriginPreSampler, TrustedOriginSampler, UntrustedOriginSampler,
    },
    decomposition::Decomposition,
};

#[allow(clippy::module_name_repetitions)]
#[derive(Debug)]
pub struct DecompositionOriginSampler<
    'd,
    M: MathsCore,
    O: UntrustedOriginSampler<'d, M>,
    D: Decomposition<M, O::Habitat>,
> {
    origin_sampler: O,
    decomposition: &'d D,
    _marker: PhantomData<M>,
}

impl<'d, M: MathsCore, O: UntrustedOriginSampler<'d, M>, D: Decomposition<M, O::Habitat>>
    DecompositionOriginSampler<'d, M, O, D>
{
    #[must_use]
    pub fn new(origin_sampler: O, decomposition: &'d D) -> Self {
        Self {
            origin_sampler,
            decomposition,
            _marker: PhantomData::<M>,
        }
    }
}

#[contract_trait]
impl<'d, M: MathsCore, O: UntrustedOriginSampler<'d, M>, D: Decomposition<M, O::Habitat>>
    UntrustedOriginSampler<'d, M> for DecompositionOriginSampler<'d, M, O, D>
{
    type Habitat = O::Habitat;
    type PreSampler = O::PreSampler;

    fn habitat(&self) -> &'d Self::Habitat {
        self.origin_sampler.habitat()
    }

    fn into_pre_sampler(self) -> OriginPreSampler<M, Self::PreSampler> {
        self.origin_sampler.into_pre_sampler()
    }

    fn full_upper_bound_size_hint(&self) -> u64 {
        #[allow(
            clippy::cast_possible_truncation,
            clippy::cast_sign_loss,
            clippy::cast_precision_loss
        )]
        {
            ((self.origin_sampler.full_upper_bound_size_hint() as f64)
                / f64::from(self.decomposition.get_subdomain().size().get())) as u64
        }
    }
}

unsafe impl<'d, M: MathsCore, O: TrustedOriginSampler<'d, M>, D: Decomposition<M, O::Habitat>>
    TrustedOriginSampler<'d, M> for DecompositionOriginSampler<'d, M, O, D>
{
}

impl<'d, M: MathsCore, O: UntrustedOriginSampler<'d, M>, D: Decomposition<M, O::Habitat>> Iterator
    for DecompositionOriginSampler<'d, M, O, D>
{
    type Item = Lineage;

    fn next(&mut self) -> Option<Self::Item> {
        #[allow(clippy::while_let_on_iterator)]
        while let Some(lineage) = self.origin_sampler.next() {
            // Forward any out-of-habitat or out-of-deme lineages
            //  (but only on the root subdomain -> no duplication)
            if !self
                .origin_sampler
                .habitat()
                .is_indexed_location_habitable(&lineage.indexed_location)
            {
                if self.decomposition.get_subdomain().is_root() {
                    return Some(lineage);
                }

                continue;
            }

            if self.decomposition.map_location_to_subdomain_rank(
                lineage.indexed_location.location(),
                self.origin_sampler.habitat(),
            ) == self.decomposition.get_subdomain().rank()
            {
                return Some(lineage);
            }
        }

        None
    }
}