{-# 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 :: (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray (ST f) = let
  !(# State# RealWorld
_, ByteArray#
ba# #) = (State# RealWorld -> (# State# RealWorld, ByteArray# #))
-> (# State# RealWorld, ByteArray# #)
forall o. (State# RealWorld -> o) -> o
runRW# \State# RealWorld
s1 ->
    let !(# State# RealWorld
s2, MutablePrimArray MutableByteArray# RealWorld
mba# #) = STRep RealWorld (MutablePrimArray RealWorld a)
f State# RealWorld
s1
    in MutableByteArray# RealWorld
-> State# RealWorld -> (# State# RealWorld, ByteArray# #)
forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# RealWorld
mba# State# RealWorld
s2
  in ByteArray# -> PrimArray a
forall a. ByteArray# -> PrimArray a
PrimArray ByteArray#
ba#

empty :: Prim a => PrimVec a
empty :: PrimVec a
empty = Int -> Int -> PrimArray a -> PrimVec a
forall a. Int -> Int -> PrimArray a -> PrimVec a
PrimVec Int
0 Int
0 (PrimArray a -> PrimVec a) -> PrimArray a -> PrimVec a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray ((forall s. ST s (MutablePrimArray s a)) -> PrimArray a)
-> (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
forall a b. (a -> b) -> a -> b
$ Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
0

head :: Prim a => PrimVec a -> a
head :: PrimVec a -> a
head (PrimVec Int
off Int
_ PrimArray a
arr) = PrimArray a -> Int -> a
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr Int
off

tail :: PrimVec a -> PrimVec a
tail :: PrimVec a -> PrimVec a
tail (PrimVec Int
off Int
len PrimArray a
arr) = Int -> Int -> PrimArray a -> PrimVec a
forall a. Int -> Int -> PrimArray a -> PrimVec a
PrimVec (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) PrimArray a
arr

take :: Prim a => Int -> PrimVec a -> PrimVec a
take :: Int -> PrimVec a -> PrimVec a
take Int
n (PrimVec Int
off Int
_ PrimArray a
arr) = Int -> Int -> PrimArray a -> PrimVec a
forall a. Int -> Int -> PrimArray a -> PrimVec a
PrimVec Int
0 Int
n (PrimArray a -> PrimVec a) -> PrimArray a -> PrimVec a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray do
  MutablePrimArray s a
marr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
n
  MutablePrimArray (PrimState (ST s)) a
-> Int -> PrimArray a -> Int -> Int -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
0 PrimArray a
arr Int
off Int
n
  MutablePrimArray s a -> ST s (MutablePrimArray s a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutablePrimArray s a
marr

drop :: Int -> PrimVec a -> PrimVec a
drop :: Int -> PrimVec a -> PrimVec a
drop Int
n (PrimVec Int
off Int
len PrimArray a
arr) = Int -> Int -> PrimArray a -> PrimVec a
forall a. Int -> Int -> PrimArray a -> PrimVec a
PrimVec (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) PrimArray a
arr

index :: Prim a => Int -> PrimVec a -> a
index :: Int -> PrimVec a -> a
index Int
n (PrimVec Int
off Int
_ PrimArray a
arr) = PrimArray a -> Int -> a
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)

cons :: Prim a => a -> PrimVec a -> PrimVec a
cons :: a -> PrimVec a -> PrimVec a
cons a
x (PrimVec Int
off Int
len PrimArray a
arr) = Int -> Int -> PrimArray a -> PrimVec a
forall a. Int -> Int -> PrimArray a -> PrimVec a
PrimVec Int
0 (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (PrimArray a -> PrimVec a) -> PrimArray a -> PrimVec a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray do
  MutablePrimArray s a
marr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  MutablePrimArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
0 a
x
  MutablePrimArray (PrimState (ST s)) a
-> Int -> PrimArray a -> Int -> Int -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
1 PrimArray a
arr Int
off Int
len
  MutablePrimArray s a -> ST s (MutablePrimArray s a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutablePrimArray s a
marr

concat :: Prim a => PrimVec a -> PrimVec a -> PrimVec a
concat :: PrimVec a -> PrimVec a -> PrimVec a
concat (PrimVec Int
off Int
len PrimArray a
arr) (PrimVec Int
off' Int
len' PrimArray a
arr') = Int -> Int -> PrimArray a -> PrimVec a
forall a. Int -> Int -> PrimArray a -> PrimVec a
PrimVec Int
0 (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len') (PrimArray a -> PrimVec a) -> PrimArray a -> PrimVec a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray do
  MutablePrimArray s a
marr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len')
  MutablePrimArray (PrimState (ST s)) a
-> Int -> PrimArray a -> Int -> Int -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
0 PrimArray a
arr Int
off Int
len
  MutablePrimArray (PrimState (ST s)) a
-> Int -> PrimArray a -> Int -> Int -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
len PrimArray a
arr' Int
off' Int
len'
  MutablePrimArray s a -> ST s (MutablePrimArray s a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutablePrimArray s a
marr

pick :: Prim a => Int -> [Int] -> PrimVec a -> PrimVec a
pick :: Int -> [Int] -> PrimVec a -> PrimVec a
pick Int
len' [Int]
ns (PrimVec Int
off Int
_ PrimArray a
arr) = Int -> Int -> PrimArray a -> PrimVec a
forall a. Int -> Int -> PrimArray a -> PrimVec a
PrimVec Int
0 Int
len' (PrimArray a -> PrimVec a) -> PrimArray a -> PrimVec a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray do
  MutablePrimArray s a
marr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len'
  [(Int, Int)] -> ((Int, Int) -> ST s ()) -> ST s ()
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Int]
ns) \(Int
new, Int
old) ->
    MutablePrimArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
new (PrimArray a -> Int -> a
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray a
arr (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
old))
  MutablePrimArray s a -> ST s (MutablePrimArray s a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutablePrimArray s a
marr

update :: Prim a => Int -> a -> PrimVec a -> PrimVec a
update :: Int -> a -> PrimVec a -> PrimVec a
update Int
n a
x (PrimVec Int
off Int
len PrimArray a
arr) = Int -> Int -> PrimArray a -> PrimVec a
forall a. Int -> Int -> PrimArray a -> PrimVec a
PrimVec Int
0 Int
len (PrimArray a -> PrimVec a) -> PrimArray a -> PrimVec a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
forall a. (forall s. ST s (MutablePrimArray s a)) -> PrimArray a
runPrimArray do
  MutablePrimArray s a
marr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  MutablePrimArray (PrimState (ST s)) a
-> Int -> PrimArray a -> Int -> Int -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
0 PrimArray a
arr Int
off Int
len
  MutablePrimArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall a (m :: Type -> Type).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
n a
x
  MutablePrimArray s a -> ST s (MutablePrimArray s a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure MutablePrimArray s a
marr