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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use core::{convert::TryFrom, fmt, num::NonZeroU32};

use serde::{Deserialize, Serialize};

#[derive(Debug)]
#[allow(clippy::module_name_repetitions)]
pub struct PartitionRankOutOfBounds(u32, u32);

impl fmt::Display for PartitionRankOutOfBounds {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(
            fmt,
            "rank {} is is not in partition range [0, {}].",
            self.0, self.1
        )
    }
}

#[allow(clippy::module_name_repetitions, clippy::unsafe_derive_deserialize)]
#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
#[serde(try_from = "PartitionRaw")]
pub struct Partition {
    rank: u32,
    size: PartitionSize,
}

impl Partition {
    /// Creates a `Partition` from a `rank` and number of partitions.
    ///
    /// # Errors
    ///
    /// Returns `PartitionRankOutOfBounds` if `rank >= size`.
    pub const fn try_new(rank: u32, size: PartitionSize) -> Result<Self, PartitionRankOutOfBounds> {
        if rank < size.get() {
            Ok(Self { rank, size })
        } else {
            Err(PartitionRankOutOfBounds(rank, size.get() - 1))
        }
    }

    /// Creates a `Partition` from a `rank` and number of partitions.
    ///
    /// # Safety
    ///
    /// The number of partitions must be strictly greater than `rank`.
    #[must_use]
    pub const unsafe fn new_unchecked(rank: u32, size: PartitionSize) -> Self {
        Self { rank, size }
    }

    #[must_use]
    pub const fn root(size: PartitionSize) -> Self {
        Self { rank: 0, size }
    }

    #[must_use]
    pub const fn monolithic() -> Self {
        Self::root(PartitionSize::MONOLITHIC)
    }

    #[must_use]
    pub const fn rank(self) -> u32 {
        self.rank
    }

    #[must_use]
    pub const fn size(self) -> PartitionSize {
        self.size
    }

    #[must_use]
    pub const fn is_root(self) -> bool {
        self.rank == 0_u32
    }
}

impl TryFrom<PartitionRaw> for Partition {
    type Error = PartitionRankOutOfBounds;

    fn try_from(raw: PartitionRaw) -> Result<Self, Self::Error> {
        Self::try_new(raw.rank, raw.size)
    }
}

#[derive(Copy, Clone, Debug, Deserialize)]
#[serde(deny_unknown_fields)]
#[serde(rename = "Partition")]
struct PartitionRaw {
    rank: u32,
    size: PartitionSize,
}

#[allow(clippy::module_name_repetitions)]
#[allow(clippy::unsafe_derive_deserialize)]
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct PartitionSize(pub NonZeroU32);

impl PartitionSize {
    pub const MONOLITHIC: Self = Self(NonZeroU32::MIN);

    #[must_use]
    pub const fn get(self) -> u32 {
        self.0.get()
    }

    #[must_use]
    pub const fn is_monolithic(self) -> bool {
        self.0.get() == 1
    }

    #[must_use]
    pub fn partitions(self) -> impl ExactSizeIterator<Item = Partition> {
        // Safety: rank is in bounds
        (0..self.get()).map(move |rank| unsafe { Partition::new_unchecked(rank, self) })
    }
}

impl fmt::Display for PartitionSize {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_fmt(format_args!("{}", self.0))
    }
}