-- | Immutable one-dimensional packed bit arrays.
-- The main advantage should be compactness in memory.

module Data.BitArray 
  ( BitArray
  , bitArrayBounds
  , lookupBit
  , unsafeLookupBit
  -- * Bit array construction \/ deconstruction
  , bitArray
  , bitArray'
  , accumBitArray
  , listBitArray
  , bits
  -- * 0\/1 versions
  , bits01
  , listBitArray01
  ) 
  where

--------------------------------------------------------------------------------

import Control.Monad
import Control.Monad.ST

import Data.Bits
import Data.Word

import Data.Array.Unboxed

import Data.BitArray.Immutable
import Data.BitArray.ST

--------------------------------------------------------------------------------

instance Eq BitArray where
  BitArray
ar1 == :: BitArray -> BitArray -> Bool
== BitArray
ar2 = BitArray -> [Bool]
bits BitArray
ar1 [Bool] -> [Bool] -> Bool
forall a. Eq a => a -> a -> Bool
== BitArray -> [Bool]
bits BitArray
ar2

instance Ord BitArray where
  compare :: BitArray -> BitArray -> Ordering
compare BitArray
ar1 BitArray
ar2 = [Bool] -> [Bool] -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (BitArray -> [Bool]
bits BitArray
ar1) (BitArray -> [Bool]
bits BitArray
ar2)

instance Show BitArray where 
  show :: BitArray -> String
show ar :: BitArray
ar@(A Int
s Int
t UArray Int Word64
a) = String
"listBitArray01 " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int
s,Int
t) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [Int] -> String
forall a. Show a => a -> String
show (BitArray -> [Int]
bits01 BitArray
ar)
  
--------------------------------------------------------------------------------

bitArrayBounds :: BitArray -> (Int,Int)
bitArrayBounds :: BitArray -> (Int, Int)
bitArrayBounds (A Int
s Int
t UArray Int Word64
_) = (Int
s,Int
t)

lookupBit :: BitArray -> Int -> Bool
lookupBit :: BitArray -> Int -> Bool
lookupBit ar :: BitArray
ar@(A Int
s Int
t UArray Int Word64
_) Int
j = if Int
jInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
s Bool -> Bool -> Bool
|| Int
jInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
t 
  then String -> Bool
forall a. HasCallStack => String -> a
error String
"BitArray/lookupBit: index out of range"
  else BitArray -> Int -> Bool
unsafeLookupBit BitArray
ar Int
j
  
unsafeLookupBit :: BitArray -> Int -> Bool
unsafeLookupBit :: BitArray -> Int -> Bool
unsafeLookupBit (A Int
s Int
t UArray Int Word64
a) Int
j = Word64 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word64
w Int
l where
  (Int
k,Int
l) = Int -> (Int, Int)
ind (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
s) 
  w :: Word64
w = UArray Int Word64
aUArray Int Word64 -> Int -> Word64
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
k 

--------------------------------------------------------------------------------

-- | Unspecified values become 'False'.
bitArray :: (Int,Int) -> [(Int,Bool)] -> BitArray
bitArray :: (Int, Int) -> [(Int, Bool)] -> BitArray
bitArray = (Bool -> Bool -> Bool)
-> Bool -> (Int, Int) -> [(Int, Bool)] -> BitArray
forall a.
(Bool -> a -> Bool) -> Bool -> (Int, Int) -> [(Int, a)] -> BitArray
accumBitArray ((Bool -> Bool -> Bool) -> Bool -> Bool -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip Bool -> Bool -> Bool
forall a b. a -> b -> a
const) Bool
False 

-- | The first argument gives the default value (instead of 'False')
bitArray' :: Bool -> (Int,Int) -> [(Int,Bool)] -> BitArray
bitArray' :: Bool -> (Int, Int) -> [(Int, Bool)] -> BitArray
bitArray' = (Bool -> Bool -> Bool)
-> Bool -> (Int, Int) -> [(Int, Bool)] -> BitArray
forall a.
(Bool -> a -> Bool) -> Bool -> (Int, Int) -> [(Int, a)] -> BitArray
accumBitArray ((Bool -> Bool -> Bool) -> Bool -> Bool -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip Bool -> Bool -> Bool
forall a b. a -> b -> a
const)

{-# SPECIALIZE accumBitArray :: (Bool -> Bool -> Bool) -> Bool -> (Int,Int) -> [(Int,Bool)] -> BitArray #-}
accumBitArray :: (Bool -> a -> Bool) -> Bool -> (Int,Int) -> [(Int,a)] -> BitArray
accumBitArray :: (Bool -> a -> Bool) -> Bool -> (Int, Int) -> [(Int, a)] -> BitArray
accumBitArray Bool -> a -> Bool
f Bool
e (Int, Int)
st [(Int, a)]
xs = (forall s. ST s BitArray) -> BitArray
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s BitArray) -> BitArray)
-> (forall s. ST s BitArray) -> BitArray
forall a b. (a -> b) -> a -> b
$ do
  STBitArray s
ar <- (Int, Int) -> Bool -> ST s (STBitArray s)
forall s. (Int, Int) -> Bool -> ST s (STBitArray s)
newBitArray (Int, Int)
st Bool
e
  [(Int, a)] -> ((Int, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Int, a)]
xs (((Int, a) -> ST s ()) -> ST s ())
-> ((Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Int
i,a
x) -> do
    Bool
b <- STBitArray s -> Int -> ST s Bool
forall s. STBitArray s -> Int -> ST s Bool
readBit STBitArray s
ar Int
i
    STBitArray s -> Int -> Bool -> ST s ()
forall s. STBitArray s -> Int -> Bool -> ST s ()
writeBit STBitArray s
ar Int
i (Bool -> a -> Bool
f Bool
b a
x)
  STBitArray s -> ST s BitArray
forall s. STBitArray s -> ST s BitArray
unsafeFreezeBitArray STBitArray s
ar
    
-- | If the list is too short, the rest of the array is filled with 'False'.
listBitArray :: (Int,Int) -> [Bool] -> BitArray
listBitArray :: (Int, Int) -> [Bool] -> BitArray
listBitArray (Int
s,Int
t) [Bool]
bs = Int -> Int -> UArray Int Word64 -> BitArray
A Int
s Int
t UArray Int Word64
a where
  a :: UArray Int Word64
a = (Int, Int) -> [Word64] -> UArray Int Word64
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
listArray (Int
0,Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [Word64]
chunks
  k :: Int
k = (Int
tInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
64) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
6
  chunks :: [Word64]
chunks = Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
k ([Word64] -> [Word64]) -> [Word64] -> [Word64]
forall a b. (a -> b) -> a -> b
$ [Bool] -> [Word64]
forall a. Num a => [Bool] -> [a]
worker ([Bool]
bs [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ Bool -> [Bool]
forall a. a -> [a]
repeat Bool
False)
  worker :: [Bool] -> [a]
worker [Bool]
bs = [Bool] -> a
forall (t :: * -> *) b. (Foldable t, Num b) => t Bool -> b
convert (Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take Int
64 [Bool]
bs) a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [Bool] -> [a]
worker (Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
drop Int
64 [Bool]
bs)
  convert :: t Bool -> b
convert t Bool
bs = (b, b) -> b
forall a b. (a, b) -> a
fst ((b, b) -> b) -> (b, b) -> b
forall a b. (a -> b) -> a -> b
$ ((b, b) -> Bool -> (b, b)) -> (b, b) -> t Bool -> (b, b)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (b, b) -> Bool -> (b, b)
forall a. Num a => (a, a) -> Bool -> (a, a)
f (b
0,b
1) t Bool
bs
  f :: (a, a) -> Bool -> (a, a)
f (a
x,a
e) Bool
b = if Bool
b then (a
xa -> a -> a
forall a. Num a => a -> a -> a
+a
e, a
ea -> a -> a
forall a. Num a => a -> a -> a
+a
e) else (a
x, a
ea -> a -> a
forall a. Num a => a -> a -> a
+a
e)   

bits :: BitArray -> [Bool]
bits :: BitArray -> [Bool]
bits (A Int
s Int
t UArray Int Word64
a) = Int -> [Bool] -> [Bool]
forall a. Int -> [a] -> [a]
take (Int
tInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ([Bool] -> [Bool]) -> [Bool] -> [Bool]
forall a b. (a -> b) -> a -> b
$ (Word64 -> [Bool]) -> [Word64] -> [Bool]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Word64 -> [Bool]
forall b. (Num b, Bits b) => b -> [Bool]
worker (UArray Int Word64 -> [Word64]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
elems UArray Int Word64
a) where
  worker :: b -> [Bool]
worker b
i = ([Bool], b) -> [Bool]
forall a b. (a, b) -> a
fst (([Bool], b) -> [Bool]) -> ([Bool], b) -> [Bool]
forall a b. (a -> b) -> a -> b
$ (([Bool], b) -> Int -> ([Bool], b))
-> ([Bool], b) -> [Int] -> ([Bool], b)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ([Bool], b) -> Int -> ([Bool], b)
forall b p. (Num b, Bits b) => ([Bool], b) -> p -> ([Bool], b)
f ([], b
i) [(Int
0::Int)..Int
63]
  f :: ([Bool], b) -> p -> ([Bool], b)
f ([Bool]
bs,b
i) p
_ = ( (b
0 b -> b -> Bool
forall a. Eq a => a -> a -> Bool
/= b
i b -> b -> b
forall a. Bits a => a -> a -> a
.&. b
0x8000000000000000) Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
: [Bool]
bs, b -> Int -> b
forall a. Bits a => a -> Int -> a
shiftL b
i Int
1)

--------------------------------------------------------------------------------

listBitArray01 :: (Int,Int) -> [Int] -> BitArray
listBitArray01 :: (Int, Int) -> [Int] -> BitArray
listBitArray01 (Int, Int)
st [Int]
is = (Int, Int) -> [Bool] -> BitArray
listBitArray (Int, Int)
st ((Int -> Bool) -> [Int] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Bool
forall a. Integral a => a -> Bool
intToBool [Int]
is)

bits01 :: BitArray -> [Int]
bits01 :: BitArray -> [Int]
bits01 = (Bool -> Int) -> [Bool] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Bool -> Int
forall a. Integral a => Bool -> a
boolToInt ([Bool] -> [Int]) -> (BitArray -> [Bool]) -> BitArray -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitArray -> [Bool]
bits
 
--------------------------------------------------------------------------------