-- | 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 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 = 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 " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
s,Int
t) forall a. [a] -> [a] -> [a]
++ String
" " forall a. [a] -> [a] -> [a]
++ 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
jforall a. Ord a => a -> a -> Bool
<Int
s Bool -> Bool -> Bool
|| Int
jforall a. Ord a => a -> a -> Bool
>Int
t 
  then 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 = forall a. Bits a => a -> Int -> Bool
testBit Word64
w Int
l where
  (Int
k,Int
l) = Int -> (Int, Int)
ind (Int
jforall a. Num a => a -> a -> a
-Int
s) 
  w :: Word64
w = UArray Int Word64
aforall (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 = forall a.
(Bool -> a -> Bool) -> Bool -> (Int, Int) -> [(Int, a)] -> BitArray
accumBitArray (forall a b c. (a -> b -> c) -> b -> a -> c
flip 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' = forall a.
(Bool -> a -> Bool) -> Bool -> (Int, Int) -> [(Int, a)] -> BitArray
accumBitArray (forall a b c. (a -> b -> c) -> b -> a -> c
flip 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 :: forall a.
(Bool -> a -> Bool) -> Bool -> (Int, Int) -> [(Int, a)] -> BitArray
accumBitArray Bool -> a -> Bool
f Bool
e (Int, Int)
st [(Int, a)]
xs = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  STBitArray s
ar <- forall s. (Int, Int) -> Bool -> ST s (STBitArray s)
newBitArray (Int, Int)
st Bool
e
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Int, a)]
xs forall a b. (a -> b) -> a -> b
$ \(Int
i,a
x) -> do
    Bool
b <- forall s. STBitArray s -> Int -> ST s Bool
readBit STBitArray s
ar Int
i
    forall s. STBitArray s -> Int -> Bool -> ST s ()
writeBit STBitArray s
ar Int
i (Bool -> a -> Bool
f Bool
b a
x)
  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 = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
listArray (Int
0,Int
kforall a. Num a => a -> a -> a
-Int
1) [Word64]
chunks
  k :: Int
k = (Int
tforall a. Num a => a -> a -> a
-Int
sforall a. Num a => a -> a -> a
+Int
64) forall a. Bits a => a -> Int -> a
`shiftR` Int
6
  chunks :: [Word64]
chunks = forall a. Int -> [a] -> [a]
take Int
k forall a b. (a -> b) -> a -> b
$ forall {a}. Num a => [Bool] -> [a]
worker ([Bool]
bs forall a. [a] -> [a] -> [a]
++ forall a. a -> [a]
repeat Bool
False)
  worker :: [Bool] -> [a]
worker [Bool]
bs = forall {t :: * -> *} {b}. (Foldable t, Num b) => t Bool -> b
convert (forall a. Int -> [a] -> [a]
take Int
64 [Bool]
bs) forall a. a -> [a] -> [a]
: [Bool] -> [a]
worker (forall a. Int -> [a] -> [a]
drop Int
64 [Bool]
bs)
  convert :: t Bool -> b
convert t Bool
bs = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall {b}. Num b => (b, b) -> Bool -> (b, b)
f (b
0,b
1) t Bool
bs
  f :: (b, b) -> Bool -> (b, b)
f (b
x,b
e) Bool
b = if Bool
b then (b
xforall a. Num a => a -> a -> a
+b
e, b
eforall a. Num a => a -> a -> a
+b
e) else (b
x, b
eforall a. Num a => a -> a -> a
+b
e)   

bits :: BitArray -> [Bool]
bits :: BitArray -> [Bool]
bits (A Int
s Int
t UArray Int Word64
a) = forall a. Int -> [a] -> [a]
take (Int
tforall a. Num a => a -> a -> a
-Int
sforall a. Num a => a -> a -> a
+Int
1) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall {b}. (Num b, Bits b) => b -> [Bool]
worker (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 = forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl 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 forall a. Eq a => a -> a -> Bool
/= b
i forall a. Bits a => a -> a -> a
.&. b
0x8000000000000000) forall a. a -> [a] -> [a]
: [Bool]
bs, 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 (forall a b. (a -> b) -> [a] -> [b]
map forall a. Integral a => a -> Bool
intToBool [Int]
is)

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