{-|

Module      :  Data.Array.BitArray
Copyright   :  (c) Claude Heiland-Allen 2012
License     :  BSD3

Maintainer  :  claude@mathr.co.uk
Stability   :  unstable
Portability :  portable

Immutable unboxed packed bit arrays using bitwise operations to
manipulate large chunks at a time much more quickly than individually
unpacking and repacking bits would allow.

-}
-- almost all is implemented with runST and the ST-based implementation
module Data.Array.BitArray
  ( BitArray()
  -- * IArray-like interface.
  , bounds
  , array
  , listArray
  , accumArray
  , (!)
  , indices
  , elems
  , assocs
  , (//)
  , accum
  , amap
  , ixmap
  -- * Constant arrays.
  , fill
  , false
  , true
  -- * Short-circuiting reductions.
  , or
  , and
  , isUniform
  , elemIndex
  -- * Aggregate operations.
  , fold
  , map
  , zipWith
  , popCount
  -- * Bounds-checked indexing.
  , (!?)
  -- * Unsafe.
  , (!!!)
  ) where

import Prelude hiding (and, or, map, zipWith)
import qualified Prelude as P

import Control.Monad (forM_)
import Control.Monad.ST (runST)
import Data.Ix (Ix, range, inRange)

import Data.Array.BitArray.Internal (BitArray)
import qualified Data.Array.BitArray.ST as ST

-- | The bounds of an array.
{-# INLINE bounds #-}
bounds :: Ix i => BitArray i -> (i, i)
bounds :: forall i. Ix i => BitArray i -> (i, i)
bounds BitArray i
a = (forall s. ST s (i, i)) -> (i, i)
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s (i, i)
forall i s. Ix i => STBitArray s i -> ST s (i, i)
ST.getBounds (STBitArray s i -> ST s (i, i))
-> ST s (STBitArray s i) -> ST s (i, i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a)

-- | Create an array from a list of (index, element) pairs.
{-# INLINE array #-}
array :: Ix i => (i, i) {- ^ bounds -} -> [(i, Bool)] {- ^ assocs -} -> BitArray i
array :: forall i. Ix i => (i, i) -> [(i, Bool)] -> BitArray i
array (i, i)
bs [(i, Bool)]
ies = (i, i) -> BitArray i
forall i. Ix i => (i, i) -> BitArray i
false (i, i)
bs BitArray i -> [(i, Bool)] -> BitArray i
forall i. Ix i => BitArray i -> [(i, Bool)] -> BitArray i
// [(i, Bool)]
ies

-- | Create an array from a list of elements.
{-# INLINE listArray #-}
listArray :: Ix i => (i, i) {- ^ bounds -} -> [Bool] {- ^ elems -} -> BitArray i
listArray :: forall i. Ix i => (i, i) -> [Bool] -> BitArray i
listArray (i, i)
bs [Bool]
es = (forall s. ST s (BitArray i)) -> BitArray i
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s (BitArray i)
forall i s. Ix i => STBitArray s i -> ST s (BitArray i)
ST.unsafeFreeze (STBitArray s i -> ST s (BitArray i))
-> ST s (STBitArray s i) -> ST s (BitArray i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (i, i) -> [Bool] -> ST s (STBitArray s i)
forall i s. Ix i => (i, i) -> [Bool] -> ST s (STBitArray s i)
ST.newListArray (i, i)
bs [Bool]
es)

-- | Create an array by accumulating a list of (index, operand) pairs
--   from a default seed with an operation.
{-# INLINE accumArray #-}
accumArray :: Ix i => (Bool -> a -> Bool) {- ^ operation -} -> Bool {- ^ default -} -> (i, i) {- ^ bounds -} -> [(i, a)] {- ^ assocs -} -> BitArray i
accumArray :: forall i a.
Ix i =>
(Bool -> a -> Bool) -> Bool -> (i, i) -> [(i, a)] -> BitArray i
accumArray Bool -> a -> Bool
f Bool
d (i, i)
bs = (Bool -> a -> Bool) -> BitArray i -> [(i, a)] -> BitArray i
forall i a.
Ix i =>
(Bool -> a -> Bool) -> BitArray i -> [(i, a)] -> BitArray i
accum Bool -> a -> Bool
f ((i, i) -> Bool -> BitArray i
forall i. Ix i => (i, i) -> Bool -> BitArray i
fill (i, i)
bs Bool
d)

-- | Bit array indexing.
{-# INLINE (!) #-}
(!) :: Ix i => BitArray i -> i -> Bool
BitArray i
a ! :: forall i. Ix i => BitArray i -> i -> Bool
! i
i = (forall s. ST s Bool) -> Bool
forall a. (forall s. ST s a) -> a
runST (do
  STBitArray s i
a' <- BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a
  STBitArray s i -> i -> ST s Bool
forall i s. Ix i => STBitArray s i -> i -> ST s Bool
ST.readArray STBitArray s i
a' i
i)

-- | Bit array indexing without bounds checking.  Unsafe.
{-# INLINE (!!!) #-}
(!!!) :: Ix i => BitArray i -> i -> Bool
BitArray i
a !!! :: forall i. Ix i => BitArray i -> i -> Bool
!!! i
i = (forall s. ST s Bool) -> Bool
forall a. (forall s. ST s a) -> a
runST (do
  STBitArray s i
a' <- BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a
  STBitArray s i -> i -> ST s Bool
forall i s. Ix i => STBitArray s i -> i -> ST s Bool
ST.unsafeReadArray STBitArray s i
a' i
i)

-- | A list of all the valid indices for this array.
{-# INLINE indices #-}
indices :: Ix i => BitArray i -> [i]
indices :: forall i. Ix i => BitArray i -> [i]
indices = (i, i) -> [i]
forall a. Ix a => (a, a) -> [a]
range ((i, i) -> [i]) -> (BitArray i -> (i, i)) -> BitArray i -> [i]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitArray i -> (i, i)
forall i. Ix i => BitArray i -> (i, i)
bounds

-- | A list of the elements in this array.
{-# INLINE elems #-}
elems :: Ix i => BitArray i -> [Bool]
elems :: forall i. Ix i => BitArray i -> [Bool]
elems BitArray i
a = (forall s. ST s [Bool]) -> [Bool]
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s [Bool]
forall i s. Ix i => STBitArray s i -> ST s [Bool]
ST.unsafeGetElems (STBitArray s i -> ST s [Bool])
-> ST s (STBitArray s i) -> ST s [Bool]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a)
  -- P.map (a !!!) (indices a) -- very slow!

-- | A list of the (index, element) pairs in this array.
{-# INLINE assocs #-}
assocs :: Ix i => BitArray i -> [(i, Bool)]
assocs :: forall i. Ix i => BitArray i -> [(i, Bool)]
assocs BitArray i
ba = (i -> (i, Bool)) -> [i] -> [(i, Bool)]
forall a b. (a -> b) -> [a] -> [b]
P.map (\i
i -> (i
i, BitArray i
ba BitArray i -> i -> Bool
forall i. Ix i => BitArray i -> i -> Bool
! i
i)) (BitArray i -> [i]
forall i. Ix i => BitArray i -> [i]
indices BitArray i
ba)

-- | A new array with updated values at the supplied indices.
{-# INLINE (//) #-}
(//) :: Ix i => BitArray i -> [(i, Bool)] {- ^ new assocs -} -> BitArray i
BitArray i
ba // :: forall i. Ix i => BitArray i -> [(i, Bool)] -> BitArray i
// [(i, Bool)]
ies = (Bool -> Bool -> Bool) -> BitArray i -> [(i, Bool)] -> BitArray i
forall i a.
Ix i =>
(Bool -> a -> Bool) -> BitArray i -> [(i, a)] -> BitArray i
accum (\Bool
_ Bool
a -> Bool
a) BitArray i
ba [(i, Bool)]
ies

-- | Accumulate with an operation and a list of (index, operand).
{-# INLINE accum #-}
accum :: Ix i => (Bool -> a -> Bool) {- ^ operation -} -> BitArray i {- ^ source -} -> [(i, a)] {- ^ assocs -} -> BitArray i
accum :: forall i a.
Ix i =>
(Bool -> a -> Bool) -> BitArray i -> [(i, a)] -> BitArray i
accum Bool -> a -> Bool
f BitArray i
a [(i, a)]
ies = (forall s. ST s (BitArray i)) -> BitArray i
forall a. (forall s. ST s a) -> a
runST (do
  STBitArray s i
a' <- BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.thaw BitArray i
a
  [(i, a)] -> ((i, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(i, a)]
ies (((i, a) -> ST s ()) -> ST s ()) -> ((i, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(i
i, a
x) -> do
    Bool
b <- STBitArray s i -> i -> ST s Bool
forall i s. Ix i => STBitArray s i -> i -> ST s Bool
ST.readArray STBitArray s i
a' i
i
    STBitArray s i -> i -> Bool -> ST s ()
forall i s. Ix i => STBitArray s i -> i -> Bool -> ST s ()
ST.writeArray STBitArray s i
a' i
i (Bool -> a -> Bool
f Bool
b a
x)
  STBitArray s i -> ST s (BitArray i)
forall i s. Ix i => STBitArray s i -> ST s (BitArray i)
ST.unsafeFreeze STBitArray s i
a')

-- | Alias for 'map'.
{-# INLINE amap #-}
amap :: Ix i => (Bool -> Bool) -> BitArray i -> BitArray i
amap :: forall i. Ix i => (Bool -> Bool) -> BitArray i -> BitArray i
amap = (Bool -> Bool) -> BitArray i -> BitArray i
forall i. Ix i => (Bool -> Bool) -> BitArray i -> BitArray i
map

-- | Create a new array by mapping indices into a source array..
{-# INLINE ixmap #-}
ixmap :: (Ix i, Ix j) => (i, i) {- ^ new bounds -} -> (i -> j) {- ^ index transformation -} -> BitArray j {- ^ source array -} -> BitArray i
ixmap :: forall i j.
(Ix i, Ix j) =>
(i, i) -> (i -> j) -> BitArray j -> BitArray i
ixmap (i, i)
bs i -> j
h BitArray j
ba = (i, i) -> [(i, Bool)] -> BitArray i
forall i. Ix i => (i, i) -> [(i, Bool)] -> BitArray i
array (i, i)
bs ((i -> (i, Bool)) -> [i] -> [(i, Bool)]
forall a b. (a -> b) -> [a] -> [b]
P.map (\i
i -> (i
i, BitArray j
ba BitArray j -> j -> Bool
forall i. Ix i => BitArray i -> i -> Bool
! i -> j
h i
i)) ((i, i) -> [i]
forall a. Ix a => (a, a) -> [a]
range (i, i)
bs))

-- | A uniform array of bits.
{-# INLINE fill #-}
fill :: Ix i => (i, i) {- ^ bounds -} -> Bool -> BitArray i
fill :: forall i. Ix i => (i, i) -> Bool -> BitArray i
fill (i, i)
bs Bool
b = (forall s. ST s (BitArray i)) -> BitArray i
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s (BitArray i)
forall i s. Ix i => STBitArray s i -> ST s (BitArray i)
ST.unsafeFreeze (STBitArray s i -> ST s (BitArray i))
-> ST s (STBitArray s i) -> ST s (BitArray i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (i, i) -> Bool -> ST s (STBitArray s i)
forall i s. Ix i => (i, i) -> Bool -> ST s (STBitArray s i)
ST.newArray (i, i)
bs Bool
b)

-- | A uniform array of 'False'.
{-# INLINE false #-}
false :: Ix i => (i, i) {- ^ bounds -} -> BitArray i
false :: forall i. Ix i => (i, i) -> BitArray i
false (i, i)
bs = (i, i) -> Bool -> BitArray i
forall i. Ix i => (i, i) -> Bool -> BitArray i
fill (i, i)
bs Bool
False

-- | A uniform array of 'True'.
{-# INLINE true #-}
true :: Ix i => (i, i) {- ^ bounds -} -> BitArray i
true :: forall i. Ix i => (i, i) -> BitArray i
true (i, i)
bs = (i, i) -> Bool -> BitArray i
forall i. Ix i => (i, i) -> Bool -> BitArray i
fill (i, i)
bs Bool
True

-- | Bounds checking combined with array indexing.
{-# INLINE (!?) #-}
(!?) :: Ix i => BitArray i -> i -> Maybe Bool
BitArray i
b !? :: forall i. Ix i => BitArray i -> i -> Maybe Bool
!? i
i
  | (i, i) -> i -> Bool
forall a. Ix a => (a, a) -> a -> Bool
inRange (BitArray i -> (i, i)
forall i. Ix i => BitArray i -> (i, i)
bounds BitArray i
b) i
i = Bool -> Maybe Bool
forall a. a -> Maybe a
Just (BitArray i
b BitArray i -> i -> Bool
forall i. Ix i => BitArray i -> i -> Bool
! i
i)
  | Bool
otherwise = Maybe Bool
forall a. Maybe a
Nothing

-- | Short-circuit bitwise reduction: True if any bit is True.
{-# INLINE or #-}
or :: Ix i => BitArray i -> Bool
or :: forall i. Ix i => BitArray i -> Bool
or BitArray i
a = (forall s. ST s Bool) -> Bool
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s Bool
forall i s. Ix i => STBitArray s i -> ST s Bool
ST.or (STBitArray s i -> ST s Bool) -> ST s (STBitArray s i) -> ST s Bool
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a)

-- | Short-circuit bitwise reduction: False if any bit is False.
{-# INLINE and #-}
and :: Ix i => BitArray i -> Bool
and :: forall i. Ix i => BitArray i -> Bool
and BitArray i
a = (forall s. ST s Bool) -> Bool
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s Bool
forall i s. Ix i => STBitArray s i -> ST s Bool
ST.and (STBitArray s i -> ST s Bool) -> ST s (STBitArray s i) -> ST s Bool
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a)

-- | Short-circuit bitwise reduction: Nothing if any bits differ.
{-# INLINE isUniform #-}
isUniform :: Ix i => BitArray i -> Maybe Bool
isUniform :: forall i. Ix i => BitArray i -> Maybe Bool
isUniform BitArray i
a = (forall s. ST s (Maybe Bool)) -> Maybe Bool
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s (Maybe Bool)
forall i s. Ix i => STBitArray s i -> ST s (Maybe Bool)
ST.isUniform (STBitArray s i -> ST s (Maybe Bool))
-> ST s (STBitArray s i) -> ST s (Maybe Bool)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a)

-- | Look up index of first matching bit.
--
--   Note that the index type is limited to Int because there
--   is no 'unindex' method in the 'Ix' class.
{-# INLINE elemIndex #-}
elemIndex :: Bool -> BitArray Int -> Maybe Int
elemIndex :: Bool -> BitArray Int -> Maybe Int
elemIndex Bool
b BitArray Int
a = (forall s. ST s (Maybe Int)) -> Maybe Int
forall a. (forall s. ST s a) -> a
runST (Bool -> STBitArray s Int -> ST s (Maybe Int)
forall s. Bool -> STBitArray s Int -> ST s (Maybe Int)
ST.elemIndex Bool
b (STBitArray s Int -> ST s (Maybe Int))
-> ST s (STBitArray s Int) -> ST s (Maybe Int)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray Int -> ST s (STBitArray s Int)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray Int
a)

-- | Bitwise reduction with an associative commutative boolean operator.
--   Implementation lifts from 'Bool' to 'Bits' and folds large chunks
--   at a time.  Each bit is used as a source exactly once.
{-# INLINE fold #-}
fold :: Ix i => (Bool -> Bool -> Bool) -> BitArray i -> Maybe Bool
fold :: forall i.
Ix i =>
(Bool -> Bool -> Bool) -> BitArray i -> Maybe Bool
fold Bool -> Bool -> Bool
f BitArray i
a = (forall s. ST s (Maybe Bool)) -> Maybe Bool
forall a. (forall s. ST s a) -> a
runST ((Bool -> Bool -> Bool) -> STBitArray s i -> ST s (Maybe Bool)
forall i s.
Ix i =>
(Bool -> Bool -> Bool) -> STBitArray s i -> ST s (Maybe Bool)
ST.fold Bool -> Bool -> Bool
f (STBitArray s i -> ST s (Maybe Bool))
-> ST s (STBitArray s i) -> ST s (Maybe Bool)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a)

-- | Bitwise map.  Implementation lifts from 'Bool' to 'Bits' and maps
--   large chunks at a time.
{-# INLINE map #-}
map :: Ix i => (Bool -> Bool) -> BitArray i -> BitArray i
map :: forall i. Ix i => (Bool -> Bool) -> BitArray i -> BitArray i
map Bool -> Bool
f BitArray i
a = (forall s. ST s (BitArray i)) -> BitArray i
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s (BitArray i)
forall i s. Ix i => STBitArray s i -> ST s (BitArray i)
ST.unsafeFreeze (STBitArray s i -> ST s (BitArray i))
-> ST s (STBitArray s i) -> ST s (BitArray i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Bool -> Bool) -> STBitArray s i -> ST s (STBitArray s i)
forall i s.
Ix i =>
(Bool -> Bool) -> STBitArray s i -> ST s (STBitArray s i)
ST.map Bool -> Bool
f (STBitArray s i -> ST s (STBitArray s i))
-> ST s (STBitArray s i) -> ST s (STBitArray s i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a)

-- | Bitwise zipWith.  Implementation lifts from 'Bool' to 'Bits' and
--   combines large chunks at a time.
--
--   The bounds of the source arrays must be identical.
{-# INLINE zipWith #-}
zipWith :: Ix i => (Bool -> Bool -> Bool) -> BitArray i -> BitArray i -> BitArray i
zipWith :: forall i.
Ix i =>
(Bool -> Bool -> Bool) -> BitArray i -> BitArray i -> BitArray i
zipWith Bool -> Bool -> Bool
f BitArray i
a BitArray i
b
  | BitArray i -> (i, i)
forall i. Ix i => BitArray i -> (i, i)
bounds BitArray i
a (i, i) -> (i, i) -> Bool
forall a. Eq a => a -> a -> Bool
== BitArray i -> (i, i)
forall i. Ix i => BitArray i -> (i, i)
bounds BitArray i
b = (forall s. ST s (BitArray i)) -> BitArray i
forall a. (forall s. ST s a) -> a
runST (do
      STBitArray s i
a' <- BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a
      STBitArray s i
b' <- BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
b
      STBitArray s i -> ST s (BitArray i)
forall i s. Ix i => STBitArray s i -> ST s (BitArray i)
ST.unsafeFreeze (STBitArray s i -> ST s (BitArray i))
-> ST s (STBitArray s i) -> ST s (BitArray i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Bool -> Bool -> Bool)
-> STBitArray s i -> STBitArray s i -> ST s (STBitArray s i)
forall i s.
Ix i =>
(Bool -> Bool -> Bool)
-> STBitArray s i -> STBitArray s i -> ST s (STBitArray s i)
ST.zipWith Bool -> Bool -> Bool
f STBitArray s i
a' STBitArray s i
b')
  | Bool
otherwise = [Char] -> BitArray i
forall a. HasCallStack => [Char] -> a
error [Char]
"zipWith bounds mismatch"

-- | Count set bits.
{-# INLINE popCount #-}
popCount :: Ix i => BitArray i -> Int
popCount :: forall i. Ix i => BitArray i -> Int
popCount BitArray i
a = (forall s. ST s Int) -> Int
forall a. (forall s. ST s a) -> a
runST (STBitArray s i -> ST s Int
forall i s. Ix i => STBitArray s i -> ST s Int
ST.popCount (STBitArray s i -> ST s Int) -> ST s (STBitArray s i) -> ST s Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BitArray i -> ST s (STBitArray s i)
forall i s. Ix i => BitArray i -> ST s (STBitArray s i)
ST.unsafeThaw BitArray i
a)