1
use std::ops::{Index, IndexMut};
2

            
3
use crate::format::BasinId;
4

            
5
#[derive(Debug)]
6
pub struct BasinMap<T> {
7
    basins: [Option<T>; 8],
8
}
9

            
10
impl<T> BasinMap<T> {
11
663
    pub const fn new() -> Self {
12
663
        Self {
13
663
            basins: [None, None, None, None, None, None, None, None],
14
663
        }
15
663
    }
16

            
17
28653
    pub fn get_or_insert_with(&mut self, index: BasinId, default: impl FnOnce() -> T) -> &mut T {
18
28653
        if self[index].is_none() {
19
331
            self[index] = Some(default());
20
28322
        }
21

            
22
28653
        self[index].as_mut().expect("always initialized above")
23
28653
    }
24

            
25
31
    pub fn get_or_default(&mut self, index: BasinId) -> &mut T
26
31
    where
27
31
        T: Default,
28
31
    {
29
31
        self.get_or_insert_with(index, T::default)
30
31
    }
31
}
32

            
33
impl<T> Index<BasinId> for BasinMap<T> {
34
    type Output = Option<T>;
35

            
36
339864
    fn index(&self, index: BasinId) -> &Self::Output {
37
339864
        &self.basins[usize::from(index.index())]
38
339864
    }
39
}
40

            
41
impl<T> IndexMut<BasinId> for BasinMap<T> {
42
81011
    fn index_mut(&mut self, index: BasinId) -> &mut Self::Output {
43
81011
        &mut self.basins[usize::from(index.index())]
44
81011
    }
45
}
46

            
47
impl<'a, T> IntoIterator for &'a BasinMap<T> {
48
    type IntoIter = Iter<'a, T>;
49
    type Item = (BasinId, &'a T);
50

            
51
690
    fn into_iter(self) -> Self::IntoIter {
52
690
        Iter {
53
690
            map: self,
54
690
            id: Some(BasinId::MIN),
55
690
        }
56
690
    }
57
}
58

            
59
#[derive(Debug)]
60
pub struct Iter<'a, T> {
61
    map: &'a BasinMap<T>,
62
    id: Option<BasinId>,
63
}
64

            
65
impl<'a, T> Iterator for Iter<'a, T> {
66
    type Item = (BasinId, &'a T);
67

            
68
1373
    fn next(&mut self) -> Option<Self::Item> {
69
6210
        while let Some(id) = self.id {
70
5520
            let next_id = id.next();
71
5520
            let basin_contents = &self.map[id];
72
5520
            self.id = next_id;
73
            // If the basin had something, return it. Otherwise, look at the
74
            // next position.
75
5520
            if let Some(contents) = basin_contents {
76
683
                return Some((id, contents));
77
4837
            }
78
        }
79

            
80
690
        None
81
1373
    }
82
}