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

module HaskellWorks.Data.PackedVector.PackedVector64
  ( PackedVector64(..)
  , empty
  , fromList
  , fromListN
  , toList
  , createFileIndex
  ) where

import Control.DeepSeq
import Data.Int
import Data.Word
import GHC.Generics
import HaskellWorks.Data.AtIndex
import HaskellWorks.Data.Bits.BitWise
import HaskellWorks.Data.Bits.LoBitsSized
import HaskellWorks.Data.PackedVector.Internal
import HaskellWorks.Data.Positioning
import HaskellWorks.Data.Unsign
import Prelude                                 hiding (length)

import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Lazy    as LBS
import qualified Data.Vector.Storable    as DVS
import qualified System.IO               as IO

data PackedVector64 = PackedVector64
  { PackedVector64 -> Vector Word64
swBuffer    :: !(DVS.Vector Word64)
  , PackedVector64 -> Word
swBitSize   :: !Word
  , PackedVector64 -> Int
swBufferLen :: !Int
  } deriving (PackedVector64 -> PackedVector64 -> Bool
(PackedVector64 -> PackedVector64 -> Bool)
-> (PackedVector64 -> PackedVector64 -> Bool) -> Eq PackedVector64
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PackedVector64 -> PackedVector64 -> Bool
$c/= :: PackedVector64 -> PackedVector64 -> Bool
== :: PackedVector64 -> PackedVector64 -> Bool
$c== :: PackedVector64 -> PackedVector64 -> Bool
Eq, Int -> PackedVector64 -> ShowS
[PackedVector64] -> ShowS
PackedVector64 -> String
(Int -> PackedVector64 -> ShowS)
-> (PackedVector64 -> String)
-> ([PackedVector64] -> ShowS)
-> Show PackedVector64
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PackedVector64] -> ShowS
$cshowList :: [PackedVector64] -> ShowS
show :: PackedVector64 -> String
$cshow :: PackedVector64 -> String
showsPrec :: Int -> PackedVector64 -> ShowS
$cshowsPrec :: Int -> PackedVector64 -> ShowS
Show, (forall x. PackedVector64 -> Rep PackedVector64 x)
-> (forall x. Rep PackedVector64 x -> PackedVector64)
-> Generic PackedVector64
forall x. Rep PackedVector64 x -> PackedVector64
forall x. PackedVector64 -> Rep PackedVector64 x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep PackedVector64 x -> PackedVector64
$cfrom :: forall x. PackedVector64 -> Rep PackedVector64 x
Generic)

instance NFData PackedVector64

empty :: PackedVector64
empty :: PackedVector64
empty =
  PackedVector64 :: Vector Word64 -> Word -> Int -> PackedVector64
PackedVector64
  { swBuffer :: Vector Word64
swBuffer    = Vector Word64
forall a. Storable a => Vector a
DVS.empty
  , swBufferLen :: Int
swBufferLen = Int
0
  , swBitSize :: Word
swBitSize   = Word
1
  }

instance Container PackedVector64 where
  type Elem PackedVector64 = Word64

instance Length PackedVector64 where
  length :: PackedVector64 -> Word64
length = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word64)
-> (PackedVector64 -> Int) -> PackedVector64 -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PackedVector64 -> Int
swBufferLen
  {-# INLINE length #-}

instance AtIndex PackedVector64 where
  atIndex :: PackedVector64 -> Position -> Elem PackedVector64
atIndex PackedVector64
v Position
i =
    let bitSize :: Word64
bitSize     = Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PackedVector64 -> Word
swBitSize PackedVector64
v) :: Count
        bitIndex :: Position
bitIndex    = Word -> Position
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PackedVector64 -> Word
swBitSize PackedVector64
v) Position -> Position -> Position
forall a. Num a => a -> a -> a
* Position
i
        (Position
q, Position
r)      = Position
bitIndex Position -> Position -> (Position, Position)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Position
64
        vv :: Vector Word64
vv          = PackedVector64 -> Vector Word64
swBuffer PackedVector64
v
    in if Position
r Position -> Position -> Bool
forall a. Ord a => a -> a -> Bool
<= Position
64 Position -> Position -> Position
forall a. Num a => a -> a -> a
- Word64 -> Position
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
bitSize
      then -- Not crossing boundary
        ((Vector Word64
vv Vector Word64 -> Position -> Elem (Vector Word64)
forall v. AtIndex v => v -> Position -> Elem v
!!! Position
q) Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.>. Position -> UnsignOf Position
forall a. Unsign a => a -> UnsignOf a
unsign Position
r) Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. Word64 -> Word64
forall a. LoBitsSized a => Word64 -> a
loBitsSized Word64
bitSize
      else -- Crossing boundary
        let loBitsSize :: Word64
loBitsSize  = Word64
64 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Position -> Word64
forall a. ToCount a => a -> Word64
toCount Position
r
            hiBitsSize :: Word64
hiBitsSize  = Word64
bitSize Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
loBitsSize
            loBits :: Word64
loBits      = ((Vector Word64
vv Vector Word64 -> Position -> Elem (Vector Word64)
forall v. AtIndex v => v -> Position -> Elem v
!!! Position
q) Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.>. Position -> UnsignOf Position
forall a. Unsign a => a -> UnsignOf a
unsign Position
r) Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. Word64 -> Word64
forall a. LoBitsSized a => Word64 -> a
loBitsSized Word64
loBitsSize
            hiBits :: Word64
hiBits      = (Vector Word64
vv Vector Word64 -> Position -> Elem (Vector Word64)
forall v. AtIndex v => v -> Position -> Elem v
!!! (Position
q Position -> Position -> Position
forall a. Num a => a -> a -> a
+ Position
1)) Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. Word64 -> Word64
forall a. LoBitsSized a => Word64 -> a
loBitsSized Word64
hiBitsSize
        in  Word64
loBits Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. (Word64
hiBits Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. Word64
loBitsSize)
  !!! :: PackedVector64 -> Position -> Elem PackedVector64
(!!!)       = PackedVector64 -> Position -> Elem PackedVector64
forall v. AtIndex v => v -> Position -> Elem v
atIndex
  {-# INLINE (!!!)   #-}
  {-# INLINE atIndex #-}

fromList :: Count -> [Word64] -> PackedVector64
fromList :: Word64 -> [Word64] -> PackedVector64
fromList Word64
wordLength [Word64]
ws = PackedVector64 :: Vector Word64 -> Word -> Int -> PackedVector64
PackedVector64
  { swBuffer :: Vector Word64
swBuffer    = [Word64] -> Vector Word64
forall a. Storable a => [a] -> Vector a
DVS.fromList (Word64 -> [Word64] -> [Word64]
forall a. PackBits a => Word64 -> [a] -> [a]
packBits Word64
wordLength [Word64]
ws)
  , swBufferLen :: Int
swBufferLen = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word64] -> Word64
forall v. Length v => v -> Word64
length [Word64]
ws)
  , swBitSize :: Word
swBitSize   = Word64 -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
wordLength
  }

fromListN :: Count -> Count -> [Word64] -> PackedVector64
fromListN :: Word64 -> Word64 -> [Word64] -> PackedVector64
fromListN Word64
vectorSize Word64
wordLength [Word64]
ws = PackedVector64 :: Vector Word64 -> Word -> Int -> PackedVector64
PackedVector64
  { swBuffer :: Vector Word64
swBuffer    = Int -> [Word64] -> Vector Word64
forall a. Storable a => Int -> [a] -> Vector a
DVS.fromListN (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
vectorSize) (Word64 -> [Word64] -> [Word64]
forall a. PackBits a => Word64 -> [a] -> [a]
packBits Word64
wordLength [Word64]
ws)
  , swBufferLen :: Int
swBufferLen = Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word64] -> Word64
forall v. Length v => v -> Word64
length [Word64]
ws)
  , swBitSize :: Word
swBitSize   = Word64 -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
wordLength
  }

toList :: PackedVector64 -> [Word64]
toList :: PackedVector64 -> [Word64]
toList PackedVector64
v = Int -> Word64 -> [Word64] -> [Word64]
forall a. UnpackBits a => Int -> Word64 -> [a] -> [a]
unpackBits (PackedVector64 -> Int
swBufferLen PackedVector64
v) (Word -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PackedVector64 -> Word
swBitSize PackedVector64
v)) (Vector Word64 -> [Word64]
forall a. Storable a => Vector a -> [a]
DVS.toList (PackedVector64 -> Vector Word64
swBuffer PackedVector64
v))

encodePacked :: Count -> [Word64] -> B.Builder
encodePacked :: Word64 -> [Word64] -> Builder
encodePacked Word64
wordSize = Word64 -> Word64 -> [Word64] -> Builder
go Word64
0 Word64
0
  where go :: Count -> Word64 -> [Word64] -> B.Builder
        go :: Word64 -> Word64 -> [Word64] -> Builder
go Word64
0           Word64
_   [] = Builder
forall a. Monoid a => a
mempty
        go Word64
_           Word64
acc [] = Word64 -> Builder
B.word64LE Word64
acc
        go Word64
bitsWritten Word64
acc (Word64
w:[Word64]
ws) =
          let totalBits :: Word64
totalBits   = Word64
bitsWritten Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
wordSize
              excessBits :: Word64
excessBits  = Word64
totalBits Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
64
              newAcc :: Word64
newAcc      = (Word64
w Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. Word64
bitsWritten) Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. Word64
acc
          in if Word64
totalBits Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
64
            then Word64 -> Builder
B.word64LE Word64
newAcc Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<>             Word64 -> Word64 -> [Word64] -> Builder
go Word64
excessBits (Word64
w Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.>. (Word64
wordSize Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
excessBits))     [Word64]
ws
            else let freeBits :: Word64
freeBits = Word64
64 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
totalBits in Word64 -> Word64 -> [Word64] -> Builder
go Word64
totalBits  (Word64
newAcc Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. Word64
freeBits Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.>. Word64
freeBits)  [Word64]
ws

createFileIndex :: IO.Handle -> Count -> Count -> [Word64] -> IO ()
createFileIndex :: Handle -> Word64 -> Word64 -> [Word64] -> IO ()
createFileIndex Handle
hOut Word64
wordSize Word64
inSize [Word64]
ws = do
  Integer
headerPos <- Handle -> IO Integer
IO.hTell Handle
hOut

  Handle -> Builder -> IO ()
B.hPutBuilder Handle
hOut (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ Builder
forall a. Monoid a => a
mempty
    Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
B.word64LE Word64
wordSize              -- Number of bits in a packed word
    Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
B.word64LE (Word64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
inSize) -- Number of entries
    Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
B.word64LE Word64
0                     -- Number of bytes in packed vector

  Integer
startPos <- Handle -> IO Integer
IO.hTell Handle
hOut

  -- TODO Write packed vector instead
  Handle -> ByteString -> IO ()
LBS.hPut Handle
hOut (Builder -> ByteString
B.toLazyByteString (Word64 -> [Word64] -> Builder
encodePacked Word64
wordSize [Word64]
ws))

  Integer
endPos <- Handle -> IO Integer
IO.hTell Handle
hOut

  let vBytes :: Integer
vBytes = Integer
endPos Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
startPos

  Handle -> SeekMode -> Integer -> IO ()
IO.hSeek Handle
hOut SeekMode
IO.AbsoluteSeek Integer
headerPos

  Handle -> Builder -> IO ()
B.hPutBuilder Handle
hOut (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ Builder
forall a. Monoid a => a
mempty
    Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
B.word64LE Word64
wordSize              -- Number of bits in a packed word
    Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
B.word64LE (Word64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
inSize) -- Number of entries
    Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
B.word64LE (Integer -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
vBytes) -- Number of bytes in packed vector

  Handle -> SeekMode -> Integer -> IO ()
IO.hSeek Handle
hOut SeekMode
IO.AbsoluteSeek Integer
endPos