{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies      #-}

module HaskellWorks.Data.EliasFano.Internal
  ( divup
  , hiSegmentToBucketBits
  , bucketBitsToHiSegment
  , bucketBoolsToBucketWords
  , bucketWordsToBucketBools
  , hiSegmentToWords
  , foldCountAndLast
  ) where

import Data.Int
import Data.Word
import HaskellWorks.Data.Bits.BitWise
import HaskellWorks.Data.Positioning
import Prelude                        hiding (length, take)

import qualified Data.Vector.Storable as DVS
import qualified Prelude              as P

foldCountAndLast :: Foldable t => t a -> (Maybe a, Count)
foldCountAndLast :: t a -> (Maybe a, Count)
foldCountAndLast = ((Maybe a, Count) -> a -> (Maybe a, Count))
-> (Maybe a, Count) -> t a -> (Maybe a, Count)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Maybe a, Count) -> a -> (Maybe a, Count)
forall a. (Maybe a, Count) -> a -> (Maybe a, Count)
go (Maybe a
forall a. Maybe a
Nothing, Count
0)
  where go :: (Maybe a, Count) -> a -> (Maybe a, Count)
        go :: (Maybe a, Count) -> a -> (Maybe a, Count)
go (Maybe a
_, Count
n) a
a = (a -> Maybe a
forall a. a -> Maybe a
Just a
a, Count
n Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1)
{-# INLINE foldCountAndLast #-}

-- | Calculates ceil (n / d) for small numbers
divup :: Word64 -> Word64 -> Word64
divup :: Count -> Count -> Count
divup Count
n Count
d = Int64 -> Count
forall a b. (Integral a, Num b) => a -> b
fromIntegral (-((-Int64
sn) Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`div` Int64
sd)) :: Word64
  where sd :: Int64
sd = Count -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Count
d :: Int64
        sn :: Int64
sn = Count -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Count
n :: Int64

bucketBoolsToBucketWords :: [Bool] -> DVS.Vector Word64
bucketBoolsToBucketWords :: [Bool] -> Vector Count
bucketBoolsToBucketWords [Bool]
bs = Int -> ([Bool] -> Maybe (Count, [Bool])) -> [Bool] -> Vector Count
forall a b.
Storable a =>
Int -> (b -> Maybe (a, b)) -> b -> Vector a
DVS.unfoldrN (([Bool] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
P.length [Bool]
bs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
64) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Bool] -> Maybe (Count, [Bool])
gen [Bool]
bs
  where gen :: [Bool] -> Maybe (Word64, [Bool])
        gen :: [Bool] -> Maybe (Count, [Bool])
gen [Bool]
cs = if Bool -> Bool
not ([Bool] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Bool]
cs) then [Bool] -> Count -> Count -> Maybe (Count, [Bool])
genWord [Bool]
cs Count
0 Count
0 else Maybe (Count, [Bool])
forall a. Maybe a
Nothing
        genWord :: [Bool] -> Count -> Word64 -> Maybe (Word64, [Bool])
        genWord :: [Bool] -> Count -> Count -> Maybe (Count, [Bool])
genWord (Bool
True :[Bool]
cs) Count
i Count
acc | Count
i Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
< Count
64 = [Bool] -> Count -> Count -> Maybe (Count, [Bool])
genWord [Bool]
cs (Count
i Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1) (Count
acc Count -> Count -> Count
forall a. BitWise a => a -> a -> a
.|. (Count
1 Count -> Count -> Count
forall a. Shift a => a -> Count -> a
.<. Count
i))
        genWord (Bool
False:[Bool]
cs) Count
i Count
acc | Count
i Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
< Count
64 = [Bool] -> Count -> Count -> Maybe (Count, [Bool])
genWord [Bool]
cs (Count
i Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1)  Count
acc
        genWord        [Bool]
cs  Count
_ Count
acc = (Count, [Bool]) -> Maybe (Count, [Bool])
forall a. a -> Maybe a
Just (Count
acc, [Bool]
cs)

bucketWordsToBucketBools :: Count -> DVS.Vector Word64 -> [Bool]
bucketWordsToBucketBools :: Count -> Vector Count -> [Bool]
bucketWordsToBucketBools Count
n Vector Count
v = ([Bool] -> [Bool], Count) -> [Bool] -> [Bool]
forall a b. (a, b) -> a
fst ((([Bool] -> [Bool], Count) -> Count -> ([Bool] -> [Bool], Count))
-> ([Bool] -> [Bool], Count)
-> Vector Count
-> ([Bool] -> [Bool], Count)
forall b a. Storable b => (a -> b -> a) -> a -> Vector b -> a
DVS.foldl ([Bool] -> [Bool], Count) -> Count -> ([Bool] -> [Bool], Count)
go ([Bool] -> [Bool]
forall a. a -> a
id, Count
n) Vector Count
v) []
  where go :: ([Bool] -> [Bool], Count) -> Word64 -> ([Bool] -> [Bool], Count)
        go :: ([Bool] -> [Bool], Count) -> Count -> ([Bool] -> [Bool], Count)
go ([Bool] -> [Bool]
bs, Count
c) Count
w | Count
c Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
> Count
0 = case Count -> Count -> Count -> ([Bool] -> [Bool], Count)
goWord Count
c Count
64 Count
w of
                                ([Bool] -> [Bool]
cs, Count
finalCount) -> ([Bool] -> [Bool]
bs ([Bool] -> [Bool]) -> ([Bool] -> [Bool]) -> [Bool] -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bool] -> [Bool]
cs, Count
finalCount)
        go ([Bool] -> [Bool]
bs, Count
_) Count
_         = ([Bool] -> [Bool]
bs, Count
0)
        goWord :: Count -> Count -> Word64 -> ([Bool] -> [Bool], Count)
        goWord :: Count -> Count -> Count -> ([Bool] -> [Bool], Count)
goWord Count
c Count
i Count
w | Count
c Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
> Count
0 Bool -> Bool -> Bool
&& Count
i Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
> Count
0 = let b :: Bool
b = (Count
w Count -> Count -> Count
forall a. BitWise a => a -> a -> a
.&. Count
1) Count -> Count -> Bool
forall a. Eq a => a -> a -> Bool
/= Count
0
                                        in case Count -> Count -> Count -> ([Bool] -> [Bool], Count)
goWord (if Bool
b then Count
c Count -> Count -> Count
forall a. Num a => a -> a -> a
- Count
1 else Count
c) (Count
i Count -> Count -> Count
forall a. Num a => a -> a -> a
- Count
1) (Count
w Count -> Count -> Count
forall a. Shift a => a -> Count -> a
.>. Count
1) of
                                              ([Bool] -> [Bool]
bs, Count
finalCount) -> ((Bool
bBool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:) ([Bool] -> [Bool]) -> ([Bool] -> [Bool]) -> [Bool] -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bool] -> [Bool]
bs, Count
finalCount)
        goWord Count
c Count
_ Count
_                  = ([Bool] -> [Bool]
forall a. a -> a
id, Count
c)

hiSegmentToWords :: [Word64] -> [Word64]
hiSegmentToWords :: [Count] -> [Count]
hiSegmentToWords = Count -> Count -> Count -> [Count] -> [Count]
go Count
0 Count
0 Count
0
  where go :: Count -> Word64 -> Word64 -> [Word64] -> [Word64]
        go :: Count -> Count -> Count -> [Count] -> [Count]
go Count
n Count
acc Count
lst us :: [Count]
us@(Count
v:[Count]
vs) = if Count
n Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
< Count
64
          then if Count
lst Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
< Count
v
            then Count -> Count -> Count -> [Count] -> [Count]
go (Count
n Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1)  Count
acc                (Count
lst Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1) [Count]
us
            else Count -> Count -> Count -> [Count] -> [Count]
go (Count
n Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1) (Count
acc Count -> Count -> Count
forall a. BitWise a => a -> a -> a
.|. (Count
1 Count -> Count -> Count
forall a. Shift a => a -> Count -> a
.<. Count
n))  Count
v        [Count]
vs
          else Count
accCount -> [Count] -> [Count]
forall a. a -> [a] -> [a]
:Count -> Count -> Count -> [Count] -> [Count]
go Count
0 Count
0 Count
lst [Count]
us
        go Count
0 Count
_   Count
_    [Count]
_         = []
        go Count
_ Count
acc Count
_    [Count]
_         = [Count
acc]

hiSegmentToBucketBits :: Word64 -> [Word64] -> [Bool]
hiSegmentToBucketBits :: Count -> [Count] -> [Bool]
hiSegmentToBucketBits Count
lastWord = Count -> [Count] -> [Bool]
go Count
0
  where go :: Word64 -> [Word64] -> [Bool]
        go :: Count -> [Count] -> [Bool]
go Count
i []     | Count
i Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
>= Count
lastWord = []
        go Count
i (Count
a:[Count]
as) | Count
i Count -> Count -> Bool
forall a. Eq a => a -> a -> Bool
== Count
a        = Bool
TrueBool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:Count -> [Count] -> [Bool]
go Count
i [Count]
as
        go Count
i (Count
a:[Count]
as) | Count
i Count -> Count -> Bool
forall a. Ord a => a -> a -> Bool
<  Count
a        = Bool
FalseBool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:Count -> [Count] -> [Bool]
go (Count
i Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1) (Count
aCount -> [Count] -> [Count]
forall a. a -> [a] -> [a]
:[Count]
as)
        go Count
i []     = Bool
FalseBool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:Count -> [Count] -> [Bool]
go (Count
i Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1) []
        go Count
_ (Count
_:[Count]
_)  = [Char] -> [Bool]
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid entry"

bucketBitsToHiSegment :: [Bool] -> [Word64]
bucketBitsToHiSegment :: [Bool] -> [Count]
bucketBitsToHiSegment = Count -> [Bool] -> [Count]
go Count
0
  where go :: Word64 -> [Bool] -> [Word64]
        go :: Count -> [Bool] -> [Count]
go Count
_ []          = []
        go Count
i (Bool
True:[Bool]
bs)   = Count
iCount -> [Count] -> [Count]
forall a. a -> [a] -> [a]
:Count -> [Bool] -> [Count]
go  Count
i      [Bool]
bs
        go Count
i (Bool
False: [Bool]
bs) =   Count -> [Bool] -> [Count]
go (Count
i Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1) [Bool]
bs