{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE UnboxedTuples       #-}
{-# OPTIONS_HADDOCK not-home #-}

module Database.LSMTree.Internal.StrictArray (
    StrictArray,
    vectorToStrictArray,
    indexStrictArray,
    sizeofStrictArray,
) where

import           Data.Elevator (Strict (Strict))
import qualified Data.Vector as V
import           GHC.Exts ((+#))
import qualified GHC.Exts as GHC
import           GHC.ST (ST (ST), runST)
import           Unsafe.Coerce (unsafeCoerce)

data StrictArray a = StrictArray !(GHC.Array# (Strict a))

vectorToStrictArray :: forall a. V.Vector a -> StrictArray a
vectorToStrictArray :: forall a. Vector a -> StrictArray a
vectorToStrictArray Vector a
v =
    (forall s. ST s (StrictArray a)) -> StrictArray a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (StrictArray a)) -> StrictArray a)
-> (forall s. ST s (StrictArray a)) -> StrictArray a
forall a b. (a -> b) -> a -> b
$ STRep s (StrictArray a) -> ST s (StrictArray a)
forall s a. STRep s a -> ST s a
ST (STRep s (StrictArray a) -> ST s (StrictArray a))
-> STRep s (StrictArray a) -> ST s (StrictArray a)
forall a b. (a -> b) -> a -> b
$ \State# s
s0 ->
    case Int#
-> Strict a
-> State# s
-> (# State# s, MutableArray# s (Strict a) #)
forall a d.
Int# -> a -> State# d -> (# State# d, MutableArray# d a #)
GHC.newArray# (case Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
v of GHC.I# Int#
l# -> Int#
l#)
                       (a -> Strict a
forall a. a -> Strict a
Strict (() -> a
forall a b. a -> b
unsafeCoerce ())) State# s
s0 of
                       -- initialise with (), will be overwritten.
      (# State# s
s1, MutableArray# s (Strict a)
a#  #) ->
        case MutableArray# s (Strict a) -> Int# -> State# s -> State# s
forall s.
MutableArray# s (Strict a) -> Int# -> State# s -> State# s
go MutableArray# s (Strict a)
a# Int#
0# State# s
s1 of
          State# s
s2 -> case MutableArray# s (Strict a)
-> State# s -> (# State# s, Array# (Strict a) #)
forall d a.
MutableArray# d a -> State# d -> (# State# d, Array# a #)
GHC.unsafeFreezeArray# MutableArray# s (Strict a)
a# State# s
s2 of
                  (# State# s
s3, Array# (Strict a)
a'# #) -> (# State# s
s3, Array# (Strict a) -> StrictArray a
forall a. Array# (Strict a) -> StrictArray a
StrictArray Array# (Strict a)
a'# #)
  where
    go :: forall s.
          GHC.MutableArray# s (Strict a)
       -> GHC.Int#
       -> GHC.State# s
       -> GHC.State# s
    go :: forall s.
MutableArray# s (Strict a) -> Int# -> State# s -> State# s
go MutableArray# s (Strict a)
a# Int#
i# State# s
s
      | Int# -> Int
GHC.I# Int#
i# Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
v
      = let x :: a
x = Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
v (Int# -> Int
GHC.I# Int#
i#)
            -- We have to use seq# here to force the array element to WHNF
            -- before putting it into the strict array. This should not be
            -- necessary. https://github.com/sgraf812/data-elevator/issues/4
         in case a -> State# s -> (# State# s, a #)
forall a d. a -> State# d -> (# State# d, a #)
GHC.seq# a
x State# s
s of
              (# State# s
s', a
x' #) ->
                case MutableArray# s (Strict a)
-> Int# -> Strict a -> State# s -> State# s
forall d a. MutableArray# d a -> Int# -> a -> State# d -> State# d
GHC.writeArray# MutableArray# s (Strict a)
a# Int#
i# (a -> Strict a
forall a. a -> Strict a
Strict a
x') State# s
s' of
                  State# s
s'' -> MutableArray# s (Strict a) -> Int# -> State# s -> State# s
forall s.
MutableArray# s (Strict a) -> Int# -> State# s -> State# s
go MutableArray# s (Strict a)
a# (Int#
i# Int# -> Int# -> Int#
+# Int#
1#) State# s
s''
      | Bool
otherwise = State# s
s

{-# INLINE indexStrictArray #-}
indexStrictArray :: StrictArray a -> Int -> a
indexStrictArray :: forall a. StrictArray a -> Int -> a
indexStrictArray (StrictArray Array# (Strict a)
a#) (GHC.I# Int#
i#) =
    case Array# (Strict a) -> Int# -> (# Strict a #)
forall a. Array# a -> Int# -> (# a #)
GHC.indexArray# Array# (Strict a)
a# Int#
i# of
      (# Strict a
x #) -> a
x

{-# INLINE sizeofStrictArray #-}
sizeofStrictArray :: StrictArray a -> Int
sizeofStrictArray :: forall a. StrictArray a -> Int
sizeofStrictArray (StrictArray Array# (Strict a)
a#) =
    Int# -> Int
GHC.I# (Array# (Strict a) -> Int#
forall a. Array# a -> Int#
GHC.sizeofArray# Array# (Strict a)
a#)