Make most shifts branchless

3 files changed, 106 insertions, 97 deletions

diff --git a/libraries/base/GHC/Base.hs b/libraries/base/GHC/Base.hs
index 8aa2312343..21ad469d27 100644
--- a/libraries/base/GHC/Base.hs
+++ b/libraries/base/GHC/Base.hs
@@ -1664,41 +1664,48 @@ split again.
 -- Note that these wrappers still produce undefined results when the
 -- second argument (the shift amount) is negative.
 
+-- | This function is used to implement branchless shifts. If the number of bits
+-- to shift is greater than or equal to the type size in bits, then the shift
+-- must return 0.  Instead of doing a test, we use a mask obtained via this
+-- function which is branchless too.
+--
+--    shift_mask m b
+--      | b < m     = 0xFF..FF
+--      | otherwise = 0
+--
+shift_mask :: Int# -> Int# -> Int#
+{-# INLINE shift_mask #-}
+shift_mask m b = negateInt# (b <# m)
+
 -- | Shift the argument left by the specified number of bits
 -- (which must be non-negative).
 shiftL# :: Word# -> Int# -> Word#
-a `shiftL#` b   | isTrue# (b >=# WORD_SIZE_IN_BITS#) = 0##
-                | otherwise                          = a `uncheckedShiftL#` b
+a `shiftL#` b = (a `uncheckedShiftL#` b) `and#` int2Word# (shift_mask WORD_SIZE_IN_BITS# b)
 
 -- | Shift the argument right by the specified number of bits
 -- (which must be non-negative).
 -- The "RL" means "right, logical" (as opposed to RA for arithmetic)
 -- (although an arithmetic right shift wouldn't make sense for Word#)
 shiftRL# :: Word# -> Int# -> Word#
-a `shiftRL#` b  | isTrue# (b >=# WORD_SIZE_IN_BITS#) = 0##
-                | otherwise                          = a `uncheckedShiftRL#` b
+a `shiftRL#` b = (a `uncheckedShiftRL#` b) `and#` int2Word# (shift_mask WORD_SIZE_IN_BITS# b)
 
 -- | Shift the argument left by the specified number of bits
 -- (which must be non-negative).
 iShiftL# :: Int# -> Int# -> Int#
-a `iShiftL#` b  | isTrue# (b >=# WORD_SIZE_IN_BITS#) = 0#
-                | otherwise                          = a `uncheckedIShiftL#` b
+a `iShiftL#` b = (a `uncheckedIShiftL#` b) `andI#` shift_mask WORD_SIZE_IN_BITS# b
 
 -- | Shift the argument right (signed) by the specified number of bits
 -- (which must be non-negative).
 -- The "RA" means "right, arithmetic" (as opposed to RL for logical)
 iShiftRA# :: Int# -> Int# -> Int#
-a `iShiftRA#` b | isTrue# (b >=# WORD_SIZE_IN_BITS#) = if isTrue# (a <# 0#)
-                                                          then (-1#)
-                                                          else 0#
+a `iShiftRA#` b | isTrue# (b >=# WORD_SIZE_IN_BITS#) = negateInt# (a <# 0#)
                 | otherwise                          = a `uncheckedIShiftRA#` b
 
 -- | Shift the argument right (unsigned) by the specified number of bits
 -- (which must be non-negative).
 -- The "RL" means "right, logical" (as opposed to RA for arithmetic)
 iShiftRL# :: Int# -> Int# -> Int#
-a `iShiftRL#` b | isTrue# (b >=# WORD_SIZE_IN_BITS#) = 0#
-                | otherwise                          = a `uncheckedIShiftRL#` b
+a `iShiftRL#` b = (a `uncheckedIShiftRL#` b) `andI#` shift_mask WORD_SIZE_IN_BITS# b
 
 -- Rules for C strings (the functions themselves are now in GHC.CString)
 {-# RULES
diff --git a/libraries/base/GHC/Int.hs b/libraries/base/GHC/Int.hs
index b5c5530688..a514f4d35b 100644
--- a/libraries/base/GHC/Int.hs
+++ b/libraries/base/GHC/Int.hs
@@ -222,8 +222,8 @@ instance Bits Int8 where
     bitSize i                 = finiteBitSize i
     isSigned _                = True
     popCount (I8# x#)         = I# (word2Int# (popCnt8# (int2Word# (int8ToInt# x#))))
-    bit                       = bitDefault
-    testBit                   = testBitDefault
+    bit i                     = bitDefault i
+    testBit a i               = testBitDefault a i
 
 -- | @since 4.6.0.0
 instance FiniteBits Int8 where
@@ -433,8 +433,8 @@ instance Bits Int16 where
     bitSize i                  = finiteBitSize i
     isSigned _                 = True
     popCount (I16# x#)         = I# (word2Int# (popCnt16# (int2Word# (int16ToInt# x#))))
-    bit                        = bitDefault
-    testBit                    = testBitDefault
+    bit i                      = bitDefault i
+    testBit a i                = testBitDefault a i
 
 -- | @since 4.6.0.0
 instance FiniteBits Int16 where
@@ -631,8 +631,8 @@ instance Bits Int32 where
     bitSize i                  = finiteBitSize i
     isSigned _                 = True
     popCount (I32# x#)         = I# (word2Int# (popCnt32# (int2Word# (int32ToInt# x#))))
-    bit                        = bitDefault
-    testBit                    = testBitDefault
+    bit i                      = bitDefault i
+    testBit a i                = testBitDefault a i
 
 -- | @since 4.6.0.0
 instance FiniteBits Int32 where
@@ -865,14 +865,14 @@ instance Bits Int64 where
     (I64# x#) `xor` (I64# y#)  = I64# (word64ToInt64# (int64ToWord64# x# `xor64#` int64ToWord64# y#))
     complement (I64# x#)       = I64# (word64ToInt64# (not64# (int64ToWord64# x#)))
     (I64# x#) `shift` (I# i#)
-        | isTrue# (i# >=# 0#)  = I64# (x# `iShiftL64#` i#)
-        | otherwise            = I64# (x# `iShiftRA64#` negateInt# i#)
+        | isTrue# (i# >=# 0#)  = I64# (x# `shiftLInt64#` i#)
+        | otherwise            = I64# (x# `shiftRAInt64#` negateInt# i#)
     (I64# x#) `shiftL` (I# i#)
-        | isTrue# (i# >=# 0#)  = I64# (x# `iShiftL64#` i#)
+        | isTrue# (i# >=# 0#)  = I64# (x# `shiftLInt64#` i#)
         | otherwise            = overflowError
     (I64# x#) `unsafeShiftL` (I# i#) = I64# (x# `uncheckedIShiftL64#` i#)
     (I64# x#) `shiftR` (I# i#)
-        | isTrue# (i# >=# 0#)  = I64# (x# `iShiftRA64#` i#)
+        | isTrue# (i# >=# 0#)  = I64# (x# `shiftRAInt64#` i#)
         | otherwise            = overflowError
     (I64# x#) `unsafeShiftR` (I# i#) = I64# (x# `uncheckedIShiftRA64#` i#)
     (I64# x#) `rotate` (I# i#)
@@ -889,23 +889,8 @@ instance Bits Int64 where
     isSigned _                 = True
     popCount (I64# x#)         =
         I# (word2Int# (popCnt64# (int64ToWord64# x#)))
-    bit                        = bitDefault
-    testBit                    = testBitDefault
-
--- give the 64-bit shift operations the same treatment as the 32-bit
--- ones (see GHC.Base), namely we wrap them in tests to catch the
--- cases when we're shifting more than 64 bits to avoid unspecified
--- behaviour in the C shift operations.
-
-iShiftL64#, iShiftRA64# :: Int64# -> Int# -> Int64#
-
-a `iShiftL64#` b  | isTrue# (b >=# 64#) = intToInt64# 0#
-                  | otherwise           = a `uncheckedIShiftL64#` b
-
-a `iShiftRA64#` b | isTrue# (b >=# 64#) = if isTrue# (a `ltInt64#` (intToInt64# 0#))
-                                          then intToInt64# (-1#)
-                                          else intToInt64# 0#
-                  | otherwise = a `uncheckedIShiftRA64#` b
+    bit i                      = bitDefault i
+    testBit a i                = testBitDefault a i
 
 -- No RULES for RealFrac methods if Int is smaller than Int64, we can't
 -- go through Int and whether going through Integer is faster is uncertain.
@@ -1090,46 +1075,63 @@ so the
 order gives us better code in the common case.
 -}
 
+-------------------------------------------------------------------------------
+
+-- unchecked shift primops may be lowered into C shift operations which have
+-- unspecified behaviour if the amount of bits to shift is greater or equal to the word
+-- size in bits.
+-- The following safe shift operations wrap unchecked primops to take this into
+-- account: 0 is consistently returned when the shift amount is too big.
+
 shiftRLInt8# :: Int8# -> Int# -> Int8#
-a `shiftRLInt8#` b | isTrue# (b >=# 8#) = intToInt8# 0#
-                   | otherwise          = a `uncheckedShiftRLInt8#` b
+a `shiftRLInt8#` b = uncheckedShiftRLInt8# a b `andInt8#` intToInt8# (shift_mask 8# b)
 
 shiftRLInt16# :: Int16# -> Int# -> Int16#
-a `shiftRLInt16#` b | isTrue# (b >=# 16#) = intToInt16# 0#
-                    | otherwise           = a `uncheckedShiftRLInt16#` b
+a `shiftRLInt16#` b = uncheckedShiftRLInt16# a b `andInt16#` intToInt16# (shift_mask 16# b)
 
 shiftRLInt32# :: Int32# -> Int# -> Int32#
-a `shiftRLInt32#` b | isTrue# (b >=# 32#) = intToInt32# 0#
-                    | otherwise           = a `uncheckedShiftRLInt32#` b
+a `shiftRLInt32#` b = uncheckedShiftRLInt32# a b `andInt32#` intToInt32# (shift_mask 32# b)
+
+
 
 shiftLInt8# :: Int8# -> Int# -> Int8#
-a `shiftLInt8#` b | isTrue# (b >=# 8#) = intToInt8# 0#
-                  | otherwise          = a `uncheckedShiftLInt8#` b
+a `shiftLInt8#` b = uncheckedShiftLInt8# a b `andInt8#` intToInt8# (shift_mask 8# b)
 
 shiftLInt16# :: Int16# -> Int# -> Int16#
-a `shiftLInt16#` b | isTrue# (b >=# 16#) = intToInt16# 0#
-                   | otherwise           = a `uncheckedShiftLInt16#` b
+a `shiftLInt16#` b = uncheckedShiftLInt16# a b `andInt16#` intToInt16# (shift_mask 16# b)
 
 shiftLInt32# :: Int32# -> Int# -> Int32#
-a `shiftLInt32#` b | isTrue# (b >=# 32#) = intToInt32# 0#
-                   | otherwise           = a `uncheckedShiftLInt32#` b
+a `shiftLInt32#` b = uncheckedShiftLInt32# a b `andInt32#` intToInt32# (shift_mask 32# b)
+
+shiftLInt64# :: Int64# -> Int# -> Int64#
+a `shiftLInt64#` b  = uncheckedIShiftL64# a b `andInt64#` intToInt64# (shift_mask 64# b)
 
 
 shiftRAInt8# :: Int8# -> Int# -> Int8#
-a `shiftRAInt8#` b | isTrue# (b >=# 8#) = if isTrue# (a `ltInt8#` (intToInt8# 0#))
-                                             then intToInt8# (-1#)
-                                             else intToInt8# 0#
+a `shiftRAInt8#` b | isTrue# (b >=# 8#) = intToInt8# (negateInt# (a `ltInt8#` (intToInt8# 0#)))
                    | otherwise          = a `uncheckedShiftRAInt8#` b
 
 shiftRAInt16# :: Int16# -> Int# -> Int16#
-a `shiftRAInt16#` b | isTrue# (b >=# 16#) = if isTrue# (a `ltInt16#` (intToInt16# 0#))
-                                               then intToInt16# (-1#)
-                                               else intToInt16# 0#
+a `shiftRAInt16#` b | isTrue# (b >=# 16#) = intToInt16# (negateInt# (a `ltInt16#` (intToInt16# 0#)))
                     | otherwise           = a `uncheckedShiftRAInt16#` b
 
 shiftRAInt32# :: Int32# -> Int# -> Int32#
-a `shiftRAInt32#` b | isTrue# (b >=# 32#) = if isTrue# (a `ltInt32#` (intToInt32# 0#))
-                                               then intToInt32# (-1#)
-                                               else intToInt32# 0#
+a `shiftRAInt32#` b | isTrue# (b >=# 32#) = intToInt32# (negateInt# (a `ltInt32#` (intToInt32# 0#)))
                     | otherwise           = a `uncheckedShiftRAInt32#` b
 
+shiftRAInt64# :: Int64# -> Int# -> Int64#
+a `shiftRAInt64#` b | isTrue# (b >=# 64#) = intToInt64# (negateInt# (a `ltInt64#` (intToInt64# 0#)))
+                    | otherwise           = a `uncheckedIShiftRA64#` b
+
+
+andInt8# :: Int8# -> Int8# -> Int8#
+x `andInt8#` y = word8ToInt8# (int8ToWord8# x `andWord8#` int8ToWord8# y)
+
+andInt16# :: Int16# -> Int16# -> Int16#
+x `andInt16#` y = word16ToInt16# (int16ToWord16# x `andWord16#` int16ToWord16# y)
+
+andInt32# :: Int32# -> Int32# -> Int32#
+x `andInt32#` y = word32ToInt32# (int32ToWord32# x `andWord32#` int32ToWord32# y)
+
+andInt64# :: Int64# -> Int64# -> Int64#
+x `andInt64#` y = word64ToInt64# (int64ToWord64# x `and64#` int64ToWord64# y)
diff --git a/libraries/base/GHC/Word.hs b/libraries/base/GHC/Word.hs
index 1dc0497856..512597aac5 100644
--- a/libraries/base/GHC/Word.hs
+++ b/libraries/base/GHC/Word.hs
@@ -210,8 +210,8 @@ instance Bits Word8 where
     bitSize i                 = finiteBitSize i
     isSigned _                = False
     popCount (W8# x#)         = I# (word2Int# (popCnt8# (word8ToWord# x#)))
-    bit                       = bitDefault
-    testBit                   = testBitDefault
+    bit i                     = bitDefault i
+    testBit a i               = testBitDefault a i
 
 -- | @since 4.6.0.0
 instance FiniteBits Word8 where
@@ -400,8 +400,8 @@ instance Bits Word16 where
     bitSize i                 = finiteBitSize i
     isSigned _                = False
     popCount (W16# x#)        = I# (word2Int# (popCnt16# (word16ToWord# x#)))
-    bit                       = bitDefault
-    testBit                   = testBitDefault
+    bit i                     = bitDefault i
+    testBit a i               = testBitDefault a i
 
 -- | @since 4.6.0.0
 instance FiniteBits Word16 where
@@ -627,8 +627,8 @@ instance Bits Word32 where
     bitSize i                 = finiteBitSize i
     isSigned _                = False
     popCount (W32# x#)        = I# (word2Int# (popCnt32# (word32ToWord# x#)))
-    bit                       = bitDefault
-    testBit                   = testBitDefault
+    bit i                     = bitDefault i
+    testBit a i               = testBitDefault a i
 
 -- | @since 4.6.0.0
 instance FiniteBits Word32 where
@@ -806,14 +806,14 @@ instance Bits Word64 where
     (W64# x#) `xor` (W64# y#)  = W64# (x# `xor64#` y#)
     complement (W64# x#)       = W64# (not64# x#)
     (W64# x#) `shift` (I# i#)
-        | isTrue# (i# >=# 0#)  = W64# (x# `shiftL64#` i#)
-        | otherwise            = W64# (x# `shiftRL64#` negateInt# i#)
+        | isTrue# (i# >=# 0#)  = W64# (x# `shiftLWord64#` i#)
+        | otherwise            = W64# (x# `shiftRLWord64#` negateInt# i#)
     (W64# x#) `shiftL`       (I# i#)
-        | isTrue# (i# >=# 0#)  = W64# (x# `shiftL64#` i#)
+        | isTrue# (i# >=# 0#)  = W64# (x# `shiftLWord64#` i#)
         | otherwise            = overflowError
     (W64# x#) `unsafeShiftL` (I# i#) = W64# (x# `uncheckedShiftL64#` i#)
     (W64# x#) `shiftR`       (I# i#)
-        | isTrue# (i# >=# 0#)  = W64# (x# `shiftRL64#` i#)
+        | isTrue# (i# >=# 0#)  = W64# (x# `shiftRLWord64#` i#)
         | otherwise            = overflowError
     (W64# x#) `unsafeShiftR` (I# i#) = W64# (x# `uncheckedShiftRL64#` i#)
     (W64# x#) `rotate` (I# i#)
@@ -826,22 +826,8 @@ instance Bits Word64 where
     bitSize i                 = finiteBitSize i
     isSigned _                = False
     popCount (W64# x#)        = I# (word2Int# (popCnt64# x#))
-    bit                       = bitDefault
-    testBit                   = testBitDefault
-
--- give the 64-bit shift operations the same treatment as the 32-bit
--- ones (see GHC.Base), namely we wrap them in tests to catch the
--- cases when we're shifting more than 64 bits to avoid unspecified
--- behaviour in the C shift operations.
-
-shiftL64#, shiftRL64# :: Word64# -> Int# -> Word64#
-
-a `shiftL64#` b  | isTrue# (b >=# 64#) = wordToWord64# 0##
-                 | otherwise           = a `uncheckedShiftL64#` b
-
-a `shiftRL64#` b | isTrue# (b >=# 64#) = wordToWord64# 0##
-                 | otherwise           = a `uncheckedShiftRL64#` b
-
+    bit i                     = bitDefault i
+    testBit a i               = testBitDefault a i
 
 -- | @since 4.6.0.0
 instance FiniteBits Word64 where
@@ -902,27 +888,41 @@ bitReverse64 (W64# w#) = W64# (bitReverse64# w#)
 
 -------------------------------------------------------------------------------
 
+-- unchecked shift primops may be lowered into C shift operations which have
+-- unspecified behaviour if the amount of bits to shift is greater or equal to the word
+-- size in bits.
+-- The following safe shift operations wrap unchecked primops to take this into
+-- account: 0 is consistently returned when the shift amount is too big.
+
 shiftRLWord8# :: Word8# -> Int# -> Word8#
-a `shiftRLWord8#` b | isTrue# (b >=# 8#) = wordToWord8# 0##
-                    | otherwise          = a `uncheckedShiftRLWord8#` b
+a `shiftRLWord8#` b = uncheckedShiftRLWord8# a b
+                       `andWord8#` wordToWord8# (int2Word# (shift_mask 8# b))
 
 shiftRLWord16# :: Word16# -> Int# -> Word16#
-a `shiftRLWord16#` b | isTrue# (b >=# 16#) = wordToWord16# 0##
-                     | otherwise           = a `uncheckedShiftRLWord16#` b
+a `shiftRLWord16#` b = uncheckedShiftRLWord16# a b
+                       `andWord16#` wordToWord16# (int2Word# (shift_mask 16# b))
 
 shiftRLWord32# :: Word32# -> Int# -> Word32#
-a `shiftRLWord32#` b | isTrue# (b >=# 32#) = wordToWord32# 0##
-                     | otherwise           = a `uncheckedShiftRLWord32#` b
+a `shiftRLWord32#` b = uncheckedShiftRLWord32# a b
+                       `andWord32#` wordToWord32# (int2Word# (shift_mask 32# b))
+
+shiftRLWord64# :: Word64# -> Int# -> Word64#
+a `shiftRLWord64#` b = uncheckedShiftRL64# a b
+                    `and64#` int64ToWord64# (intToInt64# (shift_mask 64# b))
 
 shiftLWord8# :: Word8# -> Int# -> Word8#
-a `shiftLWord8#` b | isTrue# (b >=# 8#) = wordToWord8# 0##
-                   | otherwise          = a `uncheckedShiftLWord8#` b
+a `shiftLWord8#` b = uncheckedShiftLWord8# a b
+                      `andWord8#` wordToWord8# (int2Word# (shift_mask 8# b))
 
 shiftLWord16# :: Word16# -> Int# -> Word16#
-a `shiftLWord16#` b | isTrue# (b >=# 16#) = wordToWord16# 0##
-                    | otherwise           = a `uncheckedShiftLWord16#` b
+a `shiftLWord16#` b = uncheckedShiftLWord16# a b
+                       `andWord16#` wordToWord16# (int2Word# (shift_mask 16# b))
 
 shiftLWord32# :: Word32# -> Int# -> Word32#
-a `shiftLWord32#` b | isTrue# (b >=# 32#) = wordToWord32# 0##
-                    | otherwise           = a `uncheckedShiftLWord32#` b
+a `shiftLWord32#` b = uncheckedShiftLWord32# a b
+                       `andWord32#` wordToWord32# (int2Word# (shift_mask 32# b))
+
+shiftLWord64# :: Word64# -> Int# -> Word64#
+a `shiftLWord64#` b  = uncheckedShiftL64# a b
+                    `and64#` int64ToWord64# (intToInt64# (shift_mask 64# b))