module Data.Bits.Ordered
( lsbZ
, nextActiveZ
, maybeNextActive
, maybeLsb
, popPermutation
, popComplement
, popCntSorted
, popCntMemoInt
, popCntMemoWord
, popShiftL
, popShiftR
, activeBitsL
, activeBitsS
, activeBitsV
, subseqBit
, subsequencesBitsL
, subsequencesBitsLslow
) where
import Control.Arrow
import Data.Bits
import Data.Bits.Extras
import Data.List (subsequences,foldl',unfoldr)
import Data.Ord (comparing)
import Data.Vector.Fusion.Bundle.Size
import Data.Vector.Fusion.Util
import Data.Vector.Unboxed (Unbox)
import Data.Word(Word(..))
import qualified Data.Vector.Algorithms.Intro as AI
import qualified Data.Vector.Fusion.Bundle.Monadic as BM
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU
captureNull ∷ Ranked t ⇒ t → (t → Int) → Int
{-# Inline captureNull #-}
captureNull t f = if t==0 then -1 else f t
lsbZ ∷ Ranked t ⇒ t → Int
{-# Inline lsbZ #-}
lsbZ t = captureNull t lsb
nextActiveZ :: Ranked t => Int -> t -> Int
nextActiveZ k t = lsbZ $ (t `shiftR` (k+1)) `shiftL` (k+1)
{-# Inline nextActiveZ #-}
maybeNextActive :: Ranked t => Int -> t -> Maybe Int
maybeNextActive k t = if t'==0 then Nothing else Just (lsb t')
where t' = (t `shiftR` (k+1) `shiftL` (k+1))
{-# Inline maybeNextActive #-}
maybeLsb :: Ranked t => t -> Maybe Int
maybeLsb t = if t==0 then Nothing else Just (lsb t)
{-# Inline maybeLsb #-}
activeBitsL :: Ranked t => t -> [Int]
activeBitsL = unId . SM.toList . activeBitsS
{-# Inline activeBitsL #-}
activeBitsV :: (Ranked t, VG.Vector v Int) => t -> v Int
activeBitsV = VG.unstream . flip BM.fromStream Unknown . activeBitsS
{-# Inline activeBitsV #-}
activeBitsS :: (Ranked t, Monad m) => t -> SM.Stream m Int
activeBitsS t = SM.unfoldr (fmap (id &&& (`maybeNextActive` t))) (maybeLsb t)
{-# Inline activeBitsS #-}
popCntSorted :: (Unbox n, Integral n, Bits n, Ranked n) => Int -> VU.Vector n
popCntSorted n = VU.modify (AI.sortBy (comparing (popCount &&& activeBitsL))) $ VU.enumFromN 0 (2^n)
{-# Inline popCntSorted #-}
popCntMemoInt
:: Int
-> VU.Vector Int
popCntMemoInt n
| n>limit = error $ "for safety reasons, memoization is only performed for popcounts up to " ++ show limit ++ " bits, memoize manually!"
| otherwise = _popCntMemoInt !! n
where limit = 28
{-# Inline popCntMemoInt #-}
_popCntMemoInt = map popCntSorted [0..]
{-# NoInline _popCntMemoInt #-}
popCntMemoWord
:: Int
-> VU.Vector Word
popCntMemoWord n
| n>limit = error $ "for safety reasons, memoization is only performed for popcounts up to " ++ show limit ++ " bits, memoize manually!"
| otherwise = _popCntMemoWord !! n
where limit = 28
{-# Inline popCntMemoWord #-}
_popCntMemoWord = map popCntSorted [0..]
{-# NoInline _popCntMemoWord #-}
popPermutation
:: Ranked t
=> Int
-> t
-> Maybe t
popPermutation !h' !s'
| popCount s' < 1 || h' < 2 = Nothing
| Just k <- findK (h' -2)
, Just l <- findL k (h' -1)
= let swp = setBit (clearBit s' k) l
in Just $ reverseFrom (k+1) (h' -1) swp swp
| otherwise = Nothing
where findK k
| k < 0 = Nothing
| testBit s' k && not (testBit s' (k+1)) = Just k
| otherwise = findK (k-1)
findL k l
| l <= k = Nothing
| not $ testBit s' l = Just l
| otherwise = findL k $ l-1
reverseFrom u d src tgt
| u >= h' = tgt
| otherwise = reverseFrom (u+1) (d-1) src (assignBit (assignBit tgt u (testBit src d)) d (testBit src u))
{-# Inline popPermutation #-}
popComplement
:: Ranked t
=> Int
-> t
-> t
popComplement !h !s = mask .&. complement s
where mask = (2^h -1)
{-# Inline popComplement #-}
popShiftL
:: (Ranked t)
=> t
-> t
-> t
popShiftL mask lsp = go 0 0 mask lsp where
go !acc !(k::Int) !m !l
| l==0 || m==0 = acc
| testBit m 0
, testBit l 0 = go (acc + bit k) (k+1) (unsafeShiftR m 1) (unsafeShiftR l 1)
| not $ testBit m 0 = go acc (k+1) (unsafeShiftR m 1) l
| not $ testBit l 0 = go acc (k+1) (unsafeShiftR m 1) (unsafeShiftR l 1)
{-# Inline popShiftL #-}
popShiftR
:: (Ranked t)
=> t
-> t
-> t
popShiftR mask lsp = go 0 0 mask lsp where
go !acc !k !m !l
| m==0 || l==0 = acc
| testBit m 0
, testBit l 0 = go (acc .|. bit k) (k+1) (m `unsafeShiftR` 1) (l `unsafeShiftR` 1)
| testBit m 0 = go acc (k+1) (m `unsafeShiftR` 1) (l `unsafeShiftR` 1)
| otherwise = go acc k (m `unsafeShiftR` 1) (l `unsafeShiftR` 1)
{-# Inline popShiftR #-}
subseqBit
∷ (Ord t, Ranked t)
⇒ t
→ t
→ Maybe (t,t)
{-# Inline subseqBit #-}
subseqBit mask cur
| cur > limit = Nothing
| otherwise = Just (popShiftL mask cur,cur+1)
where !limit = 2 ^ popCount mask - 1
subsequencesBitsL ∷ (Ord t, Ranked t) ⇒ t → [t]
{-# Inline subsequencesBitsL #-}
subsequencesBitsL t =
let
in unfoldr (subseqBit t) zeroBits
subsequencesBitsLslow ∷ (Ord t, Ranked t) ⇒ t → [t]
{-# Inline subsequencesBitsLslow #-}
subsequencesBitsLslow t =
let as = activeBitsL t
xs = subsequences as
setBits = foldl' setBit zeroBits
in map setBits xs
instance Ranked Int where
#if x86_64_HOST_ARCH
lsb = lsb . w64
rank = rank . w64
nlz = nlz . w64
#endif
#if i386_HOST_ARCH
lsb = lsb . w32
rank = rank . w32
nlz = nlz . w32
#endif
{-# Inline lsb #-}
{-# Inline rank #-}
{-# Inline nlz #-}
instance Ranked Word where
#if x86_64_HOST_ARCH
lsb = lsb . w64
rank = rank . w64
nlz = nlz . w64
#endif
#if i386_HOST_ARCH
lsb = lsb . w32
rank = rank . w32
nlz = nlz . w32
#endif
{-# Inline lsb #-}
{-# Inline rank #-}
{-# Inline nlz #-}