{-# OPTIONS_GHC -Wno-orphans #-}

module PrimitiveExtras.Bitmap
  ( Bitmap (..),
    empty,
    singleton,
    insert,
    invert,
    indexList,
    boolList,
    pair,
    populatedIndex,
    isPopulated,
    population,
    null,
    bits,
    populatedIndicesList,
    int64,
    allBitsList,
    allBitsUnfoldl,
    populatedBitsUnfoldl,
    indicesAmongstPopulatedBitsUnfoldl,
  )
where

import qualified DeferredFolds.Unfoldl as Unfoldl
import PrimitiveExtras.Prelude hiding (empty, insert, null, singleton, traverse_)
import PrimitiveExtras.Types

deriving instance Eq Bitmap

{-# NOINLINE maxSize #-}
maxSize :: Int
maxSize :: Int
maxSize = Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Int
forall a. HasCallStack => a
undefined :: Int)

{-# NOINLINE maxBit #-}
maxBit :: Int
maxBit :: Int
maxBit = Int -> Int
forall a. Enum a => a -> a
pred Int
maxSize

{-# NOINLINE allBitsList #-}
allBitsList :: [Int]
allBitsList :: [Int]
allBitsList = [Int
0 .. Int
maxBit]

-- * Constructors

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

{-# INLINE empty #-}
empty :: Bitmap
empty :: Bitmap
empty = Int64 -> Bitmap
Bitmap Int64
0

{-# INLINE singleton #-}
singleton :: Int -> Bitmap
singleton :: Int -> Bitmap
singleton = Int64 -> Bitmap
Bitmap (Int64 -> Bitmap) -> (Int -> Int64) -> Int -> Bitmap
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int -> Int64
forall a. Bits a => Int -> a
bit

{-# INLINE insert #-}
insert :: Int -> Bitmap -> Bitmap
insert :: Int -> Bitmap -> Bitmap
insert Int
i = Int64 -> Bitmap
Bitmap (Int64 -> Bitmap) -> (Bitmap -> Int64) -> Bitmap -> Bitmap
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int -> Int64
forall a. Bits a => Int -> a
bit Int
i Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
.|.) (Int64 -> Int64) -> (Bitmap -> Int64) -> Bitmap -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Bitmap -> Int64
int64

{-# INLINE invert #-}
invert :: Int -> Bitmap -> Bitmap
invert :: Int -> Bitmap -> Bitmap
invert Int
i = Int64 -> Bitmap
Bitmap (Int64 -> Bitmap) -> (Bitmap -> Int64) -> Bitmap -> Bitmap
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int -> Int64
forall a. Bits a => Int -> a
bit Int
i Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
`xor`) (Int64 -> Int64) -> (Bitmap -> Int64) -> Bitmap -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Bitmap -> Int64
int64

{-# INLINE indexList #-}
indexList :: [Int] -> Bitmap
indexList :: [Int] -> Bitmap
indexList = Int64 -> Bitmap
Bitmap (Int64 -> Bitmap) -> ([Int] -> Int64) -> [Int] -> Bitmap
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int64 -> Int64 -> Int64) -> Int64 -> [Int64] -> Int64
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
(.|.) Int64
0 ([Int64] -> Int64) -> ([Int] -> [Int64]) -> [Int] -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int -> Int64) -> [Int] -> [Int64]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Int64
forall a. Bits a => Int -> a
bit

{-# INLINE boolList #-}
boolList :: [Bool] -> Bitmap
boolList :: [Bool] -> Bitmap
boolList = Int64 -> Bitmap
Bitmap (Int64 -> Bitmap) -> ([Bool] -> Int64) -> [Bool] -> Bitmap
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int64 -> Int64 -> Int64) -> Int64 -> [Int64] -> Int64
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
(.|.) Int64
0 ([Int64] -> Int64) -> ([Bool] -> [Int64]) -> [Bool] -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (Int -> Bool -> Int64) -> [Int] -> [Bool] -> [Int64]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
index -> Int64 -> Int64 -> Bool -> Int64
forall a. a -> a -> Bool -> a
bool Int64
0 (Int -> Int64
forall a. Bits a => Int -> a
bit Int
index)) [Int]
allBitsList

{-# INLINE pair #-}
pair :: Int -> Int -> Bitmap
pair :: Int -> Int -> Bitmap
pair Int
i1 Int
i2 = Int64 -> Bitmap
Bitmap (Int -> Int64
forall a. Bits a => Int -> a
bit Int
i1 Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
.|. Int -> Int64
forall a. Bits a => Int -> a
bit Int
i2)

-- * Accessors

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

-- |
-- A number of non-zero bits, preceding this one.
{-# INLINE populatedIndex #-}
populatedIndex :: Int -> Bitmap -> Int
populatedIndex :: Int -> Bitmap -> Int
populatedIndex Int
i (Bitmap Int64
int) = Int64 -> Int
forall a. Bits a => a -> Int
popCount (Int64
int Int64 -> Int64 -> Int64
forall a. Bits a => a -> a -> a
.&. (Int -> Int64
forall a. Bits a => Int -> a
bit Int
i Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))

{-# INLINE isPopulated #-}
isPopulated :: Int -> Bitmap -> Bool
isPopulated :: Int -> Bitmap -> Bool
isPopulated Int
index (Bitmap Int64
int) = Int64 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int64
int Int
index

{-# INLINE population #-}
population :: Bitmap -> Int
population :: Bitmap -> Int
population (Bitmap Int64
int) = Int64 -> Int
forall a. Bits a => a -> Int
popCount Int64
int

{-# INLINE null #-}
null :: Bitmap -> Bool
null :: Bitmap -> Bool
null (Bitmap Int64
int) = Int64
int Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64
0

{-# INLINE bits #-}
bits :: Bitmap -> [Int]
bits :: Bitmap -> [Int]
bits (Bitmap Int64
int) = (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int64 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Int64
int) [Int]
allBitsList

{-# INLINE populatedIndicesList #-}
populatedIndicesList :: Bitmap -> [Int]
populatedIndicesList :: Bitmap -> [Int]
populatedIndicesList = Int -> Int -> [Int]
forall a. Enum a => a -> a -> [a]
enumFromTo Int
0 (Int -> [Int]) -> (Bitmap -> Int) -> Bitmap -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int -> Int
forall a. Enum a => a -> a
pred (Int -> Int) -> (Bitmap -> Int) -> Bitmap -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Bitmap -> Int
population

{-# INLINE int64 #-}
int64 :: Bitmap -> Int64
int64 :: Bitmap -> Int64
int64 (Bitmap Int64
int) = Int64
int

{-# NOINLINE allBitsUnfoldl #-}
allBitsUnfoldl :: Unfoldl Int
allBitsUnfoldl :: Unfoldl Int
allBitsUnfoldl = Int -> Int -> Unfoldl Int
Unfoldl.intsInRange Int
0 Int
maxBit

populatedBitsUnfoldl :: Bitmap -> Unfoldl Int
populatedBitsUnfoldl :: Bitmap -> Unfoldl Int
populatedBitsUnfoldl Bitmap
bitmap = (Int -> Bool) -> Unfoldl Int -> Unfoldl Int
forall a. (a -> Bool) -> Unfoldl a -> Unfoldl a
Unfoldl.filter ((Int -> Bitmap -> Bool) -> Bitmap -> Int -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Bitmap -> Bool
isPopulated Bitmap
bitmap) Unfoldl Int
allBitsUnfoldl

indicesAmongstPopulatedBitsUnfoldl :: Bitmap -> Unfoldl Int
indicesAmongstPopulatedBitsUnfoldl :: Bitmap -> Unfoldl Int
indicesAmongstPopulatedBitsUnfoldl Bitmap
bitmap = Int -> Int -> Unfoldl Int
Unfoldl.intsInRange Int
0 (Int -> Int
forall a. Enum a => a -> a
pred (Bitmap -> Int
population Bitmap
bitmap))