{-# LANGUAGE GADTs      #-}
{-# LANGUAGE RankNTypes #-}

module HaskellWorks.Data.Json.Simd.Index.Simple
  ( makeSimpleJsonIbBps
  , makeSimpleJsonIbBpsUnsafe
  , enabledMakeSimpleJsonIbBps
  ) where

import Control.Monad.ST
import Data.Word
import HaskellWorks.Data.Json.Simd.Internal.Index.Simple

import qualified Control.Monad.ST.Unsafe                      as ST
import qualified Data.ByteString                              as BS
import qualified Data.ByteString.Internal                     as BSI
import qualified Data.ByteString.Lazy                         as LBS
import qualified Data.Vector.Storable.Mutable                 as DVSM
import qualified Foreign.ForeignPtr                           as F
import qualified Foreign.ForeignPtr.Unsafe                    as F
import qualified Foreign.Marshal.Unsafe                       as F
import qualified Foreign.Ptr                                  as F
import qualified HaskellWorks.Data.Json.Simd.Capabilities     as C
import qualified HaskellWorks.Data.Json.Simd.Internal.Foreign as F
import qualified HaskellWorks.Data.Json.Simd.Internal.List    as L
import qualified System.IO.Unsafe                             as IO

{- HLINT ignore "Reduce duplication"  -}
{- HLINT ignore "Redundant do"        -}

makeSimpleJsonIbBps :: LBS.ByteString -> Either String [(BS.ByteString, BS.ByteString)]
makeSimpleJsonIbBps :: ByteString -> Either String [(ByteString, ByteString)]
makeSimpleJsonIbBps ByteString
lbs = if Bool
enabledMakeSimpleJsonIbBps
  then forall a b. b -> Either a b
Right (ByteString -> [(ByteString, ByteString)]
makeSimpleJsonIbBpsUnsafe ByteString
lbs)
  else forall a b. a -> Either a b
Left String
"makeSimpleJsonIbBps function is disabled"

makeSimpleJsonIbBpsUnsafe :: LBS.ByteString -> [(BS.ByteString, BS.ByteString)]
makeSimpleJsonIbBpsUnsafe :: ByteString -> [(ByteString, ByteString)]
makeSimpleJsonIbBpsUnsafe ByteString
lbs = forall a b. a -> b -> [a] -> [b] -> [(a, b)]
L.zipPadded ByteString
BS.empty ByteString
BS.empty [ByteString]
ibs [ByteString]
bps
  where chunks :: [(ByteString, ByteString, ByteString)]
chunks  = ByteString -> [(ByteString, ByteString, ByteString)]
makeIbs ByteString
lbs
        ibs :: [ByteString]
ibs     = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(ByteString
a, ByteString
_, ByteString
_) -> ByteString
a) [(ByteString, ByteString, ByteString)]
chunks
        bps :: [ByteString]
bps     = [(ByteString, ByteString, ByteString)] -> [ByteString]
ibsToIndexByteStrings [(ByteString, ByteString, ByteString)]
chunks

makeIbs :: LBS.ByteString -> [(BS.ByteString, BS.ByteString, BS.ByteString)]
makeIbs :: ByteString -> [(ByteString, ByteString, ByteString)]
makeIbs ByteString
lbs = forall a. IO a -> a
F.unsafeLocalState forall a b. (a -> b) -> a -> b
$ do
  WorkBuffers
wb <- Int -> IO WorkBuffers
allocWorkBuffers (Int
32 forall a. Num a => a -> a -> a
* Int
1024 forall a. Num a => a -> a -> a
* Int
1204)
  WorkState
ws <- IO WorkState
allocWorkState
  forall a. IO a -> IO a
IO.unsafeInterleaveIO forall a b. (a -> b) -> a -> b
$ WorkBuffers
-> WorkState
-> [ByteString]
-> IO [(ByteString, ByteString, ByteString)]
go WorkBuffers
wb WorkState
ws (ByteString -> [ByteString]
LBS.toChunks ByteString
lbs)
  where go :: WorkBuffers -> WorkState -> [BS.ByteString] -> IO [(BS.ByteString, BS.ByteString, BS.ByteString)]
        go :: WorkBuffers
-> WorkState
-> [ByteString]
-> IO [(ByteString, ByteString, ByteString)]
go WorkBuffers
_  WorkState
_  []       = forall (m :: * -> *) a. Monad m => a -> m a
return []
        go WorkBuffers
wb WorkState
ws (ByteString
bs:[ByteString]
bss) = do
          let resLen :: Int
resLen = ByteString -> Int
BS.length ByteString
bs forall a. Integral a => a -> a -> a
`div` Int
8
          ForeignPtr Word8
resIbFptr  <- forall a. Int -> IO (ForeignPtr a)
F.mallocForeignPtrBytes Int
resLen
          ForeignPtr Word8
resAFptr   <- forall a. Int -> IO (ForeignPtr a)
F.mallocForeignPtrBytes Int
resLen
          ForeignPtr Word8
resBFptr   <- forall a. Int -> IO (ForeignPtr a)
F.mallocForeignPtrBytes Int
resLen
          let resIbPtr :: Ptr UInt8
resIbPtr  = forall a b. Ptr a -> Ptr b
F.castPtr (forall a. ForeignPtr a -> Ptr a
F.unsafeForeignPtrToPtr ForeignPtr Word8
resIbFptr)
          let resAPtr :: Ptr UInt8
resAPtr   = forall a b. Ptr a -> Ptr b
F.castPtr (forall a. ForeignPtr a -> Ptr a
F.unsafeForeignPtrToPtr ForeignPtr Word8
resAFptr )
          let resBPtr :: Ptr UInt8
resBPtr   = forall a b. Ptr a -> Ptr b
F.castPtr (forall a. ForeignPtr a -> Ptr a
F.unsafeForeignPtrToPtr ForeignPtr Word8
resBFptr )
          let (ForeignPtr Word8
bsFptr, Int
bsOff, Int
bsLen) = ByteString -> (ForeignPtr Word8, Int, Int)
BSI.toForeignPtr ByteString
bs
          let bsPtr :: Ptr Any
bsPtr = forall a b. Ptr a -> Ptr b
F.castPtr (forall a. ForeignPtr a -> Ptr a
F.unsafeForeignPtrToPtr ForeignPtr Word8
bsFptr)
          Size
_ <- Ptr UInt8
-> Size
-> Ptr UInt8
-> Ptr UInt8
-> Ptr UInt8
-> Ptr UInt8
-> Ptr UInt8
-> Ptr UInt8
-> Ptr Size
-> Ptr Size
-> Ptr Size
-> Ptr Size
-> Ptr UInt8
-> Ptr UInt8
-> Ptr UInt8
-> IO Size
F.processChunk
            (forall a b. Ptr a -> Int -> Ptr b
F.plusPtr Ptr Any
bsPtr Int
bsOff) -- in_buffer:           Ptr UInt8
            (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bsLen)    -- in_length:           Size
            (WorkBuffers -> Ptr UInt8
workBuffersD WorkBuffers
wb)       -- work_bits_of_d:      Ptr UInt8
            (WorkBuffers -> Ptr UInt8
workBuffersA WorkBuffers
wb)       -- work_bits_of_a:      Ptr UInt8
            (WorkBuffers -> Ptr UInt8
workBuffersZ WorkBuffers
wb)       -- work_bits_of_z:      Ptr UInt8
            (WorkBuffers -> Ptr UInt8
workBuffersQ WorkBuffers
wb)       -- work_bits_of_q:      Ptr UInt8
            (WorkBuffers -> Ptr UInt8
workBuffersB WorkBuffers
wb)       -- work_bits_of_b:      Ptr UInt8
            (WorkBuffers -> Ptr UInt8
workBuffersE WorkBuffers
wb)       -- work_bits_of_e:      Ptr UInt8
            (WorkState -> Ptr Size
workStateZ WorkState
ws)         -- last_trailing_ones:  Ptr Size
            (WorkState -> Ptr Size
workStateO WorkState
ws)         -- quote_odds_carry:    Ptr Size
            (WorkState -> Ptr Size
workStateE WorkState
ws)         -- quote_evens_carry:   Ptr Size
            (WorkState -> Ptr Size
workStateM WorkState
ws)         -- quote_mask_carry:    Ptr UInt64
            Ptr UInt8
resIbPtr                -- result_ibs:          Ptr UInt8
            Ptr UInt8
resAPtr                 -- result_a:            Ptr UInt8
            Ptr UInt8
resBPtr                 -- result_z:            Ptr UInt8
          let r :: (ByteString, ByteString, ByteString)
r =
                ( ForeignPtr Word8 -> Int -> Int -> ByteString
BSI.fromForeignPtr ForeignPtr Word8
resIbFptr Int
0 Int
resLen
                , ForeignPtr Word8 -> Int -> Int -> ByteString
BSI.fromForeignPtr ForeignPtr Word8
resAFptr  Int
0 Int
resLen
                , ForeignPtr Word8 -> Int -> Int -> ByteString
BSI.fromForeignPtr ForeignPtr Word8
resBFptr  Int
0 Int
resLen
                )
          [(ByteString, ByteString, ByteString)]
rs <- forall a. IO a -> IO a
IO.unsafeInterleaveIO forall a b. (a -> b) -> a -> b
$ WorkBuffers
-> WorkState
-> [ByteString]
-> IO [(ByteString, ByteString, ByteString)]
go WorkBuffers
wb WorkState
ws [ByteString]
bss
          forall (m :: * -> *) a. Monad m => a -> m a
return ((ByteString, ByteString, ByteString)
rforall a. a -> [a] -> [a]
:[(ByteString, ByteString, ByteString)]
rs)

ibsToIndexByteStrings :: ()
  => [(BS.ByteString, BS.ByteString, BS.ByteString)]
  -> [BS.ByteString]
ibsToIndexByteStrings :: [(ByteString, ByteString, ByteString)] -> [ByteString]
ibsToIndexByteStrings [(ByteString, ByteString, ByteString)]
bits = forall a. IO a -> a
F.unsafeLocalState forall a b. (a -> b) -> a -> b
$ do
  BpState
bpState <- IO BpState
emptyBpState
  forall a. IO a -> IO a
IO.unsafeInterleaveIO forall a b. (a -> b) -> a -> b
$ BpState -> [Step] -> IO [ByteString]
go BpState
bpState (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(ByteString
a, ByteString
b, ByteString
c) -> ByteString -> ByteString -> ByteString -> Step
mkIndexStep ByteString
a ByteString
b ByteString
c) [(ByteString, ByteString, ByteString)]
bits)
  where go :: ()
          => BpState
          -> [Step]
          -> IO [BS.ByteString]
        go :: BpState -> [Step] -> IO [ByteString]
go BpState
s (Step
step:[Step]
steps) = do
          let bp :: ByteString
bp = BpState -> Step -> ByteString
stepToByteString BpState
s Step
step
          [ByteString]
bps <- forall a. IO a -> IO a
IO.unsafeInterleaveIO forall a b. (a -> b) -> a -> b
$ BpState -> [Step] -> IO [ByteString]
go BpState
s [Step]
steps
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ ByteString
bpforall a. a -> [a] -> [a]
:[ByteString]
bps
        go BpState
s [] = forall (m :: * -> *) a. Monad m => a -> m a
return [BpState -> Step -> ByteString
stepToByteString BpState
s Step
indexStepFinal]

mkIndexStep :: BS.ByteString -> BS.ByteString -> BS.ByteString -> Step
mkIndexStep :: ByteString -> ByteString -> ByteString -> Step
mkIndexStep ByteString
is ByteString
as ByteString
zs | Int
isLen forall a. Eq a => a -> a -> Bool
== Int
asLen Bool -> Bool -> Bool
&& Int
asLen forall a. Eq a => a -> a -> Bool
== Int
zsLen = (forall s. BpState -> MVector s Word64 -> ST s Int) -> Int -> Step
Step forall s. BpState -> MVector s Word64 -> ST s Int
go Int
isLen
  where isLen :: Int
isLen = ByteString -> Int
BS.length ByteString
is
        asLen :: Int
asLen = ByteString -> Int
BS.length ByteString
as
        zsLen :: Int
zsLen = ByteString -> Int
BS.length ByteString
zs
        (ForeignPtr Word8
isFptr, Int
_, Int
_) = ByteString -> (ForeignPtr Word8, Int, Int)
BSI.toForeignPtr ByteString
is
        (ForeignPtr Word8
asFptr, Int
_, Int
_) = ByteString -> (ForeignPtr Word8, Int, Int)
BSI.toForeignPtr ByteString
as
        (ForeignPtr Word8
zsFptr, Int
_, Int
_) = ByteString -> (ForeignPtr Word8, Int, Int)
BSI.toForeignPtr ByteString
zs
        go  :: BpState
            -> DVSM.MVector s Word64
            -> ST s Int
        go :: forall s. BpState -> MVector s Word64 -> ST s Int
go BpState
bpState MVector s Word64
bpvm = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a s. IO a -> ST s a
ST.unsafeIOToST forall a b. (a -> b) -> a -> b
$ do
          let (ForeignPtr Word64
outFptr, Int
_, Int
_) = forall s a. MVector s a -> (ForeignPtr a, Int, Int)
DVSM.unsafeToForeignPtr MVector s Word64
bpvm

          forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
F.withForeignPtr ForeignPtr Word64
outFptr forall a b. (a -> b) -> a -> b
$ \Ptr Word64
outPtr ->
            forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
F.withForeignPtr ForeignPtr Word8
isFptr forall a b. (a -> b) -> a -> b
$ \Ptr Word8
isPtr ->
              forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
F.withForeignPtr ForeignPtr Word8
asFptr forall a b. (a -> b) -> a -> b
$ \Ptr Word8
asPtr ->
                forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
F.withForeignPtr ForeignPtr Word8
zsFptr forall a b. (a -> b) -> a -> b
$ \Ptr Word8
zsPtr ->
                  forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
F.withForeignPtr (BpState -> ForeignPtr Word8
bpStateP BpState
bpState) forall a b. (a -> b) -> a -> b
$ \Ptr Word8
bpStatePtr -> do
                    Ptr UInt8
-> Ptr UInt8 -> Ptr UInt8 -> Size -> Ptr () -> Ptr UInt8 -> IO Size
F.writeBpChunk
                      (forall a b. Ptr a -> Ptr b
F.castPtr Ptr Word8
isPtr)
                      (forall a b. Ptr a -> Ptr b
F.castPtr Ptr Word8
asPtr)
                      (forall a b. Ptr a -> Ptr b
F.castPtr Ptr Word8
zsPtr)
                      (forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
is))
                      (forall a b. Ptr a -> Ptr b
F.castPtr Ptr Word8
bpStatePtr)
                      (forall a b. Ptr a -> Ptr b
F.castPtr Ptr Word64
outPtr)
mkIndexStep ByteString
_ ByteString
_ ByteString
_ = forall a. HasCallStack => String -> a
error String
"Mismatched input size"

indexStepFinal :: Step
indexStepFinal :: Step
indexStepFinal = (forall s. BpState -> MVector s Word64 -> ST s Int) -> Int -> Step
Step forall s. BpState -> MVector s Word64 -> ST s Int
go Int
2
  where go  :: BpState
            -> DVSM.MVector s Word64
            -> ST s Int
        go :: forall s. BpState -> MVector s Word64 -> ST s Int
go BpState
bpState MVector s Word64
bpvm = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a s. IO a -> ST s a
ST.unsafeIOToST forall a b. (a -> b) -> a -> b
$ do
          let (ForeignPtr Word64
outFptr, Int
_, Int
_) = forall s a. MVector s a -> (ForeignPtr a, Int, Int)
DVSM.unsafeToForeignPtr MVector s Word64
bpvm

          forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
F.withForeignPtr ForeignPtr Word64
outFptr forall a b. (a -> b) -> a -> b
$ \Ptr Word64
outPtr ->
            forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
F.withForeignPtr (BpState -> ForeignPtr Word8
bpStateP BpState
bpState) forall a b. (a -> b) -> a -> b
$ \Ptr Word8
bpStatePtr -> do
              Ptr () -> Ptr UInt8 -> IO Size
F.writeBpChunkFinal (forall a b. Ptr a -> Ptr b
F.castPtr Ptr Word8
bpStatePtr) (forall a b. Ptr a -> Ptr b
F.castPtr Ptr Word64
outPtr)

stepToByteString :: BpState -> Step -> BS.ByteString
stepToByteString :: BpState -> Step -> ByteString
stepToByteString BpState
state (Step forall s. BpState -> MVector s Word64 -> ST s Int
step Int
size) = forall a. IO a -> a
F.unsafeLocalState forall a b. (a -> b) -> a -> b
$ do
  let bsSize :: Int
bsSize = Int
size forall a. Num a => a -> a -> a
* Int
8
  ForeignPtr Word8
bpFptr <- forall a. Int -> IO (ForeignPtr a)
BSI.mallocByteString Int
bsSize
  let bpVm :: MVector RealWorld Word64
bpVm = forall a s. Storable a => ForeignPtr a -> Int -> Int -> MVector s a
DVSM.unsafeFromForeignPtr (forall a b. ForeignPtr a -> ForeignPtr b
F.castForeignPtr ForeignPtr Word8
bpFptr) Int
0 Int
size
  Int
w64Size <- forall a. ST RealWorld a -> IO a
stToIO forall a b. (a -> b) -> a -> b
$ forall s. BpState -> MVector s Word64 -> ST s Int
step BpState
state MVector RealWorld Word64
bpVm
  forall (m :: * -> *) a. Monad m => a -> m a
return (ForeignPtr Word8 -> Int -> Int -> ByteString
BSI.PS ForeignPtr Word8
bpFptr Int
0 (Int
w64Size forall a. Num a => a -> a -> a
* Int
8))

enabledMakeSimpleJsonIbBps :: Bool
enabledMakeSimpleJsonIbBps :: Bool
enabledMakeSimpleJsonIbBps = Bool
C.avx_2 Bool -> Bool -> Bool
&& Bool
C.sse_4_2 Bool -> Bool -> Bool
&& Bool
C.bmi_2