1
#![forbid(unsafe_code)]
2

            
3
use std::fmt;
4

            
5
use itertools::Itertools;
6
use log::trace;
7

            
8
use crate::BlockIndex;
9
use crate::IndexedPartition;
10

            
11
/// A partition that explicitly stores a list of blocks and their indexing into
12
/// the list of elements.
13
#[derive(Debug)]
14
pub struct BlockPartition<A: Clone + fmt::Debug = ()> {
15
    elements: Vec<usize>,
16
    blocks: Vec<Block<A>>,
17
}
18

            
19
impl<A: Clone + fmt::Debug + Default> BlockPartition<A> {
20
    /// Create an initial partition where all the states are in a single block
21
    /// 0. And all the elements in the block are marked.
22
702
    pub fn new(num_of_elements: usize) -> Self {
23
702
        debug_assert!(num_of_elements > 0, "Cannot partition the empty set");
24

            
25
702
        let blocks = vec![Block::new(0, num_of_elements)];
26
702
        let elements = (0..num_of_elements).collect();
27

            
28
702
        Self { elements, blocks }
29
702
    }
30
}
31

            
32
impl<A: Clone + fmt::Debug + Default> BlockPartition<A> {
33
    /// Create a block partition from an indexed partition.
34
800
    pub fn from_indexed_partition(partition: &IndexedPartition) -> Self {
35
800
        let mut blocks = vec![Block::new_empty(); partition.num_of_blocks()];
36

            
37
        // Figure out the number of elements per block.
38
80000
        for element in partition.iter() {
39
80000
            blocks[element].end += 1;
40
80000
        }
41

            
42
        // Compute the start index for each block.
43
800
        let mut start = 0;
44
49342
        for block in &mut blocks {
45
49342
            let end = block.end;
46
49342
            block.begin = start;
47
49342
            block.end = start; // This will be updated when adding elements.
48
49342
            start += end;
49
49342
        }
50

            
51
        // Create the elements vector.
52
800
        let mut elements = vec![0; partition.len()];
53
80000
        for (element_index, block_index) in partition.iter().enumerate() {
54
80000
            // Add the element to the block, and update the end index.
55
80000
            let block = &mut blocks[block_index];
56
80000
            let pos = block.end;
57
80000
            elements[pos] = element_index;
58
80000
            block.end = pos + 1;
59
80000
        }
60

            
61
        // Remove empty blocks.
62
49342
        blocks.retain(|block| !block.is_empty());
63

            
64
800
        Self { elements, blocks }
65
800
    }
66

            
67
    /// Return a reference to the given block.
68
16794
    pub fn block(&self, block_index: BlockIndex) -> &Block<A> {
69
16794
        &self.blocks[block_index]
70
16794
    }
71

            
72
    /// Return a mutable reference to the block's annotation.
73
15815
    pub fn block_annotation(&mut self, block_index: BlockIndex) -> &mut A {
74
15815
        self.blocks[block_index].annotation_mut()
75
15815
    }
76

            
77
    /// Splits a block into two blocks according to the given predicate. If the
78
    /// predicate holds for all or none of the elements, no split occurs.
79
3474652
    pub fn split_block<F>(&mut self, block_index: BlockIndex, predicate: F) -> Option<BlockIndex>
80
3474652
    where
81
3474652
        F: Fn(usize) -> bool,
82
    {
83
        // Size of the new block.
84
3474652
        let mut size = 0usize;
85

            
86
12294544
        for state in self.blocks[block_index].begin..self.blocks[block_index].end {
87
12294544
            if predicate(self.elements[state]) {
88
75761
                self.elements.swap(self.blocks[block_index].begin + size, state);
89
75761
                size += 1;
90
12218783
            }
91
        }
92

            
93
        // The original block are now the first [begin, begin + size) elements
94
3474652
        if size == 0 || size == self.blocks[block_index].len() {
95
            // No split occurred
96
3469964
            return None;
97
4688
        }
98

            
99
        // Create a new block for the remaining elements
100
4688
        let new_block = Block::new(self.blocks[block_index].begin + size, self.blocks[block_index].end);
101
4688
        let last_block = self.blocks.len();
102
4688
        self.blocks.push(new_block);
103

            
104
        // Update the original block
105
4688
        let new_end = self.blocks[block_index].begin + size;
106
4688
        self.blocks[block_index].end = new_end;
107

            
108
4688
        trace!(
109
            "Split block {:?} into blocks {:?} and {:?}",
110
            block_index,
111
            block_index,
112
            BlockIndex::new(last_block)
113
        );
114
4688
        Some(BlockIndex::new(last_block))
115
3474652
    }
116

            
117
    /// Returns the number of blocks in the partition.
118
152416
    pub fn num_of_blocks(&self) -> usize {
119
152416
        self.blocks.len()
120
152416
    }
121

            
122
    /// Returns an iterator over the elements of a given block.
123
6661140
    pub fn iter_block(&self, block_index: BlockIndex) -> BlockIter<'_> {
124
6661140
        BlockIter {
125
6661140
            elements: &self.elements,
126
6661140
            index: self.blocks[block_index].begin,
127
6661140
            end: self.blocks[block_index].end,
128
6661140
        }
129
6661140
    }
130

            
131
    /// Returns an iterator over all blocks in the partition.
132
    pub fn iter(&self) -> impl Iterator<Item = &Block<A>> {
133
        self.blocks.iter()
134
    }
135

            
136
    /// Returns an iterator over all blocks in the partition.
137
    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut Block<A>> {
138
        self.blocks.iter_mut()
139
    }
140

            
141
    /// Returns the number of elements in the partition.
142
    pub fn len(&self) -> usize {
143
        self.elements.len()
144
    }
145

            
146
    /// Returns true iff the partition is empty.
147
    pub fn is_empty(&self) -> bool {
148
        self.elements.is_empty()
149
    }
150
}
151

            
152
impl<B: Clone + fmt::Debug> fmt::Display for BlockPartition<B> {
153
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
154
        let format = self
155
            .blocks
156
            .iter()
157
            .map(|block| format!("{{{}}}", block.iter(&self.elements).format(", ")))
158
            .format(", ");
159

            
160
        write!(f, "{{{}}}", format)
161
    }
162
}
163
/// A block that stores a subset of the elements in a partition.
164
///
165
/// # Details
166
///
167
/// It uses `start` and `end` to indicate a range start..end of elements in the
168
/// partition.
169
#[derive(Clone, Copy, Debug)]
170
pub struct Block<A: Clone + fmt::Debug> {
171
    begin: usize,
172
    end: usize,
173
    annotation: A,
174
}
175

            
176
impl<A: Clone + fmt::Debug + Default> Block<A> {
177
5390
    pub fn new(begin: usize, end: usize) -> Self {
178
5390
        debug_assert!(begin < end, "The range of this block is incorrect {begin}..{end}");
179
5390
        Block {
180
5390
            begin,
181
5390
            end,
182
5390
            annotation: Default::default(),
183
5390
        }
184
5390
    }
185

            
186
    /// Create an empty block at the given position.
187
800
    fn new_empty() -> Self {
188
800
        Block {
189
800
            begin: 0,
190
800
            end: 0,
191
800
            annotation: Default::default(),
192
800
        }
193
800
    }
194
}
195

            
196
impl<A: Clone + fmt::Debug> Block<A> {
197
    /// Returns an iterator over the elements in this block.
198
    pub fn iter<'a>(&self, elements: &'a [usize]) -> impl Iterator<Item = usize> + 'a {
199
        BlockIter {
200
            elements,
201
            index: self.begin,
202
            end: self.end,
203
        }
204
    }
205

            
206
    /// Returns the underlying annotation of this block.
207
16788
    pub fn annotation(&self) -> &A {
208
16788
        &self.annotation
209
16788
    }
210

            
211
    /// Returns the underlying annotation of this block.
212
15815
    pub fn annotation_mut(&mut self) -> &mut A {
213
15815
        &mut self.annotation
214
15815
    }
215

            
216
    /// Returns the number of elements in the block.
217
26312
    pub fn len(&self) -> usize {
218
26312
        self.assert_consistent();
219
26312
        self.end - self.begin
220
26312
    }
221

            
222
    /// Returns true iff the block is empty.
223
49342
    pub fn is_empty(&self) -> bool {
224
49342
        self.begin == self.end
225
49342
    }
226

            
227
    /// Returns true iff the block is consistent.
228
26312
    fn assert_consistent(&self) {
229
26312
        debug_assert!(self.begin < self.end, "The range of block {self:?} is incorrect");
230
26312
    }
231
}
232

            
233
pub struct BlockIter<'a> {
234
    elements: &'a [usize],
235
    index: usize,
236
    end: usize,
237
}
238

            
239
impl Iterator for BlockIter<'_> {
240
    type Item = usize;
241

            
242
308935044
    fn next(&mut self) -> Option<Self::Item> {
243
308935044
        if self.index < self.end {
244
237639249
            let element = self.elements[self.index];
245
237639249
            self.index += 1;
246
237639249
            Some(element)
247
        } else {
248
71295795
            None
249
        }
250
308935044
    }
251
}
252

            
253
#[cfg(test)]
254
mod tests {
255
    use merc_utilities::random_test;
256
    use rand::RngExt;
257

            
258
    use super::*;
259

            
260
    #[test]
261
1
    fn test_simple_block_partition() {
262
1
        let mut partition: BlockPartition<()> = BlockPartition::new(10);
263

            
264
1
        assert_eq!(partition.num_of_blocks(), 1);
265

            
266
1
        let initial_block = BlockIndex::new(0);
267
1
        assert_eq!(partition.block(initial_block).len(), 10);
268

            
269
10
        let block_index = partition.split_block(BlockIndex::new(0), |state| state < 5).unwrap();
270

            
271
1
        assert_eq!(partition.num_of_blocks(), 2);
272
1
        assert_eq!(partition.block(initial_block).len(), 5);
273
1
        assert_eq!(partition.block(block_index).len(), 5);
274
1
    }
275

            
276
    #[test]
277
1
    fn test_random_from_indexed_partition() {
278
100
        random_test(100, |rng| {
279
100
            let mut partition = IndexedPartition::new(100);
280

            
281
10000
            for element in 0..partition.len() {
282
10000
                partition.set_block(element, BlockIndex::new(rng.random_range(0..10)));
283
10000
            }
284
100
            trace!("Input partition {partition}");
285

            
286
100
            let block_partition: BlockPartition<()> = BlockPartition::from_indexed_partition(&partition);
287
100
            trace!("Output partition {block_partition}");
288

            
289
            // Check that each block in the block partition contains only
290
            // elements from the same indexed-partition block.
291
1000
            for block in 0..block_partition.num_of_blocks() {
292
1000
                let mut elements = block_partition.iter_block(BlockIndex::new(block));
293
1000
                let first = elements.next().unwrap();
294
1000
                let expected_block = partition.block(first);
295

            
296
9000
                for element in elements {
297
9000
                    assert_eq!(partition.block(element), expected_block, "Block {block} contains elements from different indexed-partition blocks");
298
                }
299
            }
300
100
        })
301
1
    }
302
}