starnix_types/
augmented.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::ops::{Deref, DerefMut};
6
7/// An `Augmented` value is a generic wrapper that holds a primary value of type `T` and an
8/// optional auxiliary value of type `A`.
9///
10/// This is useful when a function needs to return a value that has some additional, optional
11/// context attached to it. For example, `A` can be used to override some of the fields of `T`
12/// without modifying `T` itself.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum Augmented<T, A: Clone> {
15    /// The primary value, without any auxiliary data.
16    Primary(T),
17    /// The primary value, with auxiliary data.
18    WithAux(T, A),
19}
20
21impl<T, A: Clone> Augmented<T, A> {
22    /// Maps an `Augmented<T, A>` to an `Augmented<U, A>` by applying a function to the contained
23    /// primary value, leaving the auxiliary value untouched.
24    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> Augmented<U, A> {
25        match self {
26            Self::Primary(t) => Augmented::Primary(f(t)),
27            Self::WithAux(t, aux) => Augmented::WithAux(f(t), aux),
28        }
29    }
30
31    /// Extracts the primary value from the `Augmented` value, discarding the auxiliary value if it
32    /// exists.
33    pub fn extract(self) -> T {
34        match self {
35            Self::Primary(t) => t,
36            Self::WithAux(t, _) => t,
37        }
38    }
39
40    /// Converts an `Augmented<T, A>` to an `Augmented<&T, A>`.
41    pub fn as_ref(&self) -> Augmented<&T, A> {
42        match self {
43            Self::Primary(t) => Augmented::Primary(t),
44            Self::WithAux(t, aux) => Augmented::WithAux(t, aux.clone()),
45        }
46    }
47
48    /// Converts an `Augmented<T, A>` to an `Augmented<&mut T, A>`.
49    pub fn as_mut(&mut self) -> Augmented<&mut T, A> {
50        match self {
51            Self::Primary(t) => Augmented::Primary(t),
52            Self::WithAux(t, aux) => Augmented::WithAux(t, aux.clone()),
53        }
54    }
55}
56
57impl<T, A: Clone> Augmented<&mut T, A> {
58    /// Converts an `Augmented<&mut T, A>` to an `Augmented<&T, A>`.
59    pub fn as_unmut(&self) -> Augmented<&T, A> {
60        match self {
61            Self::Primary(t) => Augmented::Primary(t),
62            Self::WithAux(t, aux) => Augmented::WithAux(t, aux.clone()),
63        }
64    }
65}
66
67impl<T, A: Clone> Augmented<Option<T>, A> {
68    /// Transposes an `Augmented<Option<T>, A>` into an `Option<Augmented<T, A>>`.
69    pub fn transpose(self) -> Option<Augmented<T, A>> {
70        match self {
71            Self::Primary(t) => Some(Augmented::Primary(t?)),
72            Self::WithAux(t, aux) => Some(Augmented::WithAux(t?, aux)),
73        }
74    }
75}
76
77impl<T, A: Clone, E> Augmented<Result<T, E>, A> {
78    /// Transposes an `Augmented<Result<T, E>, A>` into a `Result<Augmented<T, A>, E>`.
79    pub fn transpose(self) -> Result<Augmented<T, A>, E> {
80        match self {
81            Self::Primary(t) => Ok(Augmented::Primary(t?)),
82            Self::WithAux(t, aux) => Ok(Augmented::WithAux(t?, aux)),
83        }
84    }
85}
86
87impl<T, A: Clone> From<T> for Augmented<T, A> {
88    /// Creates a `Primary` `Augmented` value from a primary value.
89    fn from(t: T) -> Self {
90        Self::Primary(t)
91    }
92}
93
94impl<T, A: Clone> Deref for Augmented<T, A> {
95    type Target = T;
96
97    /// Dereferences the `Augmented` value to the primary value.
98    fn deref(&self) -> &Self::Target {
99        match self {
100            Self::Primary(t) => &t,
101            Self::WithAux(t, _) => &t,
102        }
103    }
104}
105
106impl<T, A: Clone> DerefMut for Augmented<T, A> {
107    /// Mutably dereferences the `Augmented` value to the primary value.
108    fn deref_mut(&mut self) -> &mut Self::Target {
109        match self {
110            Self::Primary(t) => t,
111            Self::WithAux(t, _) => t,
112        }
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn test_map() {
122        let primary = Augmented::<_, ()>::Primary(5);
123        let mapped_primary = primary.map(|x| x * 2);
124        assert!(matches!(mapped_primary, Augmented::<_, ()>::Primary(10)));
125
126        let with_aux = Augmented::WithAux(5, "hello");
127        let mapped_with_aux = with_aux.map(|x| x * 2);
128        assert!(matches!(mapped_with_aux, Augmented::WithAux(10, "hello")));
129    }
130
131    #[test]
132    fn test_extract() {
133        assert_eq!(Augmented::<_, ()>::Primary(5).extract(), 5);
134        assert_eq!(Augmented::WithAux(5, "hello").extract(), 5);
135    }
136
137    #[test]
138    fn test_as_ref() {
139        let primary = Augmented::<_, ()>::Primary(5);
140        let ref_primary = primary.as_ref();
141        assert!(matches!(ref_primary, Augmented::<_, ()>::Primary(&5)));
142
143        let with_aux = Augmented::WithAux(5, "hello");
144        let ref_with_aux = with_aux.as_ref();
145        assert!(matches!(ref_with_aux, Augmented::WithAux(&5, "hello")));
146    }
147
148    #[test]
149    fn test_as_mut() {
150        let mut primary = Augmented::<_, ()>::Primary(5);
151        if let Augmented::<_, ()>::Primary(val) = primary.as_mut() {
152            *val = 10;
153        }
154        assert_eq!(*primary, 10);
155
156        let mut with_aux = Augmented::WithAux(5, "hello");
157        if let Augmented::WithAux(val, aux) = with_aux.as_mut() {
158            assert_eq!(aux, "hello");
159            *val = 10;
160        }
161        assert_eq!(*with_aux, 10);
162    }
163
164    #[test]
165    fn test_as_unmut() {
166        let mut primary = Augmented::<_, ()>::Primary(5);
167        let mut_ref = primary.as_mut();
168        let unmut_ref = mut_ref.as_unmut();
169        assert!(matches!(unmut_ref, Augmented::<_, ()>::Primary(&5)));
170
171        let mut with_aux = Augmented::WithAux(5, "hello");
172        let mut_ref_aux = with_aux.as_mut();
173        let unmut_ref_aux = mut_ref_aux.as_unmut();
174        assert!(matches!(unmut_ref_aux, Augmented::WithAux(&5, "hello")));
175    }
176
177    #[test]
178    fn test_transpose_option() {
179        let primary_some = Augmented::<_, ()>::Primary(Some(5));
180        let transposed_primary_some = primary_some.transpose();
181        assert!(matches!(transposed_primary_some, Some(Augmented::<_, ()>::Primary(5))));
182
183        let primary_none: Augmented<Option<i32>, &str> = Augmented::Primary(None);
184        assert!(primary_none.transpose().is_none());
185
186        let with_aux_some = Augmented::WithAux(Some(5), "hello");
187        let transposed_with_aux_some = with_aux_some.transpose();
188        assert!(matches!(transposed_with_aux_some, Some(Augmented::WithAux(5, "hello"))));
189
190        let with_aux_none: Augmented<Option<i32>, &str> = Augmented::WithAux(None, "hello");
191        assert!(with_aux_none.transpose().is_none());
192    }
193
194    #[test]
195    fn test_transpose_result() {
196        let primary_ok: Augmented<Result<i32, &str>, &str> = Augmented::Primary(Ok(5));
197        let transposed_primary_ok = primary_ok.transpose();
198        assert!(matches!(transposed_primary_ok, Ok(Augmented::Primary(5))));
199
200        let primary_err: Augmented<Result<i32, &str>, &str> = Augmented::Primary(Err("error"));
201        assert_eq!(primary_err.transpose(), Err("error"));
202
203        let with_aux_ok: Augmented<Result<i32, &str>, &str> = Augmented::WithAux(Ok(5), "hello");
204        let transposed_with_aux_ok = with_aux_ok.transpose();
205        assert!(matches!(transposed_with_aux_ok, Ok(Augmented::WithAux(5, "hello"))));
206
207        let with_aux_err: Augmented<Result<i32, &str>, &str> =
208            Augmented::WithAux(Err("error"), "hello");
209        assert_eq!(with_aux_err.transpose(), Err("error"));
210    }
211
212    #[test]
213    fn test_from() {
214        let augmented: Augmented<i32, &str> = 5.into();
215        assert!(matches!(augmented, Augmented::Primary(5)));
216    }
217
218    #[test]
219    fn test_deref() {
220        assert_eq!(*Augmented::<_, ()>::Primary(5), 5);
221        assert_eq!(*Augmented::WithAux(5, "hello"), 5);
222    }
223
224    #[test]
225    fn test_deref_mut() {
226        let mut primary = Augmented::<_, ()>::Primary(5);
227        *primary = 10;
228        assert_eq!(*primary, 10);
229
230        let mut with_aux = Augmented::WithAux(5, "hello");
231        *with_aux = 10;
232        assert_eq!(*with_aux, 10);
233    }
234
235    #[derive(Clone, Copy, Debug, PartialEq)]
236    struct Data {
237        a: i32,
238        b: i32,
239    }
240
241    #[derive(Clone, Copy, Debug, PartialEq)]
242    struct DataOverride {
243        b: Option<i32>,
244    }
245
246    fn get_b(augmented: &Augmented<Data, DataOverride>) -> i32 {
247        match augmented {
248            Augmented::Primary(data) => data.b,
249            Augmented::WithAux(data, aux) => aux.b.unwrap_or(data.b),
250        }
251    }
252
253    #[test]
254    fn test_override_example() {
255        let primary = Augmented::Primary(Data { a: 1, b: 2 });
256        assert_eq!(get_b(&primary), 2);
257
258        let with_aux_override =
259            Augmented::WithAux(Data { a: 1, b: 2 }, DataOverride { b: Some(3) });
260        assert_eq!(get_b(&with_aux_override), 3);
261
262        let with_aux_no_override =
263            Augmented::WithAux(Data { a: 1, b: 2 }, DataOverride { b: None });
264        assert_eq!(get_b(&with_aux_no_override), 2);
265    }
266}
starnix_types/augmented.rs

starnix_types/
augmented.rs
