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
/// Defines a partition based on an explicit indexing of elements to their block
14
/// number.
15
#[derive(Debug)]
16
pub struct IndexedPartition {
17
    partition: Vec<BlockIndex>,
18

            
19
    num_of_blocks: usize,
20
}
21

            
22
impl IndexedPartition {
23
    /// Create a new partition where all elements are in a single block.
24
31911
    pub fn new(num_of_elements: usize) -> IndexedPartition {
25
31911
        IndexedPartition {
26
31911
            partition: vec![BlockIndex::new(0); num_of_elements],
27
31911
            num_of_blocks: 1,
28
31911
        }
29
31911
    }
30

            
31
    /// Create a new partition with the given partitioning.
32
    pub fn with_partition(partition: Vec<BlockIndex>, num_of_blocks: usize) -> IndexedPartition {
33
        IndexedPartition {
34
            partition,
35
            num_of_blocks,
36
        }
37
    }
38

            
39
    /// Iterates over the blocks in the partition.
40
    pub fn iter(&self) -> impl Iterator<Item = BlockIndex> + '_ {
41
        self.partition.iter().copied()
42
    }
43

            
44
    /// Sets the block number of the given element
45
483252
    pub fn set_block(&mut self, element_index: usize, block_number: BlockIndex) {
46
        // TODO: This assumes that the blocks are dense, otherwise it overestimates the number of blocks.
47
483252
        self.num_of_blocks = self.num_of_blocks.max(block_number.value() + 1);
48

            
49
483252
        self.partition[element_index] = block_number;
50
483252
    }
51

            
52
    /// Returns the number of blocks in the partition.
53
8316
    pub fn len(&self) -> usize {
54
8316
        self.partition.len()
55
8316
    }
56

            
57
    /// Returns whether the partition is empty.
58
    pub fn is_empty(&self) -> bool {
59
        self.partition.is_empty()
60
    }
61

            
62
    /// Returns the number of blocks in the partition.
63
228437
    pub fn num_of_blocks(&self) -> usize {
64
228437
        self.num_of_blocks
65
228437
    }
66

            
67
    /// Returns the underlying partition vector.
68
5852836
    pub fn partition(&self) -> &Vec<BlockIndex> {
69
5852836
        &self.partition
70
5852836
    }
71
}
72

            
73
/// Reorders the blocks of the given partition according to the given permutation.
74
pub fn reorder_partition<P>(partition: IndexedPartition, permutation: P) -> IndexedPartition
75
where
76
    P: Fn(BlockIndex) -> BlockIndex,
77
{
78
    let mut new_partition = IndexedPartition::new(partition.partition.len());
79

            
80
    for (element_index, block) in partition.iter().enumerate() {
81
        new_partition.set_block(element_index, permutation(block));
82
    }
83

            
84
    new_partition
85
}
86

            
87
impl fmt::Display for IndexedPartition {
88
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
89
        write!(f, "{{ ")?;
90

            
91
        let mut first = true;
92

            
93
        for block_index in 0..self.partition.len() {
94
            // Print all elements with the same block number.
95
            let mut first_block = true;
96
            for (element_index, _) in self.iter().enumerate().filter(|(_, value)| *value == block_index) {
97
                if !first_block {
98
                    write!(f, ", ")?;
99
                } else {
100
                    if !first {
101
                        write!(f, ", ")?;
102
                    }
103

            
104
                    write!(f, "{{")?;
105
                }
106

            
107
                write!(f, "{element_index}")?;
108
                first_block = false;
109
            }
110

            
111
            if !first_block {
112
                write!(f, "}}")?;
113
                first = false;
114
            }
115
        }
116

            
117
        write!(f, " }}")
118
    }
119
}