{-# LANGUAGE CPP, ForeignFunctionInterface, BangPatterns #-}
#if __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif
#ifndef HLINT
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
#endif
--------------------------------------------------------------------
-- |
-- Copyright :  (c) Edward Kmett 2013-2014
-- License   :  BSD3
-- Maintainer:  Edward Kmett <ekmett@gmail.com>
-- Stability :  experimental
-- Portability: non-portable
--
-- Calculate a number of fiddly bit operations using fast de Bruijn
-- multiplication tables.
--------------------------------------------------------------------
module Data.Bits.Extras
  ( Ranked(..)
  , log2
  , msb
  , w8
  , w16
  , w32
  , w64
  , assignBit
  , zeroBits
  , oneBits
  , srl
  ) where

import Data.Bits
import Data.Int
import Data.Word
import Foreign.Ptr
import Foreign.Storable
import GHC.Base

-- TODO: generalize to 64 bits, etc.
log2 :: Word32 -> Int
log2 :: Word32 -> Int
log2 !Word32
n0 = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> Word8 -> Int
forall a b. (a -> b) -> a -> b
$ Word32 -> Word8
go (Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR (Word32
n5 Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word32
0x7C4ACDD) Int
27) where
  go :: Word32 -> Word8
  go :: Word32 -> Word8
go !Word32
i = IO Word8 -> Word8
forall a. IO a -> a
inlinePerformIO (IO Word8 -> Word8) -> IO Word8 -> Word8
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
debruijn_log32 (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i)
  !n1 :: Word32
n1 = Word32
n0 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
n0 Int
1
  !n2 :: Word32
n2 = Word32
n1 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
n1 Int
2
  !n3 :: Word32
n3 = Word32
n2 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
n2 Int
4
  !n4 :: Word32
n4 = Word32
n3 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
n3 Int
8
  !n5 :: Word32
n5 = Word32
n4 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
n4 Int
16
{-# INLINE log2 #-}

class (Num t, FiniteBits t) => Ranked t where
  -- | Calculate the least significant set bit using a debruijn multiplication table.
  -- /NB:/ The result of this function is undefined when given 0.
  lsb :: t -> Int
  lsb t
n = t -> Int
forall t. Ranked t => t -> Int
rank t
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  {-# INLINE lsb #-}

  -- | Calculate the number of trailing 0 bits.
  rank :: t -> Int
  rank t
0 = Int
0
  rank t
n = t -> Int
forall t. Ranked t => t -> Int
lsb t
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  {-# INLINE rank #-}

  -- | Calculate the number of leading zeros.
  nlz :: t -> Int

instance Ranked Word64 where
  lsb :: Word64 -> Int
lsb Word64
n = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> Word8 -> Int
forall a b. (a -> b) -> a -> b
$ Word64 -> Word8
go (Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR ((Word64
n Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. (-Word64
n)) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
0x07EDD5E59A4E28C2) Int
58) where
    go :: Word64 -> Word8
    go :: Word64 -> Word8
go Word64
i = IO Word8 -> Word8
forall a. IO a -> a
inlinePerformIO (IO Word8 -> Word8) -> IO Word8 -> Word8
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
debruijn_lsb64 (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
i)
  {-# INLINE lsb #-}

  nlz :: Word64 -> Int
nlz Word64
x0 = Word64 -> Int
forall a. Bits a => a -> Int
popCount (Word64 -> Word64
forall a. Bits a => a -> a
complement Word64
x6) where
     x1 :: Word64
x1 = Word64
x0 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
x0 Int
1
     x2 :: Word64
x2 = Word64
x1 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
x1 Int
2
     x3 :: Word64
x3 = Word64
x2 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
x2 Int
4
     x4 :: Word64
x4 = Word64
x3 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
x3 Int
8
     x5 :: Word64
x5 = Word64
x4 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
x4 Int
16
     x6 :: Word64
x6 = Word64
x5 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftR Word64
x5 Int
32
  {-# INLINE nlz #-}

instance Ranked Word32 where
  lsb :: Word32 -> Int
lsb Word32
n = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> Word8 -> Int
forall a b. (a -> b) -> a -> b
$ Word32 -> Word8
go (Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR ((Word32
n Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. (-Word32
n)) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
* Word32
0x077CB531) Int
27) where
    go :: Word32 -> Word8
    go :: Word32 -> Word8
go Word32
i = IO Word8 -> Word8
forall a. IO a -> a
inlinePerformIO (IO Word8 -> Word8) -> IO Word8 -> Word8
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Int -> IO Word8
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr Word8
debruijn_lsb32 (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i)
  {-# INLINE lsb #-}

{-
  rank n = fromIntegral $ go (unsafeShiftR ((n .&. (-n)) * 0x4279976B) 26) where
    go :: Word32 -> Word8
    go i = inlinePerformIO $ peekElemOff debruijn_rank32 (fromIntegral i)
  {-# INLINE rank #-}
-}

  nlz :: Word32 -> Int
nlz Word32
x0 = Word32 -> Int
forall a. Bits a => a -> Int
popCount (Word32 -> Word32
forall a. Bits a => a -> a
complement Word32
x5) where
     x1 :: Word32
x1 = Word32
x0 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
x0 Int
1
     x2 :: Word32
x2 = Word32
x1 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
x1 Int
2
     x3 :: Word32
x3 = Word32
x2 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
x2 Int
4
     x4 :: Word32
x4 = Word32
x3 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
x3 Int
8
     x5 :: Word32
x5 = Word32
x4 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
unsafeShiftR Word32
x4 Int
16
  {-# INLINE nlz #-}


instance Ranked Word16 where
  lsb :: Word16 -> Int
lsb = Word32 -> Int
forall t. Ranked t => t -> Int
lsb (Word32 -> Int) -> (Word16 -> Word32) -> Word16 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word16 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE lsb #-}

  rank :: Word16 -> Int
rank = Word32 -> Int
forall t. Ranked t => t -> Int
rank (Word32 -> Int) -> (Word16 -> Word32) -> Word16 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word16 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE rank #-}

  nlz :: Word16 -> Int
nlz Word16
x0 = Word16 -> Int
forall a. Bits a => a -> Int
popCount (Word16 -> Word16
forall a. Bits a => a -> a
complement Word16
x4) where
     x1 :: Word16
x1 = Word16
x0 Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
unsafeShiftR Word16
x0 Int
1
     x2 :: Word16
x2 = Word16
x1 Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
unsafeShiftR Word16
x1 Int
2
     x3 :: Word16
x3 = Word16
x2 Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
unsafeShiftR Word16
x2 Int
4
     x4 :: Word16
x4 = Word16
x3 Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
unsafeShiftR Word16
x3 Int
8
  {-# INLINE nlz #-}

instance Ranked Word8 where
  lsb :: Word8 -> Int
lsb = Word32 -> Int
forall t. Ranked t => t -> Int
lsb (Word32 -> Int) -> (Word8 -> Word32) -> Word8 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE lsb #-}

  rank :: Word8 -> Int
rank = Word32 -> Int
forall t. Ranked t => t -> Int
rank (Word32 -> Int) -> (Word8 -> Word32) -> Word8 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE rank #-}

  nlz :: Word8 -> Int
nlz Word8
x0 = Word8 -> Int
forall a. Bits a => a -> Int
popCount (Word8 -> Word8
forall a. Bits a => a -> a
complement Word8
x3) where
     x1 :: Word8
x1 = Word8
x0 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR Word8
x0 Int
1
     x2 :: Word8
x2 = Word8
x1 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR Word8
x1 Int
2
     x3 :: Word8
x3 = Word8
x2 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
unsafeShiftR Word8
x2 Int
4
  {-# INLINE nlz #-}

instance Ranked Int64 where
  lsb :: Int64 -> Int
lsb = Word64 -> Int
forall t. Ranked t => t -> Int
lsb (Word64 -> Int) -> (Int64 -> Word64) -> Int64 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Word64
forall a. Integral a => a -> Word64
w64
  {-# INLINE lsb #-}

  rank :: Int64 -> Int
rank = Word64 -> Int
forall t. Ranked t => t -> Int
rank (Word64 -> Int) -> (Int64 -> Word64) -> Int64 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Word64
forall a. Integral a => a -> Word64
w64
  {-# INLINE rank #-}

  nlz :: Int64 -> Int
nlz = Word64 -> Int
forall t. Ranked t => t -> Int
nlz (Word64 -> Int) -> (Int64 -> Word64) -> Int64 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Word64
forall a. Integral a => a -> Word64
w64
  {-# INLINE nlz #-}

instance Ranked Int32 where
  lsb :: Int32 -> Int
lsb = Word32 -> Int
forall t. Ranked t => t -> Int
lsb (Word32 -> Int) -> (Int32 -> Word32) -> Int32 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE lsb #-}

  rank :: Int32 -> Int
rank = Word32 -> Int
forall t. Ranked t => t -> Int
rank (Word32 -> Int) -> (Int32 -> Word32) -> Int32 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE rank #-}

  nlz :: Int32 -> Int
nlz = Word32 -> Int
forall t. Ranked t => t -> Int
nlz (Word32 -> Int) -> (Int32 -> Word32) -> Int32 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE nlz #-}

instance Ranked Int16 where
  lsb :: Int16 -> Int
lsb = Word32 -> Int
forall t. Ranked t => t -> Int
lsb (Word32 -> Int) -> (Int16 -> Word32) -> Int16 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE lsb #-}

  rank :: Int16 -> Int
rank = Word32 -> Int
forall t. Ranked t => t -> Int
rank (Word32 -> Int) -> (Int16 -> Word32) -> Int16 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE rank #-}

  nlz :: Int16 -> Int
nlz = Word16 -> Int
forall t. Ranked t => t -> Int
nlz (Word16 -> Int) -> (Int16 -> Word16) -> Int16 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> Word16
forall a. Integral a => a -> Word16
w16
  {-# INLINE nlz #-}

instance Ranked Int8 where
  lsb :: Int8 -> Int
lsb = Word32 -> Int
forall t. Ranked t => t -> Int
lsb (Word32 -> Int) -> (Int8 -> Word32) -> Int8 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE lsb #-}

  rank :: Int8 -> Int
rank = Word32 -> Int
forall t. Ranked t => t -> Int
rank (Word32 -> Int) -> (Int8 -> Word32) -> Int8 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> Word32
forall a. Integral a => a -> Word32
w32
  {-# INLINE rank #-}

  nlz :: Int8 -> Int
nlz = Word8 -> Int
forall t. Ranked t => t -> Int
nlz (Word8 -> Int) -> (Int8 -> Word8) -> Int8 -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> Word8
forall a. Integral a => a -> Word8
w8
  {-# INLINE nlz #-}

------------------------------------------------------------------------------
-- Util
------------------------------------------------------------------------------

w8 :: Integral a => a -> Word8
w8 :: a -> Word8
w8 = a -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE w8 #-}

w16 :: Integral a => a -> Word16
w16 :: a -> Word16
w16 = a -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE w16 #-}

w32 :: Integral a => a -> Word32
w32 :: a -> Word32
w32 = a -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE w32 #-}

w64 :: Integral a => a -> Word64
w64 :: a -> Word64
w64 = a -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE w64 #-}

-- | Calculate the most significant set bit.
msb :: Ranked t => t -> Int
msb :: t -> Int
msb t
n = t -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize t
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- t -> Int
forall t. Ranked t => t -> Int
nlz t
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
{-# INLINE msb #-}

assignBit :: Bits b => b -> Int -> Bool -> b
assignBit :: b -> Int -> Bool -> b
assignBit b
b Int
n  Bool
True = b
b b -> Int -> b
forall a. Bits a => a -> Int -> a
`setBit` Int
n
assignBit b
b Int
n Bool
False = b
b b -> Int -> b
forall a. Bits a => a -> Int -> a
`clearBit` Int
n
{-# INLINE assignBit #-}

oneBits :: Bits b => b
oneBits :: b
oneBits  = b -> b
forall a. Bits a => a -> a
complement b
forall a. Bits a => a
zeroBits

-- | Shift Right Logical (i.e., without sign extension)
--
-- /NB:/ When used on negative 'Integer's, hilarity may ensue.
srl :: Bits b => b -> Int -> b
srl :: b -> Int -> b
srl b
b Int
n = (b
b b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
n) b -> b -> b
forall a. Bits a => a -> a -> a
.&. b -> Int -> b
forall a. Bits a => a -> Int -> a
rotateR (b
forall a. Bits a => a
oneBits b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftL` Int
n) Int
n
{-# INLINE srl #-}

------------------------------------------------------------------------------
-- de Bruijn Multiplication Tables
------------------------------------------------------------------------------

foreign import ccall "static &debruijn_lsb64"  debruijn_lsb64  :: Ptr Word8
foreign import ccall "static &debruijn_lsb32"  debruijn_lsb32  :: Ptr Word8
-- foreign import ccall "static &debruijn_rank32" debruijn_rank32 :: Ptr Word8
foreign import ccall "static &debruijn_log32"  debruijn_log32  :: Ptr Word8

#ifndef HLINT
inlinePerformIO :: IO a -> a
inlinePerformIO :: IO a -> a
inlinePerformIO (IO State# RealWorld -> (# State# RealWorld, a #)
m) = case State# RealWorld -> (# State# RealWorld, a #)
m State# RealWorld
realWorld# of
  (# State# RealWorld
_, a
r #) -> a
r
{-# INLINE inlinePerformIO #-}
#endif