Implement BitOr and BitOrAssign for the NonZero integer types

2 files changed, 87 insertions, 0 deletions

diff --git a/src/libcore/num/mod.rs b/src/libcore/num/mod.rs
index 889ce2f211e..b46f399035a 100644
--- a/src/libcore/num/mod.rs
+++ b/src/libcore/num/mod.rs
@@ -8,6 +8,7 @@ use crate::convert::Infallible;
 use crate::fmt;
 use crate::intrinsics;
 use crate::mem;
+use crate::ops::{BitOr, BitOrAssign};
 use crate::str::FromStr;
 
 // Used because the `?` operator is not allowed in a const context.
@@ -110,6 +111,57 @@ assert_eq!(size_of::<Option<core::num::", stringify!($Ty), ">>(), size_of::<", s
                 }
             }
 
+            #[stable(feature = "nonzero_bitor", since = "1.43.0")]
+            impl BitOr for $Ty {
+                type Output = Self;
+                #[inline]
+                fn bitor(self, rhs: Self) -> Self::Output {
+                    // Safety: since `self` and `rhs` are both nonzero, the
+                    // result of the bitwise-or will be nonzero.
+                    unsafe { $Ty::new_unchecked(self.get() | rhs.get()) }
+                }
+            }
+
+            #[stable(feature = "nonzero_bitor", since = "1.43.0")]
+            impl BitOr<$Int> for $Ty {
+                type Output = Self;
+                #[inline]
+                fn bitor(self, rhs: $Int) -> Self::Output {
+                    // Safety: since `self` is nonzero, the result of the
+                    // bitwise-or will be nonzero regardless of the value of
+                    // `rhs`.
+                    unsafe { $Ty::new_unchecked(self.get() | rhs) }
+                }
+            }
+
+            #[stable(feature = "nonzero_bitor", since = "1.43.0")]
+            impl BitOr<$Ty> for $Int {
+                type Output = $Ty;
+                #[inline]
+                fn bitor(self, rhs: $Ty) -> Self::Output {
+                    // Safety: since `rhs` is nonzero, the result of the
+                    // bitwise-or will be nonzero regardless of the value of
+                    // `self`.
+                    unsafe { $Ty::new_unchecked(self | rhs.get()) }
+                }
+            }
+
+            #[stable(feature = "nonzero_bitor", since = "1.43.0")]
+            impl BitOrAssign for $Ty {
+                #[inline]
+                fn bitor_assign(&mut self, rhs: Self) {
+                    *self = *self | rhs;
+                }
+            }
+
+            #[stable(feature = "nonzero_bitor", since = "1.43.0")]
+            impl BitOrAssign<$Int> for $Ty {
+                #[inline]
+                fn bitor_assign(&mut self, rhs: $Int) {
+                    *self = *self | rhs;
+                }
+            }
+
             impl_nonzero_fmt! {
                 #[$stability] (Debug, Display, Binary, Octal, LowerHex, UpperHex) for $Ty
             }
diff --git a/src/libcore/tests/nonzero.rs b/src/libcore/tests/nonzero.rs
index 6c5d19845e4..0227a66b863 100644
--- a/src/libcore/tests/nonzero.rs
+++ b/src/libcore/tests/nonzero.rs
@@ -141,3 +141,38 @@ fn test_from_str() {
         Some(IntErrorKind::Overflow)
     );
 }
+
+#[test]
+fn test_nonzero_bitor() {
+    let nz_alt = NonZeroU8::new(0b1010_1010).unwrap();
+    let nz_low = NonZeroU8::new(0b0000_1111).unwrap();
+
+    let both_nz: NonZeroU8 = nz_alt | nz_low;
+    assert_eq!(both_nz.get(), 0b1010_1111);
+
+    let rhs_int: NonZeroU8 = nz_low | 0b1100_0000u8;
+    assert_eq!(rhs_int.get(), 0b1100_1111);
+
+    let rhs_zero: NonZeroU8 = nz_alt | 0u8;
+    assert_eq!(rhs_zero.get(), 0b1010_1010);
+
+    let lhs_int: NonZeroU8 = 0b0110_0110u8 | nz_alt;
+    assert_eq!(lhs_int.get(), 0b1110_1110);
+
+    let lhs_zero: NonZeroU8 = 0u8 | nz_low;
+    assert_eq!(lhs_zero.get(), 0b0000_1111);
+}
+
+#[test]
+fn test_nonzero_bitor_assign() {
+    let mut target = NonZeroU8::new(0b1010_1010).unwrap();
+
+    target |= NonZeroU8::new(0b0000_1111).unwrap();
+    assert_eq!(target.get(), 0b1010_1111);
+
+    target |= 0b0001_0000;
+    assert_eq!(target.get(), 0b1011_1111);
+
+    target |= 0;
+    assert_eq!(target.get(), 0b1011_1111);
+}