{-# LANGUAGE DeriveGeneric #-}

module HaskellWorks.Data.Json.Standard.Cursor.IbBp
  ( IbBp(..)
  , slowToIbBp
  , simdToIbBp
  ) where

import Control.Monad
import Control.Monad.ST                    (ST)
import Data.Word
import GHC.Generics
import HaskellWorks.Data.Vector.AsVector64

import qualified Data.ByteString                                                    as BS
import qualified Data.ByteString.Lazy                                               as LBS
import qualified Data.Vector.Storable                                               as DVS
import qualified Data.Vector.Storable.Mutable                                       as DVSM
import qualified HaskellWorks.Data.ByteString.Lazy                                  as LBS
import qualified HaskellWorks.Data.Json.Simd.Index.Standard                         as STSI
import qualified HaskellWorks.Data.Json.Standard.Cursor.Internal.BlankedJson        as J
import qualified HaskellWorks.Data.Json.Standard.Cursor.Internal.ToBalancedParens64 as J
import qualified HaskellWorks.Data.Json.Standard.Cursor.Internal.ToInterestBits64   as J
import qualified HaskellWorks.Data.Vector.Storable                                  as DVS

data IbBp = IbBp
  { IbBp -> Vector Word64
ib :: DVS.Vector Word64
  , IbBp -> Vector Word64
bp :: DVS.Vector Word64
  } deriving (forall x. IbBp -> Rep IbBp x)
-> (forall x. Rep IbBp x -> IbBp) -> Generic IbBp
forall x. Rep IbBp x -> IbBp
forall x. IbBp -> Rep IbBp x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep IbBp x -> IbBp
$cfrom :: forall x. IbBp -> Rep IbBp x
Generic

slowToIbBp :: BS.ByteString -> IbBp
slowToIbBp :: ByteString -> IbBp
slowToIbBp ByteString
bs = IbBp :: Vector Word64 -> Vector Word64 -> IbBp
IbBp
  { ib :: Vector Word64
ib = BlankedJson -> Vector Word64
forall a. ToInterestBits64 a => a -> Vector Word64
J.toInterestBits64 BlankedJson
blankedJson
  , bp :: Vector Word64
bp = BlankedJson -> Vector Word64
forall a. ToBalancedParens64 a => a -> Vector Word64
J.toBalancedParens64 BlankedJson
blankedJson
  }
  where blankedJson :: BlankedJson
blankedJson = ByteString -> BlankedJson
forall a. ToBlankedJson a => a -> BlankedJson
J.toBlankedJsonTyped ByteString
bs

simdToIbBp :: BS.ByteString -> IbBp
simdToIbBp :: ByteString -> IbBp
simdToIbBp ByteString
bs = case ByteString -> Either String [(ByteString, ByteString)]
STSI.makeStandardJsonIbBps (Int -> ByteString -> ByteString
LBS.rechunkPadded Int
chunkSize (ByteString -> ByteString
LBS.fromStrict ByteString
bs)) of
  Left String
msg     -> String -> IbBp
forall a. HasCallStack => String -> a
error String
msg
  Right [(ByteString, ByteString)]
chunks -> (Vector Word64 -> Vector Word64 -> IbBp)
-> (Vector Word64, Vector Word64) -> IbBp
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Vector Word64 -> Vector Word64 -> IbBp
IbBp ((Vector Word64, Vector Word64) -> IbBp)
-> (Vector Word64, Vector Word64) -> IbBp
forall a b. (a -> b) -> a -> b
$ Int
-> (forall s.
    (ByteString, ByteString) -> MVector s Word64 -> ST s Int)
-> Int
-> (forall s.
    (ByteString, ByteString) -> MVector s Word64 -> ST s Int)
-> [(ByteString, ByteString)]
-> (Vector Word64, Vector Word64)
forall b c a.
(Storable b, Storable c) =>
Int
-> (forall s. a -> MVector s b -> ST s Int)
-> Int
-> (forall s. a -> MVector s c -> ST s Int)
-> [a]
-> (Vector b, Vector c)
DVS.construct2N Int
maxSize (ByteString -> MVector s Word64 -> ST s Int
forall s. ByteString -> MVector s Word64 -> ST s Int
go (ByteString -> MVector s Word64 -> ST s Int)
-> ((ByteString, ByteString) -> ByteString)
-> (ByteString, ByteString)
-> MVector s Word64
-> ST s Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString, ByteString) -> ByteString
forall a b. (a, b) -> a
fst) Int
maxSize (ByteString -> MVector s Word64 -> ST s Int
forall s. ByteString -> MVector s Word64 -> ST s Int
go (ByteString -> MVector s Word64 -> ST s Int)
-> ((ByteString, ByteString) -> ByteString)
-> (ByteString, ByteString)
-> MVector s Word64
-> ST s Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString, ByteString) -> ByteString
forall a b. (a, b) -> b
snd) [(ByteString, ByteString)]
chunks
  where chunkSize :: Int
chunkSize = Int
16384
        maxSize :: Int
maxSize = (ByteString -> Int
BS.length ByteString
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8
        go :: BS.ByteString -> DVSM.MVector s Word64 -> ST s Int
        go :: ByteString -> MVector s Word64 -> ST s Int
go ByteString
bytes MVector s Word64
mv = do
          let source :: Vector Word64
source = ByteString -> Vector Word64
forall a. AsVector64 a => a -> Vector Word64
asVector64 ByteString
bytes
          let target :: MVector s Word64
target = Int -> MVector s Word64 -> MVector s Word64
forall a s. Storable a => Int -> MVector s a -> MVector s a
DVSM.take (Vector Word64 -> Int
forall a. Storable a => Vector a -> Int
DVS.length Vector Word64
source) MVector s Word64
mv
          Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Vector Word64 -> Int
forall a. Storable a => Vector a -> Int
DVS.length Vector Word64
source Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ MVector (PrimState (ST s)) Word64 -> Vector Word64 -> ST s ()
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> Vector a -> m ()
DVS.copy MVector s Word64
MVector (PrimState (ST s)) Word64
target Vector Word64
source
          Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector Word64 -> Int
forall a. Storable a => Vector a -> Int
DVS.length Vector Word64
source)