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
1502
    pub fn new(num_of_elements: usize) -> Self {
23
1502
        debug_assert!(num_of_elements > 0, "Cannot partition the empty set");
24

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

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

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

            
38
        // Figure out the number of elements per block.
39
86628
        for (_, block_index) in partition.iter_elements() {
40
86628
            blocks[block_index].end += 1;
41
86628
        }
42

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

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

            
62
        // Remove empty blocks.
63
84481
        blocks.retain(|block| !block.is_empty());
64

            
65
1300
        Self { elements, blocks }
66
1300
    }
67

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

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

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

            
87
156164783
        for state in self.blocks[block_index].begin..self.blocks[block_index].end {
88
156164783
            if predicate(self.elements[state]) {
89
1409282
                self.elements.swap(self.blocks[block_index].begin + size, state);
90
1409282
                size += 1;
91
154755501
            }
92
        }
93

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
260
    use super::*;
261

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

            
266
1
        assert_eq!(partition.num_of_blocks(), 1);
267

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

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

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

            
278
    #[test]
279
1
    fn test_random_from_indexed_partition() {
280
100
        random_test(100, |rng| {
281
100
            let subset = (0..100).sample(rng, 25);
282
100
            let mut partition = IndexedPartition::with_subset(100, subset.iter().copied());
283

            
284
2500
            for element in subset {
285
2500
                partition.set_block(element, BlockIndex::new(rng.random_range(0..10)));
286
2500
            }
287
100
            trace!("Input partition {partition}");
288

            
289
100
            let block_partition: BlockPartition<()> = BlockPartition::from_indexed_partition(&partition);
290
100
            trace!("Output partition {block_partition}");
291

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

            
299
1563
                for element in elements {
300
1563
                    assert_eq!(
301
1563
                        partition.block(element),
302
                        expected_block,
303
                        "Block {block} contains elements from different indexed-partition blocks"
304
                    );
305
                }
306
            }
307
100
        })
308
1
    }
309
}