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

module HaskellWorks.Data.EliasFano
  ( EliasFano(..)
  , fromWord64s
  , toWord64s
  , empty
  , size
  ) where

import Control.DeepSeq
import Data.Word
import GHC.Generics
import HaskellWorks.Data.AtIndex                 hiding (end)
import HaskellWorks.Data.Bits.BitWise
import HaskellWorks.Data.Bits.Log2
import HaskellWorks.Data.EliasFano.Internal
import HaskellWorks.Data.Positioning
import HaskellWorks.Data.RankSelect.Base.Select1
import HaskellWorks.Data.RankSelect.CsPoppy
import Prelude                                   hiding (length, take)

import qualified Data.Vector.Storable                          as DVS
import qualified HaskellWorks.Data.PackedVector.PackedVector64 as PV

data EliasFano = EliasFano
  { EliasFano -> CsPoppy
efBucketBits :: !CsPoppy              -- 1 marks bucket, 0 marks skip to next
  , EliasFano -> PackedVector64
efLoSegments :: !PV.PackedVector64    -- Lower segment of each entry
  , EliasFano -> Count
efLoBitCount :: !Count                -- Number of bits in each lower segment
  , EliasFano -> Count
efCount      :: !Count                -- Number of entries
  } deriving (EliasFano -> EliasFano -> Bool
(EliasFano -> EliasFano -> Bool)
-> (EliasFano -> EliasFano -> Bool) -> Eq EliasFano
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: EliasFano -> EliasFano -> Bool
$c/= :: EliasFano -> EliasFano -> Bool
== :: EliasFano -> EliasFano -> Bool
$c== :: EliasFano -> EliasFano -> Bool
Eq, Int -> EliasFano -> ShowS
[EliasFano] -> ShowS
EliasFano -> String
(Int -> EliasFano -> ShowS)
-> (EliasFano -> String)
-> ([EliasFano] -> ShowS)
-> Show EliasFano
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [EliasFano] -> ShowS
$cshowList :: [EliasFano] -> ShowS
show :: EliasFano -> String
$cshow :: EliasFano -> String
showsPrec :: Int -> EliasFano -> ShowS
$cshowsPrec :: Int -> EliasFano -> ShowS
Show, (forall x. EliasFano -> Rep EliasFano x)
-> (forall x. Rep EliasFano x -> EliasFano) -> Generic EliasFano
forall x. Rep EliasFano x -> EliasFano
forall x. EliasFano -> Rep EliasFano x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep EliasFano x -> EliasFano
$cfrom :: forall x. EliasFano -> Rep EliasFano x
Generic)

instance NFData EliasFano

size :: EliasFano -> Count
size :: EliasFano -> Count
size = EliasFano -> Count
efCount

empty :: EliasFano
empty :: EliasFano
empty = EliasFano :: CsPoppy -> PackedVector64 -> Count -> Count -> EliasFano
EliasFano
  { efBucketBits :: CsPoppy
efBucketBits  = Vector Count -> CsPoppy
makeCsPoppy Vector Count
forall a. Storable a => Vector a
DVS.empty
  , efLoSegments :: PackedVector64
efLoSegments  = PackedVector64
PV.empty
  , efLoBitCount :: Count
efLoBitCount  = Count
0
  , efCount :: Count
efCount       = Count
0
  }

fromWord64s :: [Word64] -> EliasFano
fromWord64s :: [Count] -> EliasFano
fromWord64s [Count]
ws = case [Count] -> (Maybe Count, Count)
forall (t :: * -> *) a. Foldable t => t a -> (Maybe a, Count)
foldCountAndLast [Count]
ws of
  (Just Count
end', Count
_) -> EliasFano :: CsPoppy -> PackedVector64 -> Count -> Count -> EliasFano
EliasFano
    { efBucketBits :: CsPoppy
efBucketBits  = Vector Count -> CsPoppy
makeCsPoppy (Vector Count -> CsPoppy)
-> ([Count] -> Vector Count) -> [Count] -> CsPoppy
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Count] -> Vector Count
forall a. Storable a => [a] -> Vector a
DVS.fromList ([Count] -> CsPoppy) -> [Count] -> CsPoppy
forall a b. (a -> b) -> a -> b
$ [Count] -> [Count]
hiSegmentToWords [Count]
his
    , efLoSegments :: PackedVector64
efLoSegments  = Count -> [Count] -> PackedVector64
PV.fromList Count
loBits' [Count]
los
    , efLoBitCount :: Count
efLoBitCount  = Count
loBits'
    , efCount :: Count
efCount       = Count
length'
    }
    where length' :: Count
length'   = [Count] -> Count
forall v. Length v => v -> Count
length [Count]
ws
          loBits' :: Count
loBits'   = Int -> Count
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Count -> Int
forall a. Log2 a => a -> Int
log2 (Count
end' Count -> Count -> Count
`divup` Count
length')) :: Count
          hiMask :: Count
hiMask    = Count
forall a. Bounded a => a
maxBound Count -> Count -> Count
forall a. Shift a => a -> Count -> a
.<. Count
loBits' :: Word64
          loMask :: Count
loMask    = Count -> Count
forall a. BitWise a => a -> a
comp Count
hiMask :: Word64
          his :: [Count]
his       = (Count -> Count -> Count
forall a. Shift a => a -> Count -> a
.>. Count
loBits') (Count -> Count) -> (Count -> Count) -> Count -> Count
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Count -> Count -> Count
forall a. BitWise a => a -> a -> a
.&. Count
hiMask) (Count -> Count) -> [Count] -> [Count]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Count]
ws
          los :: [Count]
los       = (Count -> Count -> Count
forall a. BitWise a => a -> a -> a
.&. Count
loMask) (Count -> Count) -> [Count] -> [Count]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Count]
ws
  (Maybe Count
Nothing, Count
_) -> EliasFano
empty

toWord64s :: EliasFano -> [Word64]
toWord64s :: EliasFano -> [Count]
toWord64s EliasFano
ef = (Count -> Count -> Count) -> (Count, Count) -> Count
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Count -> Count -> Count
combine ((Count, Count) -> Count) -> [(Count, Count)] -> [Count]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Count] -> [Count] -> [(Count, Count)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Bool] -> [Count]
bucketBitsToHiSegment [Bool]
bucketBits) (PackedVector64 -> [Count]
PV.toList (EliasFano -> PackedVector64
efLoSegments EliasFano
ef))
  where combine :: Count -> Count -> Count
combine Count
hi Count
lo = (Count
hi Count -> Count -> Count
forall a. Shift a => a -> Count -> a
.<. EliasFano -> Count
efLoBitCount EliasFano
ef) Count -> Count -> Count
forall a. BitWise a => a -> a -> a
.|. Count
lo
        bucketBits :: [Bool]
bucketBits = Count -> Vector Count -> [Bool]
bucketWordsToBucketBools (EliasFano -> Count
efCount EliasFano
ef) (CsPoppy -> Vector Count
csPoppyBits (EliasFano -> CsPoppy
efBucketBits EliasFano
ef))

instance Container EliasFano where
  type Elem EliasFano = Word64

instance Length EliasFano where
  length :: EliasFano -> Count
length = EliasFano -> Count
efCount
  {-# INLINE length #-}

instance AtIndex EliasFano where
  !!! :: EliasFano -> Position -> Elem EliasFano
(!!!)         = EliasFano -> Position -> Elem EliasFano
forall v. AtIndex v => v -> Position -> Elem v
atIndex
  atIndex :: EliasFano -> Position -> Elem EliasFano
atIndex EliasFano
ef Position
i  = PackedVector64 -> Position -> Elem PackedVector64
forall v. AtIndex v => v -> Position -> Elem v
atIndex (EliasFano -> PackedVector64
efLoSegments EliasFano
ef) Position
i Count -> Count -> Count
forall a. BitWise a => a -> a -> a
.|. ((CsPoppy -> Count -> Count
forall v. Select1 v => v -> Count -> Count
select1 (EliasFano -> CsPoppy
efBucketBits EliasFano
ef) Count
j Count -> Count -> Count
forall a. Num a => a -> a -> a
- Count
j) Count -> Count -> Count
forall a. Shift a => a -> Count -> a
.<. EliasFano -> Count
efLoBitCount EliasFano
ef)
    where j :: Count
j = Position -> Count
forall a b. (Integral a, Num b) => a -> b
fromIntegral Position
i Count -> Count -> Count
forall a. Num a => a -> a -> a
+ Count
1
  {-# INLINE (!!!)   #-}
  {-# INLINE atIndex #-}