From c55cda301c943270b7eb2b4765bedbcce56edb90 Mon Sep 17 00:00:00 2001
From: Brandon Pickering <brandonpickering95@gmail.com>
Date: Sun, 6 Dec 2020 07:42:00 -0800
Subject: [PATCH] Make assignment operators panic safe

---
 src/lib.rs    | 15 +++++----------
 tests/test.rs | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index d74094c..bf94e89 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -377,8 +377,7 @@ impl<T: Float + AddAssign> AddAssign for NotNan<T> {
 /// Panics if the provided value is NaN.
 impl<T: Float + AddAssign> AddAssign<T> for NotNan<T> {
     fn add_assign(&mut self, other: T) {
-        self.0 += other;
-        assert!(!self.0.is_nan(), "Addition resulted in NaN");
+        *self = *self + other;
     }
 }
 
@@ -425,8 +424,7 @@ impl<T: Float + SubAssign> SubAssign for NotNan<T> {
 /// Panics if the provided value is NaN or the computation results in NaN
 impl<T: Float + SubAssign> SubAssign<T> for NotNan<T> {
     fn sub_assign(&mut self, other: T) {
-        self.0 -= other;
-        assert!(!self.0.is_nan(), "Subtraction resulted in NaN");
+        *self = *self - other;
     }
 }
 
@@ -460,8 +458,7 @@ impl<T: Float + MulAssign> MulAssign for NotNan<T> {
 /// Panics if the provided value is NaN.
 impl<T: Float + MulAssign> MulAssign<T> for NotNan<T> {
     fn mul_assign(&mut self, other: T) {
-        self.0 *= other;
-        assert!(!self.0.is_nan(), "Multiplication resulted in NaN");
+        *self = *self * other;
     }
 }
 
@@ -507,8 +504,7 @@ impl<T: Float + DivAssign> DivAssign for NotNan<T> {
 /// Panics if the provided value is NaN or the computation results in NaN
 impl<T: Float + DivAssign> DivAssign<T> for NotNan<T> {
     fn div_assign(&mut self, other: T) {
-        self.0 /= other;
-        assert!(!self.0.is_nan(), "Division resulted in NaN");
+        *self = *self / other;
     }
 }
 
@@ -542,8 +538,7 @@ impl<T: Float + RemAssign> RemAssign for NotNan<T> {
 /// Panics if the provided value is NaN or the computation results in NaN
 impl<T: Float + RemAssign> RemAssign<T> for NotNan<T> {
     fn rem_assign(&mut self, other: T) {
-        self.0 %= other;
-        assert!(!self.0.is_nan(), "Rem resulted in NaN");
+        *self = *self % other;
     }
 }
 
diff --git a/tests/test.rs b/tests/test.rs
index 2645681..8d134e0 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -590,3 +590,18 @@ fn not_nan_usage_in_const_context() {
     const A: NotNan<f32> = unsafe { NotNan::unchecked_new(111f32) };
     assert_eq!(A, NotNan::new(111f32).unwrap());
 }
+
+#[test]
+fn not_nan_panic_safety() {
+    let catch_op = |mut num, op: fn(&mut NotNan<_>)| {
+        let mut num_ref = panic::AssertUnwindSafe(&mut num);
+        let _ = panic::catch_unwind(move || op(*num_ref));
+        num
+    };
+
+    assert!(!catch_op(not_nan(f32::INFINITY), |a| *a += f32::NEG_INFINITY).is_nan());
+    assert!(!catch_op(not_nan(f32::INFINITY), |a| *a -= f32::INFINITY).is_nan());
+    assert!(!catch_op(not_nan(0.0), |a| *a *= f32::INFINITY).is_nan());
+    assert!(!catch_op(not_nan(0.0), |a| *a /= 0.0).is_nan());
+    assert!(!catch_op(not_nan(0.0), |a| *a %= 0.0).is_nan());
+}