Add `FiniteBits(count{Leading,Trailing}Zeros)`

This exposes the newly added CLZ/CTZ primops from
e0c1767d0ea8d12e0a4badf43682a08784e379c6 (re #9340)
via two new methods `countLeadingZeros` and `countTrailingZeros`
in the `Data.Bits.FiniteBits` class.

The original proposal can be found at

  http://www.haskell.org/pipermail/libraries/2014-August/023567.html

Test Plan: successful validate

Reviewers: ekmett, tibbe

GHC Trac Issues: #9532

Differential Revision: https://phabricator.haskell.org/D158

8 files changed, 186 insertions, 2 deletions

diff --git a/libraries/base/Data/Bits.hs b/libraries/base/Data/Bits.hs
index 81b180bb7e..a751176441 100644
--- a/libraries/base/Data/Bits.hs
+++ b/libraries/base/Data/Bits.hs
@@ -39,7 +39,11 @@ module Data.Bits (
     rotateL, rotateR,
     popCount
   ),
-  FiniteBits(finiteBitSize),
+  FiniteBits(
+    finiteBitSize,
+    countLeadingZeros,
+    countTrailingZeros
+  ),
 
   bitDefault,
   testBitDefault,
@@ -288,6 +292,65 @@ class Bits b => FiniteBits b where
     -- /Since: 4.7.0.0/
     finiteBitSize :: b -> Int
 
+    -- | Count number of zero bits preceding the most significant set bit.
+    --
+    -- @
+    -- 'countLeadingZeros' ('zeroBits' :: a) = finiteBitSize ('zeroBits' :: a)
+    -- 'countLeadingZeros' . 'negate' = 'const' 0
+    -- @
+    --
+    -- 'countLeadingZeros' can be used to compute log base 2 via
+    --
+    -- @
+    -- logBase2 x = 'finiteBitSize' x - 1 - 'countLeadingZeros' x
+    -- @
+    --
+    -- Note: The default implementation for this method is intentionally
+    -- naive. However, the instances provided for the primitive
+    -- integral types are implemented using CPU specific machine
+    -- instructions.
+    --
+    -- /Since: 4.8.0.0/
+    countLeadingZeros :: b -> Int
+    countLeadingZeros x = (w-1) - go (w-1)
+      where
+        go i | i < 0       = i -- no bit set
+             | testBit x i = i
+             | otherwise   = go (i-1)
+
+        w = finiteBitSize x
+
+    -- | Count number of zero bits following the least significant set bit.
+    --
+    -- @
+    -- 'countTrailingZeros' ('zeroBits' :: a) = finiteBitSize ('zeroBits' :: a)
+    -- 'countTrailingZeros' . 'negate' = 'countTrailingZeros'
+    -- @
+    --
+    -- The related
+    -- <http://en.wikipedia.org/wiki/Find_first_set find-first-set operation>
+    -- can be expressed in terms of 'countTrailingZeros' as follows
+    --
+    -- @
+    -- findFirstSet x = 1 + 'countTrailingZeros' x
+    -- @
+    --
+    -- Note: The default implementation for this method is intentionally
+    -- naive. However, the instances provided for the primitive
+    -- integral types are implemented using CPU specific machine
+    -- instructions.
+    --
+    -- /Since: 4.8.0.0/
+    countTrailingZeros :: b -> Int
+    countTrailingZeros x = go 0
+      where
+        go i | i >= w      = i
+             | testBit x i = i
+             | otherwise   = go (i+1)
+
+        w = finiteBitSize x
+
+
 -- The defaults below are written with lambdas so that e.g.
 --     bit = bitDefault
 -- is fully applied, so inlining will happen
@@ -356,7 +419,8 @@ instance Bits Bool where
 
 instance FiniteBits Bool where
     finiteBitSize _ = 1
-
+    countTrailingZeros x = if x then 0 else 1
+    countLeadingZeros  x = if x then 0 else 1
 
 instance Bits Int where
     {-# INLINE shift #-}
@@ -396,6 +460,8 @@ instance Bits Int where
 
 instance FiniteBits Int where
     finiteBitSize _ = WORD_SIZE_IN_BITS
+    countLeadingZeros  (I# x#) = I# (word2Int# (clz# (int2Word# x#)))
+    countTrailingZeros (I# x#) = I# (word2Int# (ctz# (int2Word# x#)))
 
 instance Bits Word where
     {-# INLINE shift #-}
@@ -429,6 +495,8 @@ instance Bits Word where
 
 instance FiniteBits Word where
     finiteBitSize _ = WORD_SIZE_IN_BITS
+    countLeadingZeros  (W# x#) = I# (word2Int# (clz# x#))
+    countTrailingZeros (W# x#) = I# (word2Int# (ctz# x#))
 
 instance Bits Integer where
    (.&.) = andInteger
diff --git a/libraries/base/GHC/Int.hs b/libraries/base/GHC/Int.hs
index 467b3f4e30..a9743ce1a0 100644
--- a/libraries/base/GHC/Int.hs
+++ b/libraries/base/GHC/Int.hs
@@ -165,6 +165,8 @@ instance Bits Int8 where
 
 instance FiniteBits Int8 where
     finiteBitSize _ = 8
+    countLeadingZeros  (I8# x#) = I# (word2Int# (clz8# (int2Word# x#)))
+    countTrailingZeros (I8# x#) = I# (word2Int# (ctz8# (int2Word# x#)))
 
 {-# RULES
 "fromIntegral/Int8->Int8" fromIntegral = id :: Int8 -> Int8
@@ -324,6 +326,8 @@ instance Bits Int16 where
 
 instance FiniteBits Int16 where
     finiteBitSize _ = 16
+    countLeadingZeros  (I16# x#) = I# (word2Int# (clz16# (int2Word# x#)))
+    countTrailingZeros (I16# x#) = I# (word2Int# (ctz16# (int2Word# x#)))
 
 {-# RULES
 "fromIntegral/Word8->Int16"  fromIntegral = \(W8# x#) -> I16# (word2Int# x#)
@@ -489,6 +493,8 @@ instance Bits Int32 where
 
 instance FiniteBits Int32 where
     finiteBitSize _ = 32
+    countLeadingZeros  (I32# x#) = I# (word2Int# (clz32# (int2Word# x#)))
+    countTrailingZeros (I32# x#) = I# (word2Int# (ctz32# (int2Word# x#)))
 
 {-# RULES
 "fromIntegral/Word8->Int32"  fromIntegral = \(W8# x#) -> I32# (word2Int# x#)
@@ -871,6 +877,13 @@ uncheckedIShiftRA64# = uncheckedIShiftRA#
 
 instance FiniteBits Int64 where
     finiteBitSize _ = 64
+#if WORD_SIZE_IN_BITS < 64
+    countLeadingZeros  (I64# x#) = I# (word2Int# (clz64# (int64ToWord64# x#)))
+    countTrailingZeros (I64# x#) = I# (word2Int# (ctz64# (int64ToWord64# x#)))
+#else
+    countLeadingZeros  (I64# x#) = I# (word2Int# (clz64# (int2Word# x#)))
+    countTrailingZeros (I64# x#) = I# (word2Int# (ctz64# (int2Word# x#)))
+#endif
 
 instance Real Int64 where
     toRational x = toInteger x % 1
diff --git a/libraries/base/GHC/Word.hs b/libraries/base/GHC/Word.hs
index 86978dc9c4..6721d07950 100644
--- a/libraries/base/GHC/Word.hs
+++ b/libraries/base/GHC/Word.hs
@@ -154,6 +154,8 @@ instance Bits Word8 where
 
 instance FiniteBits Word8 where
     finiteBitSize _ = 8
+    countLeadingZeros  (W8# x#) = I# (word2Int# (clz8# x#))
+    countTrailingZeros (W8# x#) = I# (word2Int# (ctz8# x#))
 
 {-# RULES
 "fromIntegral/Word8->Word8"   fromIntegral = id :: Word8 -> Word8
@@ -301,6 +303,8 @@ instance Bits Word16 where
 
 instance FiniteBits Word16 where
     finiteBitSize _ = 16
+    countLeadingZeros  (W16# x#) = I# (word2Int# (clz16# x#))
+    countTrailingZeros (W16# x#) = I# (word2Int# (ctz16# x#))
 
 -- | Swap bytes in 'Word16'.
 --
@@ -495,6 +499,8 @@ instance Bits Word32 where
 
 instance FiniteBits Word32 where
     finiteBitSize _ = 32
+    countLeadingZeros  (W32# x#) = I# (word2Int# (clz32# x#))
+    countTrailingZeros (W32# x#) = I# (word2Int# (ctz32# x#))
 
 {-# RULES
 "fromIntegral/Word8->Word32"   fromIntegral = \(W8# x#) -> W32# x#
@@ -767,6 +773,8 @@ uncheckedShiftRL64# = uncheckedShiftRL#
 
 instance FiniteBits Word64 where
     finiteBitSize _ = 64
+    countLeadingZeros  (W64# x#) = I# (word2Int# (clz64# x#))
+    countTrailingZeros (W64# x#) = I# (word2Int# (ctz64# x#))
 
 instance Show Word64 where
     showsPrec p x = showsPrec p (toInteger x)
diff --git a/libraries/base/changelog.md b/libraries/base/changelog.md
index b9768119be..97a82429ec 100644
--- a/libraries/base/changelog.md
+++ b/libraries/base/changelog.md
@@ -22,6 +22,9 @@
 
   * Re-export `Data.Word.Word` from `Prelude`
 
+  * Add `countLeadingZeros` and `countTrailingZeros` methods to
+    `Data.Bits.FiniteBits` class
+
 ## 4.7.0.1  *Jul 2014*
 
   * Bundled with GHC 7.8.3
diff --git a/libraries/base/tests/.gitignore b/libraries/base/tests/.gitignore
index 3115fd753c..973ab9d661 100644
--- a/libraries/base/tests/.gitignore
+++ b/libraries/base/tests/.gitignore
@@ -270,3 +270,4 @@
 /unicode002
 /weak001
 /T9395
+/T9532
diff --git a/libraries/base/tests/T9532.hs b/libraries/base/tests/T9532.hs
new file mode 100644
index 0000000000..e99a42b572
--- /dev/null
+++ b/libraries/base/tests/T9532.hs
@@ -0,0 +1,89 @@
+-- Tests Data.Bits.FiniteBits(count{Leading,Trailing}Zeros)` -- c.f. T9340.hs
+
+import Control.Monad
+import Data.Bits
+import Data.Int
+import Data.Typeable
+import Data.Word
+import Numeric (showHex)
+
+-- Reference Implementations
+
+-- count trailing zeros
+ctzRI :: FiniteBits a => a -> Word
+ctzRI x = fromIntegral $ go 0
+  where
+    go i | i >= w      = i
+         | testBit x i = i
+         | otherwise   = go (i+1)
+
+    w = finiteBitSize x
+
+-- count leading zeros
+clzRI :: FiniteBits a => a -> Word
+clzRI x = fromIntegral $ (w-1) - go (w-1)
+  where
+    go i | i < 0       = i -- no bit set
+         | testBit x i = i
+         | otherwise   = go (i-1)
+
+    w = finiteBitSize x
+
+-- Test Driver
+main :: IO ()
+main = do
+    forM_ testpats $ \w64 -> do
+        checkCLZ (fromIntegral w64 :: Word)
+        checkCLZ (fromIntegral w64 :: Word8)
+        checkCLZ (fromIntegral w64 :: Word16)
+        checkCLZ (fromIntegral w64 :: Word32)
+        checkCLZ (fromIntegral w64 :: Word64)
+
+        checkCLZ (fromIntegral w64 :: Int)
+        checkCLZ (fromIntegral w64 :: Int8)
+        checkCLZ (fromIntegral w64 :: Int16)
+        checkCLZ (fromIntegral w64 :: Int32)
+        checkCLZ (fromIntegral w64 :: Int64)
+
+        checkCTZ (fromIntegral w64 :: Word)
+        checkCTZ (fromIntegral w64 :: Word8)
+        checkCTZ (fromIntegral w64 :: Word16)
+        checkCTZ (fromIntegral w64 :: Word32)
+        checkCTZ (fromIntegral w64 :: Word64)
+
+        checkCTZ (fromIntegral w64 :: Int)
+        checkCTZ (fromIntegral w64 :: Int8)
+        checkCTZ (fromIntegral w64 :: Int16)
+        checkCTZ (fromIntegral w64 :: Int32)
+        checkCTZ (fromIntegral w64 :: Int64)
+
+    putStrLn $ concat ["tested ", show (length testpats), " patterns"]
+
+  where
+    -- try to construct some interesting patterns
+    testpats :: [Word64]
+    testpats = [ bit i - 1 | i <- [0..63] ] ++
+               [ complement (bit i - 1) | i <- [0..63] ] ++
+               [ bit i .|. bit j | i <- [0..63], j <- [0..i] ]
+
+    -- Compare impl-under-test with reference-impl
+    checkCLZ :: (Typeable a, Show a, Integral a, FiniteBits a) => a -> IO ()
+    checkCLZ v = unless (vri == viut) $ do
+        putStrLn $ concat [ "FAILED: clz (0x", showHex v " :: ", tyName
+                          , ") ==> (RI=", show vri, " vs. IUT=", show viut, ")"
+                          ]
+      where
+        tyName = show (typeOf v)
+        vri    = clzRI v
+        viut   = fromIntegral (countLeadingZeros v)
+
+    -- Compare impl-under-test with reference-impl
+    checkCTZ :: (Typeable a, Show a, Integral a, FiniteBits a) => a -> IO ()
+    checkCTZ v = unless (vri == viut) $ do
+        putStrLn $ concat [ "FAILED: ctz (0x", showHex v " :: ", tyName
+                          , ") ==> (RI=", show vri, " vs. IUT=", show viut, ")"
+                          ]
+      where
+        tyName = show (typeOf v)
+        vri    = ctzRI v
+        viut   = fromIntegral (countTrailingZeros v)
diff --git a/libraries/base/tests/T9532.stdout b/libraries/base/tests/T9532.stdout
new file mode 100644
index 0000000000..455b0abc18
--- /dev/null
+++ b/libraries/base/tests/T9532.stdout
@@ -0,0 +1 @@
+tested 2208 patterns
diff --git a/libraries/base/tests/all.T b/libraries/base/tests/all.T
index 8b18d63f01..5fe862f449 100644
--- a/libraries/base/tests/all.T
+++ b/libraries/base/tests/all.T
@@ -169,3 +169,4 @@ test('T8766',
 
 test('T9111', normal, compile, [''])
 test('T9395', normal, compile_and_run, [''])
+test('T9532', normal, compile_and_run, [''])