Add ctlz_nonzero & cttz_nonzero intrinsics

LLVM currently doesn't remove the "bypass if argument is zero" assembly inside branches where the value is known to be non-zero, pessimizing code that uses uN::leading_zeros

4 files changed, 80 insertions, 2 deletions

diff --git a/src/libcore/intrinsics.rs b/src/libcore/intrinsics.rs
index 3566bbdebc2..8188c15a282 100644
--- a/src/libcore/intrinsics.rs
+++ b/src/libcore/intrinsics.rs
@@ -1229,6 +1229,23 @@ extern "rust-intrinsic" {
     /// ```
     pub fn ctlz<T>(x: T) -> T;
 
+    /// Like `ctlz`, but extra-unsafe as it returns `undef` when
+    /// given an `x` with value `0`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(core_intrinsics)]
+    ///
+    /// use std::intrinsics::ctlz_nonzero;
+    ///
+    /// let x = 0b0001_1100_u8;
+    /// let num_leading = unsafe { ctlz_nonzero(x) };
+    /// assert_eq!(num_leading, 3);
+    /// ```
+    #[cfg(not(stage0))]
+    pub fn ctlz_nonzero<T>(x: T) -> T;
+
     /// Returns the number of trailing unset bits (zeroes) in an integer type `T`.
     ///
     /// # Examples
@@ -1256,6 +1273,23 @@ extern "rust-intrinsic" {
     /// ```
     pub fn cttz<T>(x: T) -> T;
 
+    /// Like `cttz`, but extra-unsafe as it returns `undef` when
+    /// given an `x` with value `0`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(core_intrinsics)]
+    ///
+    /// use std::intrinsics::cttz_nonzero;
+    ///
+    /// let x = 0b0011_1000_u8;
+    /// let num_trailing = unsafe { cttz_nonzero(x) };
+    /// assert_eq!(num_trailing, 3);
+    /// ```
+    #[cfg(not(stage0))]
+    pub fn cttz_nonzero<T>(x: T) -> T;
+
     /// Reverses the bytes in an integer type `T`.
     pub fn bswap<T>(x: T) -> T;
 
diff --git a/src/librustc_trans/intrinsic.rs b/src/librustc_trans/intrinsic.rs
index f2e6aa3ef00..de908bb24a7 100644
--- a/src/librustc_trans/intrinsic.rs
+++ b/src/librustc_trans/intrinsic.rs
@@ -267,7 +267,7 @@ pub fn trans_intrinsic_call<'a, 'tcx>(bcx: &Builder<'a, 'tcx>,
             };
             bcx.call(expect, &[llargs[0], C_i32(ccx, rw), llargs[1], C_i32(ccx, cache_type)], None)
         },
-        "ctlz" | "cttz" | "ctpop" | "bswap" |
+        "ctlz" | "ctlz_nonzero" | "cttz" | "cttz_nonzero" | "ctpop" | "bswap" |
         "add_with_overflow" | "sub_with_overflow" | "mul_with_overflow" |
         "overflowing_add" | "overflowing_sub" | "overflowing_mul" |
         "unchecked_div" | "unchecked_rem" | "unchecked_shl" | "unchecked_shr" => {
@@ -280,6 +280,12 @@ pub fn trans_intrinsic_call<'a, 'tcx>(bcx: &Builder<'a, 'tcx>,
                             let llfn = ccx.get_intrinsic(&format!("llvm.{}.i{}", name, width));
                             bcx.call(llfn, &[llargs[0], y], None)
                         }
+                        "ctlz_nonzero" | "cttz_nonzero" => {
+                            let y = C_bool(bcx.ccx, true);
+                            let llvm_name = &format!("llvm.{}.i{}", &name[..4], width);
+                            let llfn = ccx.get_intrinsic(llvm_name);
+                            bcx.call(llfn, &[llargs[0], y], None)
+                        }
                         "ctpop" => bcx.call(ccx.get_intrinsic(&format!("llvm.ctpop.i{}", width)),
                                         &llargs, None),
                         "bswap" => {
diff --git a/src/librustc_typeck/check/intrinsic.rs b/src/librustc_typeck/check/intrinsic.rs
index daf202cd797..4d9f50b0fc0 100644
--- a/src/librustc_typeck/check/intrinsic.rs
+++ b/src/librustc_typeck/check/intrinsic.rs
@@ -272,7 +272,8 @@ pub fn check_intrinsic_type<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
             "volatile_store" =>
                 (1, vec![ tcx.mk_mut_ptr(param(0)), param(0) ], tcx.mk_nil()),
 
-            "ctpop" | "ctlz" | "cttz" | "bswap" => (1, vec![param(0)], param(0)),
+            "ctpop" | "ctlz" | "ctlz_nonzero" | "cttz" | "cttz_nonzero" | "bswap" =>
+                (1, vec![param(0)], param(0)),
 
             "add_with_overflow" | "sub_with_overflow"  | "mul_with_overflow" =>
                 (1, vec![param(0), param(0)],
diff --git a/src/test/run-pass/intrinsics-integer.rs b/src/test/run-pass/intrinsics-integer.rs
index 759dc515456..4896f02da20 100644
--- a/src/test/run-pass/intrinsics-integer.rs
+++ b/src/test/run-pass/intrinsics-integer.rs
@@ -14,7 +14,9 @@ mod rusti {
     extern "rust-intrinsic" {
         pub fn ctpop<T>(x: T) -> T;
         pub fn ctlz<T>(x: T) -> T;
+        pub fn ctlz_nonzero<T>(x: T) -> T;
         pub fn cttz<T>(x: T) -> T;
+        pub fn cttz_nonzero<T>(x: T) -> T;
         pub fn bswap<T>(x: T) -> T;
     }
 }
@@ -68,6 +70,21 @@ pub fn main() {
         assert_eq!(ctlz(100u32), 25); assert_eq!(ctlz(100i32), 25);
         assert_eq!(ctlz(100u64), 57); assert_eq!(ctlz(100i64), 57);
 
+        assert_eq!(ctlz_nonzero(1u8), 7); assert_eq!(ctlz_nonzero(1i8), 7);
+        assert_eq!(ctlz_nonzero(1u16), 15); assert_eq!(ctlz_nonzero(1i16), 15);
+        assert_eq!(ctlz_nonzero(1u32), 31); assert_eq!(ctlz_nonzero(1i32), 31);
+        assert_eq!(ctlz_nonzero(1u64), 63); assert_eq!(ctlz_nonzero(1i64), 63);
+
+        assert_eq!(ctlz_nonzero(10u8), 4); assert_eq!(ctlz_nonzero(10i8), 4);
+        assert_eq!(ctlz_nonzero(10u16), 12); assert_eq!(ctlz_nonzero(10i16), 12);
+        assert_eq!(ctlz_nonzero(10u32), 28); assert_eq!(ctlz_nonzero(10i32), 28);
+        assert_eq!(ctlz_nonzero(10u64), 60); assert_eq!(ctlz_nonzero(10i64), 60);
+
+        assert_eq!(ctlz_nonzero(100u8), 1); assert_eq!(ctlz_nonzero(100i8), 1);
+        assert_eq!(ctlz_nonzero(100u16), 9); assert_eq!(ctlz_nonzero(100i16), 9);
+        assert_eq!(ctlz_nonzero(100u32), 25); assert_eq!(ctlz_nonzero(100i32), 25);
+        assert_eq!(ctlz_nonzero(100u64), 57); assert_eq!(ctlz_nonzero(100i64), 57);
+
         assert_eq!(cttz(-1i8 as u8), 0); assert_eq!(cttz(-1i8), 0);
         assert_eq!(cttz(-1i16 as u16), 0); assert_eq!(cttz(-1i16), 0);
         assert_eq!(cttz(-1i32 as u32), 0); assert_eq!(cttz(-1i32), 0);
@@ -93,6 +110,26 @@ pub fn main() {
         assert_eq!(cttz(100u32), 2); assert_eq!(cttz(100i32), 2);
         assert_eq!(cttz(100u64), 2); assert_eq!(cttz(100i64), 2);
 
+        assert_eq!(cttz_nonzero(-1i8 as u8), 0); assert_eq!(cttz_nonzero(-1i8), 0);
+        assert_eq!(cttz_nonzero(-1i16 as u16), 0); assert_eq!(cttz_nonzero(-1i16), 0);
+        assert_eq!(cttz_nonzero(-1i32 as u32), 0); assert_eq!(cttz_nonzero(-1i32), 0);
+        assert_eq!(cttz_nonzero(-1i64 as u64), 0); assert_eq!(cttz_nonzero(-1i64), 0);
+
+        assert_eq!(cttz_nonzero(1u8), 0); assert_eq!(cttz_nonzero(1i8), 0);
+        assert_eq!(cttz_nonzero(1u16), 0); assert_eq!(cttz_nonzero(1i16), 0);
+        assert_eq!(cttz_nonzero(1u32), 0); assert_eq!(cttz_nonzero(1i32), 0);
+        assert_eq!(cttz_nonzero(1u64), 0); assert_eq!(cttz_nonzero(1i64), 0);
+
+        assert_eq!(cttz_nonzero(10u8), 1); assert_eq!(cttz_nonzero(10i8), 1);
+        assert_eq!(cttz_nonzero(10u16), 1); assert_eq!(cttz_nonzero(10i16), 1);
+        assert_eq!(cttz_nonzero(10u32), 1); assert_eq!(cttz_nonzero(10i32), 1);
+        assert_eq!(cttz_nonzero(10u64), 1); assert_eq!(cttz_nonzero(10i64), 1);
+
+        assert_eq!(cttz_nonzero(100u8), 2); assert_eq!(cttz_nonzero(100i8), 2);
+        assert_eq!(cttz_nonzero(100u16), 2); assert_eq!(cttz_nonzero(100i16), 2);
+        assert_eq!(cttz_nonzero(100u32), 2); assert_eq!(cttz_nonzero(100i32), 2);
+        assert_eq!(cttz_nonzero(100u64), 2); assert_eq!(cttz_nonzero(100i64), 2);
+
         assert_eq!(bswap(0x0Au8), 0x0A); // no-op
         assert_eq!(bswap(0x0Ai8), 0x0A); // no-op
         assert_eq!(bswap(0x0A0Bu16), 0x0B0A);