From 21dfae004cf778b0a45e9f38fb12c797f5e3d977 Mon Sep 17 00:00:00 2001
From: Fabian Schuiki <fabian@schuiki.ch>
Date: Wed, 3 Jan 2018 16:25:13 +0100
Subject: [PATCH] Add checked shifts

Add traits `CheckedShl` and `CheckedShr` that correspond to the standard
library's `checked_shl` and `checked_shr` functions. Implement the trait
on all primitive integer types by default, akin to what the standard
library does.

The stdlib is somewhat inconsistent when it comes to the type of the
shift amount. The `checked_*` functions have a `u32` shift amount, but
the `std::ops::{Shl,Shr}` traits are generic over the shift amount. Also
the stdlib implements these traits for all primitive integer types as
right-hand sides. Our implementation mimics this behaviour.
---
 src/int.rs         |  2 ++
 src/ops/checked.rs | 53 +++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/src/int.rs b/src/int.rs
index 4f9221f..18ce873 100644
--- a/src/int.rs
+++ b/src/int.rs
@@ -21,6 +21,8 @@ pub trait PrimInt
     + CheckedSub<Output=Self>
     + CheckedMul<Output=Self>
     + CheckedDiv<Output=Self>
+    + CheckedShl<usize, Output=Self>
+    + CheckedShr<usize, Output=Self>
     + Saturating
 {
     /// Returns the number of ones in the binary representation of `self`.
diff --git a/src/ops/checked.rs b/src/ops/checked.rs
index 45e6716..6527b5b 100644
--- a/src/ops/checked.rs
+++ b/src/ops/checked.rs
@@ -1,4 +1,4 @@
-use std::ops::{Add, Sub, Mul, Div};
+use std::ops::{Add, Sub, Mul, Div, Shl, Shr};
 
 /// Performs addition that returns `None` instead of wrapping around on
 /// overflow.
@@ -90,3 +90,54 @@ checked_impl!(CheckedDiv, checked_div, i32);
 checked_impl!(CheckedDiv, checked_div, i64);
 checked_impl!(CheckedDiv, checked_div, isize);
 
+/// Performs a left shift that returns `None` on overflow.
+pub trait CheckedShl<RHS>: Sized + Shl<RHS, Output=Self> {
+    /// Shifts a number to the left, checking for overflow. If overflow happens, `None` is
+    /// returned.
+    fn checked_shl(&self, rhs: &RHS) -> Option<Self>;
+}
+
+macro_rules! checked_shift_impl {
+    ($trait_name:ident, $method:ident, $rhs:ty, $t:ty) => {
+        impl $trait_name<$rhs> for $t {
+            #[inline]
+            fn $method(&self, rhs: &$rhs) -> Option<$t> {
+                // Note the cast to `u32` here: The standard library is somewhat inconsistent here.
+                // The `Shl<T>` and `Shr<T>` trait are generic over the right-hand side `T`, but
+                // the checked versions of the shifts all operate on `u32`.
+
+                // TODO: Maybe we should use a conversion that can fail here. This would allow us
+                // to catch the case where `rhs` exceeds the `u32` accepted by the stdlib, and
+                // return a `None` instead.
+                <$t>::$method(*self, *rhs as u32)
+            }
+        }
+    }
+}
+
+macro_rules! checked_shift_impl_all {
+    ($trait_name:ident, $method:ident, $($t:ty)*) => ($(
+        checked_shift_impl! { $trait_name, $method, u8   , $t }
+        checked_shift_impl! { $trait_name, $method, u16  , $t }
+        checked_shift_impl! { $trait_name, $method, u32  , $t }
+        checked_shift_impl! { $trait_name, $method, u64  , $t }
+        checked_shift_impl! { $trait_name, $method, usize, $t }
+
+        checked_shift_impl! { $trait_name, $method, i8   , $t }
+        checked_shift_impl! { $trait_name, $method, i16  , $t }
+        checked_shift_impl! { $trait_name, $method, i32  , $t }
+        checked_shift_impl! { $trait_name, $method, i64  , $t }
+        checked_shift_impl! { $trait_name, $method, isize, $t }
+    )*)
+}
+
+checked_shift_impl_all!(CheckedShl, checked_shl, u8 u16 u32 u64 usize i8 i16 i32 i64 isize);
+
+/// Performs a right shift that returns `None` on overflow.
+pub trait CheckedShr<RHS>: Sized + Shr<RHS, Output=Self> {
+    /// Shifts a number to the left, checking for overflow. If overflow happens, `None` is
+    /// returned.
+    fn checked_shr(&self, rhs: &RHS) -> Option<Self>;
+}
+
+checked_shift_impl_all!(CheckedShr, checked_shr, u8 u16 u32 u64 usize i8 i16 i32 i64 isize);