1 // Copyright (C) 2024 The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 //! # safemath library
16 //!
17 //! This library provides an API to safely work with unsigned integers. At a high level, all math
18 //! operations are checked by default rather than having to remember to call specific `checked_*`
19 //! functions, so that the burden is on the programmer if they want to perform unchecked math
20 //! rather than the other way around:
21 //!
22 //! ```
23 //! use safemath::SafeNum;
24 //!
25 //! let safe = SafeNum::from(0);
26 //! let result = safe - 1;
27 //! assert!(u32::try_from(result).is_err());
28 //!
29 //! let safe_chain = (SafeNum::from(BIG_NUMBER) * HUGE_NUMBER) / MAYBE_ZERO;
30 //! // If any operation would have caused an overflow or division by zero,
31 //! // the number is flagged and the lexical location is specified for logging.
32 //! if safe_chain.has_error() {
33 //!     eprintln!("safe_chain error = {:#?}", safe_chain);
34 //! }
35 //! ```
36 //!
37 //! In addition to checked-by-default arithmetic, the API exposed here support
38 //! more natural usage than the `checked_*` functions by allowing chaining
39 //! of operations without having to check the result at each step.
40 //! This is similar to how floating-point `NaN` works - you can continue to use the
41 //! value, but continued operations will just propagate `NaN`.
42 //!
43 //! ## Supported Operations
44 //!
45 //! ### Arithmetic
46 //! The basic arithmetic operations are supported:
47 //! addition, subtraction, multiplication, division, and remainder.
48 //! The right hand side may be another SafeNum or any integer,
49 //! and the result is always another SafeNum.
50 //! If the operation would result in an overflow or division by zero,
51 //! or if converting the right hand element to a `u64` would cause an error,
52 //! the result is an error-tagged SafeNum that tracks the lexical origin of the error.
53 //!
54 //! ### Conversion from and to SafeNum
55 //! SafeNums support conversion to and from all integer types.
56 //! Conversion to SafeNum from signed integers and from usize and u128
57 //! can fail, generating an error value that is then propagated.
58 //! Conversion from SafeNum to all integers is only exposed via `try_from`
59 //! in order to force the user to handle potential resultant errors.
60 //!
61 //! E.g.
62 //! ```
63 //! fn call_func(_: u32, _: u32) {
64 //! }
65 //!
66 //! fn do_a_thing(a: SafeNum) -> Result<(), safemath::Error> {
67 //!     call_func(16, a.try_into()?);
68 //!     Ok(())
69 //! }
70 //! ```
71 //!
72 //! ### Comparison
73 //! SafeNums can be checked for equality against each other.
74 //! Valid numbers are equal to other numbers of the same magnitude.
75 //! Errored SafeNums are only equal to themselves.
76 //! Note that because errors propagate from their first introduction in an
77 //! arithmetic chain this can lead to surprising results.
78 //!
79 //! E.g.
80 //! ```
81 //! let overflow = SafeNum::MAX + 1;
82 //! let otherflow = SafeNum::MAX + 1;
83 //!
84 //! assert_ne!(overflow, otherflow);
85 //! assert_eq!(overflow + otherflow, overflow);
86 //! assert_eq!(otherflow + overflow, otherflow);
87 //! ```
88 //!
89 //! Inequality comparison operators are deliberately not provided.
90 //! By necessity they would have similar caveats to floating point comparisons,
91 //! which are easy to use incorrectly and unintuitive to use correctly.
92 //!
93 //! The required alternative is to convert to a real integer type before comparing,
94 //! forcing any errors upwards.
95 //!
96 //! E.g.
97 //! ```
98 //! impl From<safemath::Error> for &'static str {
99 //!     fn from(_: safemath::Error) -> Self {
100 //!         "checked arithmetic error"
101 //!     }
102 //! }
103 //!
104 //! fn my_op(a: SafeNum, b: SafeNum, c: SafeNum, d: SafeNum) -> Result<bool, &'static str> {
105 //!     Ok(safemath::Primitive::try_from(a)? < b.try_into()?
106 //!        && safemath::Primitive::try_from(c)? >= d.try_into()?)
107 //! }
108 //! ```
109 //!
110 //! ### Miscellaneous
111 //! SafeNums also provide helper methods to round up or down
112 //! to the nearest multiple of another number
113 //! and helper predicate methods that indicate whether the SafeNum
114 //! is valid or is tracking an error.
115 //!
116 //! Also provided are constants `SafeNum::MAX`, `SafeNum::MIN`, and `SafeNum::ZERO`.
117 //!
118 //! Warning: SafeNums can help prevent, isolate, and detect arithmetic overflow
119 //!          but they are not a panacea. In particular, chains of different operations
120 //!          are not guaranteed to be associative or commutative.
121 //!
122 //! E.g.
123 //! ```
124 //! let a = SafeNum::MAX - 1 + 1;
125 //! let b = SafeNum::MAX + 1 - 1;
126 //! assert_ne!(a, b);
127 //! assert!(a.is_valid());
128 //! assert!(b.has_error());
129 //!
130 //! let c = (SafeNum::MAX + 31) / 31;
131 //! let d = SafeNum::MAX / 31 + 31 / 31;
132 //! assert_ne!(c, d);
133 //! assert!(c.has_error());
134 //! assert!(d.is_valid());
135 //! ```
136 //!
137 //! Note:    SafeNum arithmetic is much slower than arithmetic on integer primitives.
138 //!          If you are concerned about performance, be sure to run benchmarks.
139 
140 #![cfg_attr(not(test), no_std)]
141 
142 use core::convert::TryFrom;
143 use core::fmt;
144 use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Rem, RemAssign, Sub, SubAssign};
145 use core::panic::Location;
146 
147 pub type Primitive = u64;
148 pub type Error = &'static Location<'static>;
149 
150 #[derive(Copy, Clone, PartialEq, Eq)]
151 pub struct SafeNum(Result<Primitive, Error>);
152 
153 impl fmt::Debug for SafeNum {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result154     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
155         match self.0 {
156             Ok(val) => write!(f, "{}", val),
157             Err(location) => write!(f, "error at {}", location),
158         }
159     }
160 }
161 
162 impl SafeNum {
163     pub const MAX: SafeNum = SafeNum(Ok(u64::MAX));
164     pub const MIN: SafeNum = SafeNum(Ok(u64::MIN));
165     pub const ZERO: SafeNum = SafeNum(Ok(0));
166 
167     /// Round `self` down to the nearest multiple of `rhs`.
168     #[track_caller]
round_down<T>(self, rhs: T) -> Self where Self: Rem<T, Output = Self>,169     pub fn round_down<T>(self, rhs: T) -> Self
170     where
171         Self: Rem<T, Output = Self>,
172     {
173         self - (self % rhs)
174     }
175 
176     /// Round `self` up to the nearest multiple of `rhs`.
177     #[track_caller]
round_up<T>(self, rhs: T) -> Self where Self: Add<T, Output = Self>, T: Copy + Into<Self>,178     pub fn round_up<T>(self, rhs: T) -> Self
179     where
180         Self: Add<T, Output = Self>,
181         T: Copy + Into<Self>,
182     {
183         ((self + rhs) - 1).round_down(rhs)
184     }
185 
186     /// Returns whether self is the result of an operation that has errored.
has_error(&self) -> bool187     pub const fn has_error(&self) -> bool {
188         self.0.is_err()
189     }
190 
191     /// Returns whether self represents a valid, non-overflowed integer.
is_valid(&self) -> bool192     pub const fn is_valid(&self) -> bool {
193         self.0.is_ok()
194     }
195 }
196 
197 macro_rules! try_conversion_func {
198     ($other_type:tt) => {
199         impl TryFrom<SafeNum> for $other_type {
200             type Error = Error;
201 
202             #[track_caller]
203             fn try_from(val: SafeNum) -> Result<Self, Self::Error> {
204                 Self::try_from(val.0?).map_err(|_| Location::caller())
205             }
206         }
207     };
208 }
209 
210 macro_rules! conversion_func {
211     ($from_type:tt) => {
212         impl From<$from_type> for SafeNum {
213             fn from(val: $from_type) -> SafeNum {
214                 Self(Ok(val.into()))
215             }
216         }
217 
218         try_conversion_func!($from_type);
219     };
220 }
221 
222 macro_rules! conversion_func_maybe_error {
223     ($from_type:tt) => {
224         impl From<$from_type> for SafeNum {
225             #[track_caller]
226             fn from(val: $from_type) -> Self {
227                 Self(Primitive::try_from(val).map_err(|_| Location::caller()))
228             }
229         }
230 
231         try_conversion_func!($from_type);
232     };
233 }
234 
235 macro_rules! arithmetic_impl {
236     ($trait_name:ident, $op:ident, $assign_trait_name:ident, $assign_op:ident, $func:ident) => {
237         impl<T: Into<SafeNum>> $trait_name<T> for SafeNum {
238             type Output = Self;
239             #[track_caller]
240             fn $op(self, rhs: T) -> Self {
241                 let rhs: Self = rhs.into();
242 
243                 match (self.0, rhs.0) {
244                     (Err(_), _) => self,
245                     (_, Err(_)) => rhs,
246                     (Ok(lhs), Ok(rhs)) => Self(lhs.$func(rhs).ok_or_else(Location::caller)),
247                 }
248             }
249         }
250 
251         impl<T> $assign_trait_name<T> for SafeNum
252         where
253             Self: $trait_name<T, Output = Self>,
254         {
255             #[track_caller]
256             fn $assign_op(&mut self, rhs: T) {
257                 *self = self.$op(rhs)
258             }
259         }
260     };
261 }
262 
263 conversion_func!(u8);
264 conversion_func!(u16);
265 conversion_func!(u32);
266 conversion_func!(u64);
267 conversion_func_maybe_error!(usize);
268 conversion_func_maybe_error!(u128);
269 conversion_func_maybe_error!(i8);
270 conversion_func_maybe_error!(i16);
271 conversion_func_maybe_error!(i32);
272 conversion_func_maybe_error!(i64);
273 conversion_func_maybe_error!(i128);
274 conversion_func_maybe_error!(isize);
275 arithmetic_impl!(Add, add, AddAssign, add_assign, checked_add);
276 arithmetic_impl!(Sub, sub, SubAssign, sub_assign, checked_sub);
277 arithmetic_impl!(Mul, mul, MulAssign, mul_assign, checked_mul);
278 arithmetic_impl!(Div, div, DivAssign, div_assign, checked_div);
279 arithmetic_impl!(Rem, rem, RemAssign, rem_assign, checked_rem);
280 
281 #[cfg(test)]
282 mod test {
283     use super::*;
284 
285     #[test]
test_addition()286     fn test_addition() {
287         let a: SafeNum = 2100.into();
288         let b: SafeNum = 12.into();
289         assert_eq!(a + b, 2112.into());
290     }
291 
292     #[test]
test_subtraction()293     fn test_subtraction() {
294         let a: SafeNum = 667.into();
295         let b: SafeNum = 1.into();
296         assert_eq!(a - b, 666.into());
297     }
298 
299     #[test]
test_multiplication()300     fn test_multiplication() {
301         let a: SafeNum = 17.into();
302         let b: SafeNum = 3.into();
303         assert_eq!(a * b, 51.into());
304     }
305 
306     #[test]
test_division()307     fn test_division() {
308         let a: SafeNum = 1066.into();
309         let b: SafeNum = 41.into();
310         assert_eq!(a / b, 26.into());
311     }
312 
313     #[test]
test_remainder()314     fn test_remainder() {
315         let a: SafeNum = 613.into();
316         let b: SafeNum = 10.into();
317         assert_eq!(a % b, 3.into());
318     }
319 
320     #[test]
test_addition_poison()321     fn test_addition_poison() {
322         let base: SafeNum = 2.into();
323         let poison = base + SafeNum::MAX;
324         assert!(u64::try_from(poison).is_err());
325 
326         let a = poison - 1;
327         let b = poison - 2;
328 
329         assert_eq!(a, poison);
330         assert_eq!(b, poison);
331     }
332 
333     #[test]
test_subtraction_poison()334     fn test_subtraction_poison() {
335         let base: SafeNum = 2.into();
336         let poison = base - SafeNum::MAX;
337         assert!(u64::try_from(poison).is_err());
338 
339         let a = poison + 1;
340         let b = poison + 2;
341 
342         assert_eq!(a, poison);
343         assert_eq!(b, poison);
344     }
345 
346     #[test]
test_multiplication_poison()347     fn test_multiplication_poison() {
348         let base: SafeNum = 2.into();
349         let poison = base * SafeNum::MAX;
350         assert!(u64::try_from(poison).is_err());
351 
352         let a = poison / 2;
353         let b = poison / 4;
354 
355         assert_eq!(a, poison);
356         assert_eq!(b, poison);
357     }
358 
359     #[test]
test_division_poison()360     fn test_division_poison() {
361         let base: SafeNum = 2.into();
362         let poison = base / 0;
363         assert!(u64::try_from(poison).is_err());
364 
365         let a = poison * 2;
366         let b = poison * 4;
367 
368         assert_eq!(a, poison);
369         assert_eq!(b, poison);
370     }
371 
372     #[test]
test_remainder_poison()373     fn test_remainder_poison() {
374         let base: SafeNum = 2.into();
375         let poison = base % 0;
376         assert!(u64::try_from(poison).is_err());
377 
378         let a = poison * 2;
379         let b = poison * 4;
380 
381         assert_eq!(a, poison);
382         assert_eq!(b, poison);
383     }
384 
385     macro_rules! conversion_test {
386         ($name:ident) => {
387             mod $name {
388                 use super::*;
389                 use core::convert::TryInto;
390 
391                 #[test]
392                 fn test_between_safenum() {
393                     let var: $name = 16;
394                     let sn: SafeNum = var.into();
395                     let res: $name = sn.try_into().unwrap();
396                     assert_eq!(var, res);
397                 }
398 
399                 #[test]
400                 fn test_arithmetic_safenum() {
401                     let primitive: $name = ((((0 + 11) * 11) / 3) % 32) - 3;
402                     let safe = ((((SafeNum::ZERO + $name::try_from(11u8).unwrap())
403                         * $name::try_from(11u8).unwrap())
404                         / $name::try_from(3u8).unwrap())
405                         % $name::try_from(32u8).unwrap())
406                         - $name::try_from(3u8).unwrap();
407                     assert_eq!($name::try_from(safe).unwrap(), primitive);
408                 }
409             }
410         };
411     }
412 
413     conversion_test!(u8);
414     conversion_test!(u16);
415     conversion_test!(u32);
416     conversion_test!(u64);
417     conversion_test!(u128);
418     conversion_test!(usize);
419     conversion_test!(i8);
420     conversion_test!(i16);
421     conversion_test!(i32);
422     conversion_test!(i64);
423     conversion_test!(i128);
424     conversion_test!(isize);
425 
426     macro_rules! correctness_tests {
427         ($name:ident, $operation:ident, $assign_operation:ident) => {
428             mod $operation {
429                 use super::*;
430                 use core::ops::$name;
431 
432                 #[test]
433                 fn test_correctness() {
434                     let normal = 300u64;
435                     let safe: SafeNum = normal.into();
436                     let rhs = 7u64;
437                     assert_eq!(
438                         u64::try_from(safe.$operation(rhs)).unwrap(),
439                         normal.$operation(rhs)
440                     );
441                 }
442 
443                 #[test]
444                 fn test_assign() {
445                     let mut var: SafeNum = 2112.into();
446                     let rhs = 666u64;
447                     let expect = var.$operation(rhs);
448                     var.$assign_operation(rhs);
449                     assert_eq!(var, expect);
450                 }
451 
452                 #[test]
453                 fn test_assign_poison() {
454                     let mut var = SafeNum::MIN - 1;
455                     let expected = var - 1;
456                     var.$assign_operation(2);
457                     // Poison saturates and doesn't perform additional changes
458                     assert_eq!(var, expected);
459                 }
460             }
461         };
462     }
463 
464     correctness_tests!(Add, add, add_assign);
465     correctness_tests!(Sub, sub, sub_assign);
466     correctness_tests!(Mul, mul, mul_assign);
467     correctness_tests!(Div, div, div_assign);
468     correctness_tests!(Rem, rem, rem_assign);
469 
470     #[test]
test_round_down()471     fn test_round_down() {
472         let x: SafeNum = 255.into();
473         assert_eq!(x.round_down(32), 224.into());
474         assert_eq!((x + 1).round_down(64), 256.into());
475         assert_eq!(x.round_down(256), SafeNum::ZERO);
476         assert!(x.round_down(SafeNum::MIN).has_error());
477     }
478 
479     #[test]
test_round_up()480     fn test_round_up() {
481         let x: SafeNum = 255.into();
482         assert_eq!(x.round_up(32), 256.into());
483         assert_eq!(x.round_up(51), x);
484         assert_eq!(SafeNum::ZERO.round_up(x), SafeNum::ZERO);
485         assert!(SafeNum::MAX.round_up(32).has_error());
486     }
487 }
488