module VectorExtras.Generic where

import Data.Vector.Generic
import qualified Data.Vector.Generic.Mutable as Mut
import qualified VectorExtras.Generic.Mutable as Mut
import VectorExtras.Prelude

-- |
-- Notice: It is your responsibility to ensure that the indices
-- in the assoc list are within bounds.
fromAssocListWithGen :: Vector v a => Int -> (Int -> a) -> [(Int, a)] -> v a
fromAssocListWithGen :: Int -> (Int -> a) -> [(Int, a)] -> v a
fromAssocListWithGen Int
size Int -> a
genDefVal [(Int, a)]
list =
  (forall s. ST s (v a)) -> v a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (v a)) -> v a) -> (forall s. ST s (v a)) -> v a
forall a b. (a -> b) -> a -> b
$ do
    Mutable v s a
mv <- Int -> (Int -> a) -> ST s (Mutable v (PrimState (ST s)) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> (Int -> a) -> m (v (PrimState m) a)
Mut.generate Int
size Int -> a
genDefVal
    Mutable v s a -> [(Int, a)] -> ST s ()
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> [(Int, a)] -> ST s ()
Mut.writeAssocList Mutable v s a
mv [(Int, a)]
list
    Mutable v (PrimState (ST s)) a -> ST s (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
unsafeFreeze Mutable v s a
Mutable v (PrimState (ST s)) a
mv

-- |
-- Notice: It is your responsibility to ensure that the indices
-- in the assoc list are within bounds.
fromAssocListWithDef :: Vector v a => Int -> a -> [(Int, a)] -> v a
fromAssocListWithDef :: Int -> a -> [(Int, a)] -> v a
fromAssocListWithDef Int
size a
defVal [(Int, a)]
list =
  (forall s. ST s (v a)) -> v a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (v a)) -> v a) -> (forall s. ST s (v a)) -> v a
forall a b. (a -> b) -> a -> b
$ do
    Mutable v s a
mv <- Int -> a -> ST s (Mutable v (PrimState (ST s)) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
Mut.replicate Int
size a
defVal
    Mutable v s a -> [(Int, a)] -> ST s ()
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> [(Int, a)] -> ST s ()
Mut.writeAssocList Mutable v s a
mv [(Int, a)]
list
    Mutable v (PrimState (ST s)) a -> ST s (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
unsafeFreeze Mutable v s a
Mutable v (PrimState (ST s)) a
mv

-- |
-- >>> fromReverseListN 3 [1,2,3] :: Data.Vector.Vector Int
-- [3,2,1]
{-# INLINE fromReverseListN #-}
fromReverseListN :: Vector v a => Int -> [a] -> v a
fromReverseListN :: Int -> [a] -> v a
fromReverseListN Int
size [a]
list =
  Int -> (forall s. Mutable v s a -> ST s ()) -> v a
forall (v :: * -> *) a.
Vector v a =>
Int -> (forall s. Mutable v s a -> ST s ()) -> v a
initialized Int
size ((forall s. Mutable v s a -> ST s ()) -> v a)
-> (forall s. Mutable v s a -> ST s ()) -> v a
forall a b. (a -> b) -> a -> b
$ \Mutable v s a
mv -> Mutable v s a -> Int -> [a] -> ST s ()
forall (v :: * -> * -> *) a s.
MVector v a =>
v s a -> Int -> [a] -> ST s ()
Mut.writeListInReverseOrderStartingFrom Mutable v s a
mv (Int -> Int
forall a. Enum a => a -> a
pred Int
size) [a]
list

{-# INLINE initialized #-}
initialized :: Vector v a => Int -> (forall s. Mutable v s a -> ST s ()) -> v a
initialized :: Int -> (forall s. Mutable v s a -> ST s ()) -> v a
initialized Int
size forall s. Mutable v s a -> ST s ()
initialize = (forall s. ST s (v a)) -> v a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (v a)) -> v a) -> (forall s. ST s (v a)) -> v a
forall a b. (a -> b) -> a -> b
$ do
  Mutable v s a
mv <- Int -> ST s (Mutable v (PrimState (ST s)) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
Mut.unsafeNew Int
size
  Mutable v s a -> ST s ()
forall s. Mutable v s a -> ST s ()
initialize Mutable v s a
mv
  Mutable v (PrimState (ST s)) a -> ST s (v a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
unsafeFreeze Mutable v s a
Mutable v (PrimState (ST s)) a
mv