{-# LANGUAGE CPP           #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module HaskellWorks.Data.Dsv.Internal.Vector
  ( empty64
  , constructNS
  , ltWord
  , indexCsvChunk
  , oddsMask
  ) where

import Control.Monad.ST
import Data.Bits.Pdep
import Data.Word
import Foreign.Storable                          (Storable)
import GHC.Int
import GHC.Prim
import GHC.Word                                  hiding (ltWord)
import HaskellWorks.Data.Bits.BitWise
import HaskellWorks.Data.Bits.PopCount.PopCount1
import HaskellWorks.Data.Positioning

import qualified Data.Vector.Storable         as DVS
import qualified Data.Vector.Storable.Mutable as DVSM

empty64 :: DVS.Vector Word64
empty64 :: Vector Word64
empty64 = Int -> Word64 -> Vector Word64
forall a. Storable a => Int -> a -> Vector a
DVS.replicate Int
64 Word64
0
{-# NOINLINE empty64 #-}

constructNS :: forall a s. Storable a => Int -> s -> (s -> DVS.Vector a -> (s, a)) -> (s, DVS.Vector a)
constructNS :: forall a s.
Storable a =>
Int -> s -> (s -> Vector a -> (s, a)) -> (s, Vector a)
constructNS Int
n s
s s -> Vector a -> (s, a)
f = (forall s. ST s (s, MVector s a)) -> (s, Vector a)
forall (f :: * -> *) a.
(Traversable f, Storable a) =>
(forall s. ST s (f (MVector s a))) -> f (Vector a)
DVS.createT (Int -> s -> ST s (s, MVector s a)
forall q. Int -> s -> ST q (s, MVector q a)
go Int
0 s
s)
  where go :: forall q. Int -> s -> ST q (s, DVS.MVector q a)
        go :: forall q. Int -> s -> ST q (s, MVector q a)
go Int
n1 s
s1 = do
          MVector q a
mv :: DVS.MVector q a <- Int -> ST q (MVector (PrimState (ST q)) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew Int
n
          Vector a
u <- MVector (PrimState (ST q)) a -> ST q (Vector a)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
DVS.unsafeFreeze MVector q a
MVector (PrimState (ST q)) a
mv
          let (s
s2, a
w) = s -> Vector a -> (s, a)
f s
s1 (Int -> Vector a -> Vector a
forall a. Storable a => Int -> Vector a -> Vector a
DVS.take Int
n1 Vector a
u)
          MVector (PrimState (ST q)) a -> Int -> a -> ST q ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector q a
MVector (PrimState (ST q)) a
mv Int
n1 a
w
          (s, MVector q a) -> ST q (s, MVector q a)
forall a. a -> ST q a
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s2, MVector q a
mv)
{-# INLINE constructNS #-}

ltWord :: Word64 -> Word64 -> Word64
#if MIN_VERSION_base(4,17,2)
ltWord :: Word64 -> Word64 -> Word64
ltWord (W64# Word64#
a#) (W64# Word64#
b#) = Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int# -> Int
I# (Word64# -> Word64# -> Int#
ltWord64# Word64#
a# Word64#
b#))
#else
ltWord (W64# a#) (W64# b#) = fromIntegral (I64# (ltWord# a# b#))
#endif
{-# INLINE ltWord #-}

indexCsvChunk ::
     Count
  -> Word64
  -> DVS.Vector Word64
  -> DVS.Vector Word64
  -> DVS.Vector Word64
  -> (DVS.Vector Word64, DVS.Vector Word64, Word64, Word64)
indexCsvChunk :: Word64
-> Word64
-> Vector Word64
-> Vector Word64
-> Vector Word64
-> (Vector Word64, Vector Word64, Word64, Word64)
indexCsvChunk Word64
qqCount Word64
qqCarry Vector Word64
mks Vector Word64
nls Vector Word64
qqs = (forall s. ST s (Vector Word64, Vector Word64, Word64, Word64))
-> (Vector Word64, Vector Word64, Word64, Word64)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Word64, Vector Word64, Word64, Word64))
 -> (Vector Word64, Vector Word64, Word64, Word64))
-> (forall s. ST s (Vector Word64, Vector Word64, Word64, Word64))
-> (Vector Word64, Vector Word64, Word64, Word64)
forall a b. (a -> b) -> a -> b
$ do
  MVector s Word64
tmks <- Int -> ST s (MVector (PrimState (ST s)) Word64)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew Int
len
  MVector s Word64
tnls <- Int -> ST s (MVector (PrimState (ST s)) Word64)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
DVSM.unsafeNew Int
len
  (Word64
newCount, Word64
newCarry) <- Int
-> Word64
-> Word64
-> MVector s Word64
-> MVector s Word64
-> ST s (Word64, Word64)
forall z.
Int
-> Word64
-> Word64
-> MVector z Word64
-> MVector z Word64
-> ST z (Word64, Word64)
go Int
0 Word64
qqCount Word64
qqCarry MVector s Word64
tmks MVector s Word64
tnls
  Vector Word64
rmks <- MVector (PrimState (ST s)) Word64 -> ST s (Vector Word64)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
DVS.unsafeFreeze MVector s Word64
MVector (PrimState (ST s)) Word64
tmks
  Vector Word64
rnls <- MVector (PrimState (ST s)) Word64 -> ST s (Vector Word64)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
DVS.unsafeFreeze MVector s Word64
MVector (PrimState (ST s)) Word64
tnls
  (Vector Word64, Vector Word64, Word64, Word64)
-> ST s (Vector Word64, Vector Word64, Word64, Word64)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector Word64
rmks, Vector Word64
rnls, Word64
newCount, Word64
newCarry)
  where len :: Int
len = Vector Word64 -> Int
forall a. Storable a => Vector a -> Int
DVS.length Vector Word64
mks
        go :: Int -> Word64 -> Word64 -> DVSM.MVector z Word64 -> DVSM.MVector z Word64 -> ST z (Count, Word64)
        go :: forall z.
Int
-> Word64
-> Word64
-> MVector z Word64
-> MVector z Word64
-> ST z (Word64, Word64)
go Int
i Word64
pc Word64
carry MVector z Word64
tmks MVector z Word64
tnls | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = do
          let qq :: Word64
qq = Vector Word64 -> Int -> Word64
forall a. Storable a => Vector a -> Int -> a
DVS.unsafeIndex Vector Word64
qqs Int
i
          let mk :: Word64
mk = Vector Word64 -> Int -> Word64
forall a. Storable a => Vector a -> Int -> a
DVS.unsafeIndex Vector Word64
mks Int
i
          let nl :: Word64
nl = Vector Word64 -> Int -> Word64
forall a. Storable a => Vector a -> Int -> a
DVS.unsafeIndex Vector Word64
nls Int
i

          let enters :: Word64
enters = Word64 -> Word64 -> Word64
forall a. Pdep a => a -> a -> a
pdep (Word64
oddsMask Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. (Word64
0x1 Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&.      Word64
pc)) Word64
qq
          let leaves :: Word64
leaves = Word64 -> Word64 -> Word64
forall a. Pdep a => a -> a -> a
pdep (Word64
oddsMask Word64 -> Word64 -> Word64
forall a. Shift a => a -> Word64 -> a
.<. (Word64
0x1 Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. Word64 -> Word64
forall a. BitWise a => a -> a
comp Word64
pc)) Word64
qq

          let compLeaves :: Word64
compLeaves    = Word64 -> Word64
forall a. BitWise a => a -> a
comp Word64
leaves
          let preQuoteMask :: Word64
preQuoteMask  = Word64
enters Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
compLeaves
          let quoteMask :: Word64
quoteMask     = Word64
preQuoteMask Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
carry
          let newCarry :: Word64
newCarry      = Word64
quoteMask Word64 -> Word64 -> Word64
`ltWord` (Word64
enters Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. Word64
compLeaves Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. Word64
carry)

          MVector (PrimState (ST z)) Word64 -> Int -> Word64 -> ST z ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector z Word64
MVector (PrimState (ST z)) Word64
tmks Int
i ((Word64
nl Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.|. Word64
mk) Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. Word64
quoteMask)
          MVector (PrimState (ST z)) Word64 -> Int -> Word64 -> ST z ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
DVSM.unsafeWrite MVector z Word64
MVector (PrimState (ST z)) Word64
tnls Int
i ( Word64
nl         Word64 -> Word64 -> Word64
forall a. BitWise a => a -> a -> a
.&. Word64
quoteMask)

          Int
-> Word64
-> Word64
-> MVector z Word64
-> MVector z Word64
-> ST z (Word64, Word64)
forall z.
Int
-> Word64
-> Word64
-> MVector z Word64
-> MVector z Word64
-> ST z (Word64, Word64)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Word64 -> Word64
forall v. PopCount1 v => v -> Word64
popCount1 Word64
qq Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Word64
pc) Word64
newCarry MVector z Word64
tmks MVector z Word64
tnls
        go Int
_ Word64
pc Word64
carry MVector z Word64
_ MVector z Word64
_ = (Word64, Word64) -> ST z (Word64, Word64)
forall a. a -> ST z a
forall (m :: * -> *) a. Monad m => a -> m a
return (Word64
pc, Word64
carry)
{-# INLINE indexCsvChunk #-}

oddsMask :: Word64
oddsMask :: Word64
oddsMask = Word64
0x5555555555555555
{-# INLINE oddsMask #-}