content: {-# LANGUAGE UnboxedTuples #-}
  {-# OPTIONS_HADDOCK not-home #-}
  -- |
  -- Copyright: (c) 2021 Xy Ren
  -- License: BSD3
  -- Maintainer: xy.r@outlook.com
  -- Stability: unstable
  -- Portability: non-portable (GHC only)
  module Data.PrimVec (PrimVec, empty, head, tail, take, drop, index, cons, concat, pick, update) where
  
  import           Data.Foldable            (for_)
  import           Data.Primitive           (Prim, indexPrimArray)
  import           Data.Primitive.PrimArray (MutablePrimArray (MutablePrimArray), PrimArray (PrimArray), copyPrimArray,
                                             newPrimArray, writePrimArray)
  import           GHC.Exts                 (runRW#, unsafeFreezeByteArray#)
  import           GHC.ST                   (ST (ST))
  import           Prelude                  hiding (concat, drop, head, tail, take)
  
  -- | Slices of 'PrimArray'.
  data PrimVec a = PrimVec
    {-# UNPACK #-} !Int
    {-# UNPACK #-} !Int
    {-# UNPACK #-} !(PrimArray a)
  
  runPrimArray :: (∀ s. ST s (MutablePrimArray s a)) -> PrimArray a
  runPrimArray (ST f) = let
    !(# _, ba# #) = runRW# \s1 ->
      let !(# s2, MutablePrimArray mba# #) = f s1
      in unsafeFreezeByteArray# mba# s2
    in PrimArray ba#
  
  empty :: Prim a => PrimVec a
  empty = PrimVec 0 0 $ runPrimArray $ newPrimArray 0
  
  head :: Prim a => PrimVec a -> a
  head (PrimVec off _ arr) = indexPrimArray arr off
  
  tail :: PrimVec a -> PrimVec a
  tail (PrimVec off len arr) = PrimVec (off + 1) (len - 1) arr
  
  take :: Prim a => Int -> PrimVec a -> PrimVec a
  take n (PrimVec off _ arr) = PrimVec 0 n $ runPrimArray do
    marr <- newPrimArray n
    copyPrimArray marr 0 arr off n
    pure marr
  
  drop :: Int -> PrimVec a -> PrimVec a
  drop n (PrimVec off len arr) = PrimVec (off + n) (len - n) arr
  
  index :: Prim a => Int -> PrimVec a -> a
  index n (PrimVec off _ arr) = indexPrimArray arr (off + n)
  
  cons :: Prim a => a -> PrimVec a -> PrimVec a
  cons x (PrimVec off len arr) = PrimVec 0 (len + 1) $ runPrimArray do
    marr <- newPrimArray (len + 1)
    writePrimArray marr 0 x
    copyPrimArray marr 1 arr off len
    pure marr
  
  concat :: Prim a => PrimVec a -> PrimVec a -> PrimVec a
  concat (PrimVec off len arr) (PrimVec off' len' arr') = PrimVec 0 (len + len') $ runPrimArray do
    marr <- newPrimArray (len + len')
    copyPrimArray marr 0 arr off len
    copyPrimArray marr len arr' off' len'
    pure marr
  
  pick :: Prim a => Int -> [Int] -> PrimVec a -> PrimVec a
  pick len' ns (PrimVec off _ arr) = PrimVec 0 len' $ runPrimArray do
    marr <- newPrimArray len'
    for_ (zip [0 ..] ns) \(new, old) ->
      writePrimArray marr new (indexPrimArray arr (off + old))
    pure marr
  
  update :: Prim a => Int -> a -> PrimVec a -> PrimVec a
  update n x (PrimVec off len arr) = PrimVec 0 len $ runPrimArray do
    marr <- newPrimArray len
    copyPrimArray marr 0 arr off len
    writePrimArray marr n x
    pure marr
  