-- |Some extra functionality for Bits instances, specifically an optimized version of 'findBits' on 'Natural' (using 'countTrailingZeros').
-- Mainly exposed for testing purposes, but may be useful on its own.
{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE MagicHash #-}

module Text.Regex.Anagram.Bits
  where

import           Data.Bits
import           Data.Int
import           Data.Word
import           Numeric.Natural (Natural)

#ifdef VERSION_ghc_prim
import qualified GHC.Integer.GMP.Internals as GHC
import qualified GHC.Natural as GHC
import qualified GHC.Types as GHC
#include "MachDeps.h"
#endif

-- |value with all lower n bits set
allBits :: (Enum b, Bits b) => Int -> b
allBits :: Int -> b
allBits = b -> b
forall a. Enum a => a -> a
pred (b -> b) -> (Int -> b) -> Int -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> b
forall a. Bits a => Int -> a
bit

class Bits b => FindBits b where
  -- |list of all set bits in a value
  findBits :: b -> [Int]
  default findBits :: FiniteBits b => b -> [Int]
  findBits b
w
    | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== b -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize b
w = []
    | Bool
otherwise = Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: b -> [Int]
forall b. FindBits b => b -> [Int]
findBits (b -> Int -> b
forall a. Bits a => a -> Int -> a
clearBit b
w Int
i)
    where i :: Int
i = b -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros b
w

instance FindBits Int
instance FindBits Int8
instance FindBits Int16
instance FindBits Int32
instance FindBits Int64
instance FindBits Word
instance FindBits Word8
instance FindBits Word16
instance FindBits Word32
instance FindBits Word64

instance FindBits Natural where
#ifdef VERSION_ghc_prim
  findBits :: Natural -> [Int]
findBits (GHC.NatS# GmpLimb#
w) = Word -> [Int]
forall b. FindBits b => b -> [Int]
findBits (GmpLimb# -> Word
GHC.W# GmpLimb#
w)
  findBits (GHC.NatJ# BigNat
b) = (Int -> [Int]) -> [Int] -> [Int]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Int -> [Int]
fbi [Int
0..Int -> Int
forall a. Enum a => a -> a
pred (Int# -> Int
GHC.I# (BigNat -> Int#
GHC.sizeofBigNat# BigNat
b))]
    where
    fbi :: Int -> [Int]
fbi i :: Int
i@(GHC.I# Int#
i#) = let w :: Word
w = GmpLimb# -> Word
GHC.W# (BigNat -> Int# -> GmpLimb#
GHC.indexBigNat# BigNat
b Int#
i#) in (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
*Word -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize Word
wInt -> Int -> Int
forall a. Num a => a -> a -> a
+) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Word -> [Int]
forall b. FindBits b => b -> [Int]
findBits Word
w
#else
  findBits = fb 0 . unBitVec where
    fb i x
      | x == B.zeroBits = []
      | B.testBit x i = i : fb (succ i) (B.clearBit x i)
      | otherwise = fb (succ i) x
#endif