option_entry/
lib.rs

1pub struct VacantEntry<'a, T> {
2    option: &'a mut Option<T>,
3}
4
5impl<'a, T> VacantEntry<'a, T> {
6    fn new(option: &'a mut Option<T>) -> Self {
7        assert!(option.is_none());
8        Self { option }
9    }
10
11    pub fn insert(self, value: T) -> &'a mut T {
12        self.option.insert(value)
13    }
14
15    pub fn insert_entry(self, value: T) -> OccupiedEntry<'a, T> {
16        let Self { option } = self;
17        *option = Some(value);
18        OccupiedEntry::new(option)
19    }
20}
21
22pub struct OccupiedEntry<'a, T> {
23    option: &'a mut Option<T>,
24}
25
26impl<'a, T> OccupiedEntry<'a, T> {
27    fn new(option: &'a mut Option<T>) -> Self {
28        assert!(option.is_some());
29        Self { option }
30    }
31
32    pub fn get(&self) -> &T {
33        self.option.as_ref().expect("OccupiedEntry is None?")
34    }
35
36    pub fn get_mut(&mut self) -> &mut T {
37        self.option.as_mut().expect("OccupiedEntry is None?")
38    }
39
40    pub fn insert(&mut self, value: T) -> T {
41        self.option.replace(value).expect("OccupiedEntry is None?")
42    }
43
44    pub fn into_mut(self) -> &'a mut T {
45        self.option.as_mut().expect("OccupiedEntry is None?")
46    }
47
48    pub fn remove(self) -> T {
49        self.option.take().expect("OccupiedEntry is None?")
50    }
51}
52
53pub enum Entry<'a, T> {
54    Vacant(VacantEntry<'a, T>),
55    Occupied(OccupiedEntry<'a, T>),
56}
57
58impl<'a, T> Entry<'a, T> {
59    fn into_option_mut(self) -> &'a mut Option<T> {
60        match self {
61            Entry::Vacant(VacantEntry { option }) => option,
62            Entry::Occupied(OccupiedEntry { option }) => option,
63        }
64    }
65
66    pub fn and_modify(self, f: impl FnOnce(&mut T)) -> Self {
67        let option = self.into_option_mut();
68        if let Some(value) = option {
69            f(value);
70        }
71        option.entry()
72    }
73
74    pub fn insert_entry(self, value: T) -> OccupiedEntry<'a, T> {
75        match self {
76            Entry::Occupied(mut entry) => {
77                entry.insert(value);
78                entry
79            }
80            Entry::Vacant(entry) => entry.insert_entry(value),
81        }
82    }
83
84    pub fn or_default(self) -> &'a mut T
85    where
86        T: Default,
87    {
88        self.into_option_mut().get_or_insert_default()
89    }
90
91    pub fn or_insert(self, default: T) -> &'a mut T {
92        match self {
93            Entry::Occupied(entry) => entry.into_mut(),
94            Entry::Vacant(entry) => entry.insert(default),
95        }
96    }
97
98    pub fn or_insert_with(self, default: impl FnOnce() -> T) -> &'a mut T {
99        match self {
100            Entry::Occupied(entry) => entry.into_mut(),
101            Entry::Vacant(entry) => entry.insert(default()),
102        }
103    }
104}
105
106mod private {
107    pub trait Sealed {}
108}
109
110pub trait OptionEntry: private::Sealed {
111    type T;
112    fn entry(&mut self) -> Entry<'_, Self::T>;
113}
114
115impl<T> private::Sealed for Option<T> {}
116
117impl<T> OptionEntry for Option<T> {
118    type T = T;
119
120    fn entry(&mut self) -> Entry<'_, Self::T> {
121        if self.is_none() {
122            Entry::Vacant(VacantEntry::new(self))
123        } else {
124            Entry::Occupied(OccupiedEntry::new(self))
125        }
126    }
127}