1
#![forbid(unsafe_code)]
2

            
3
use std::fmt;
4

            
5
use merc_utilities::TagIndex;
6

            
7
/// A zero sized tag for the block.
8
pub struct BlockTag {}
9

            
10
/// The index for blocks.
11
pub type BlockIndex = TagIndex<usize, BlockTag>;
12

            
13
/// Special value used for elements that do not occur in the partition.
14
pub const NOT_IN_PARTITION: usize = usize::MAX;
15

            
16
/// Defines a partition based on an explicit indexing of elements to their block
17
/// number.
18
///
19
/// # Details
20
///
21
/// This partition always stores 0..max_value elements in the partition, but it
22
/// uses a special value to indicate that an element is not in the partition. So
23
/// this is not memory efficient.
24
#[derive(Clone, Debug)]
25
pub struct IndexedPartition {
26
    /// Stores a mapping from element index to block number.
27
    partition: Vec<BlockIndex>,
28

            
29
    /// Keeps track of the number of blocks defined.
30
    num_of_blocks: usize,
31
}
32

            
33
impl IndexedPartition {
34
    /// Create a new partition with 0 to `num_of_elements` in a single block.
35
102000
    pub fn new(num_of_elements: usize) -> IndexedPartition {
36
102000
        IndexedPartition {
37
102000
            partition: vec![BlockIndex::new(0); num_of_elements],
38
102000
            num_of_blocks: 1,
39
102000
        }
40
102000
    }
41

            
42
    /// Create a new partition where the given indices are in a single block.
43
6905
    pub fn with_subset<I>(num_of_elements: usize, included_indices: I) -> IndexedPartition
44
6905
    where
45
6905
        I: IntoIterator<Item = usize>,
46
    {
47
6905
        let mut partition = vec![BlockIndex::new(NOT_IN_PARTITION); num_of_elements];
48
6905
        let mut has_included_indices = false;
49

            
50
2029337
        for index in included_indices {
51
2029337
            partition[index] = BlockIndex::new(0);
52
2029337
            has_included_indices = true;
53
2029337
        }
54

            
55
6905
        IndexedPartition {
56
6905
            partition,
57
6905
            num_of_blocks: usize::from(has_included_indices),
58
6905
        }
59
6905
    }
60

            
61
    /// Create a new partition with the given partitioning.
62
    pub fn with_partition(partition: Vec<BlockIndex>, num_of_blocks: usize) -> IndexedPartition {
63
        debug_assert!(
64
            partition
65
                .iter()
66
                .all(|&block| block.value() < num_of_blocks || block.value() == NOT_IN_PARTITION),
67
            "Block numbers must be less than the number of blocks, or equal to NOT_IN_PARTITION"
68
        );
69

            
70
        IndexedPartition {
71
            partition,
72
            num_of_blocks,
73
        }
74
    }
75

            
76
    /// Iterates over the blocks in the partition, blocks can be repeated.
77
20500
    pub fn iter(&self) -> impl Iterator<Item = BlockIndex> + '_ {
78
20500
        self.iter_elements().map(|(_, block)| block)
79
20500
    }
80

            
81
    /// Iterates over element-block pairs in the partition.
82
61500
    pub fn iter_elements(&self) -> impl Iterator<Item = (usize, BlockIndex)> + '_ {
83
61500
        self.partition
84
61500
            .iter()
85
61500
            .enumerate()
86
6150000
            .filter_map(|(element_index, &block)| (block.value() != NOT_IN_PARTITION).then_some((element_index, block)))
87
61500
    }
88

            
89
    /// Sets the block number of the given element
90
    ///
91
    /// # Details
92
    ///
93
    /// This assumes that the blocks numbers are dense, otherwise the partition
94
    /// overestimates the total number of blocks present returned from [Self::num_of_blocks].
95
151389863
    pub fn set_block(&mut self, element_index: usize, block_number: BlockIndex) {
96
151389863
        debug_assert!(
97
151389863
            block_number.value() != NOT_IN_PARTITION,
98
            "Block number cannot be NOT_IN_PARTITION"
99
        );
100

            
101
151389863
        self.num_of_blocks = self.num_of_blocks.max(block_number.value() + 1);
102
151389863
        self.partition[element_index] = block_number;
103
151389863
    }
104

            
105
    /// Returns the block number of the given element.
106
19355682045
    pub fn block(&self, element_index: usize) -> BlockIndex {
107
19355682045
        self.partition[element_index]
108
19355682045
    }
109

            
110
    /// Returns the number of elements in the partition.
111
1212814
    pub fn len(&self) -> usize {
112
1212814
        self.partition.len()
113
1212814
    }
114

            
115
    /// Returns whether the partition is empty.
116
    pub fn is_empty(&self) -> bool {
117
        self.partition.is_empty()
118
    }
119

            
120
    /// Returns the number of blocks in the partition.
121
31381675
    pub fn num_of_blocks(&self) -> usize {
122
31381675
        self.num_of_blocks
123
31381675
    }
124
}
125

            
126
/// Reorders the blocks of the given partition according to the given permutation.
127
pub fn reorder_partition<P>(partition: IndexedPartition, permutation: P) -> IndexedPartition
128
where
129
    P: Fn(BlockIndex) -> BlockIndex,
130
{
131
    let mut new_partition = partition.clone();
132

            
133
    for (element_index, block) in partition.iter_elements() {
134
        new_partition.set_block(element_index, permutation(block));
135
    }
136

            
137
    new_partition
138
}
139

            
140
impl fmt::Display for IndexedPartition {
141
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
142
        write!(f, "{{ ")?;
143

            
144
        let mut first = true;
145

            
146
        for block_index in self.iter() {
147
            // Print all elements with the same block number.
148
            let mut first_block = true;
149
            for (element_index, _) in self.iter_elements().filter(|&(_, value)| value == block_index) {
150
                if !first_block {
151
                    write!(f, ", ")?;
152
                } else {
153
                    if !first {
154
                        write!(f, ", ")?;
155
                    }
156

            
157
                    write!(f, "{{")?;
158
                }
159

            
160
                write!(f, "{element_index}")?;
161
                first_block = false;
162
            }
163

            
164
            if !first_block {
165
                write!(f, "}}")?;
166
                first = false;
167
            }
168
        }
169

            
170
        write!(f, " }}")
171
    }
172
}