-- | The functions in this module perform no bounds checking.

module Util.Internal.Array
    ( Array
    , empty
    , singleton
    , snoc
    , index
    , head, last
    , adjust
    , take
    , fromList2
    , fromTail
    , toTail
    ) where

import qualified Data.List.NonEmpty as L
import Prelude hiding (head, last, take)

import Data.Primitive.SmallArray

type Array a = SmallArray a

empty :: Array a
empty :: Array a
empty = Array a
forall a. Monoid a => a
mempty

singleton :: a -> Array a
singleton :: a -> Array a
singleton a
x = (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> Array a)
-> (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a b. (a -> b) -> a -> b
$ Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
1 a
x

snoc :: Array a -> a -> Array a
snoc :: Array a -> a -> Array a
snoc Array a
arr a
x = (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> Array a)
-> (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
    let size :: Int
size = Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr
    SmallMutableArray s a
arr' <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
x
    SmallMutableArray (PrimState (ST s)) a
-> Int -> Array a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
arr' Int
0 Array a
arr Int
0 Int
size
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
arr'

index :: Int -> Array a -> a
index :: Int -> Array a -> a
index = (Array a -> Int -> a) -> Int -> Array a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Array a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray

head :: Array a -> a
head :: Array a -> a
head Array a
arr = Array a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray Array a
arr Int
0

last :: Array a -> a
last :: Array a -> a
last Array a
arr = Array a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray Array a
arr (Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | Update the element at the specified index.
adjust :: Int -> (a -> a) -> Array a -> Array a
adjust :: Int -> (a -> a) -> Array a -> Array a
adjust Int
i a -> a
f Array a
arr = (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> Array a)
-> (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
arr' <- Array a
-> Int -> Int -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray Array a
arr Int
0 (Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr)
    let x :: a
x = Array a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray Array a
arr Int
i
    SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
arr' Int
i (a -> a
f a
x)
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
arr'
{-# INLINE adjust #-}

take :: Int -> Array a -> Array a
take :: Int -> Array a -> Array a
take Int
n Array a
arr = Array a -> Int -> Int -> Array a
forall a. SmallArray a -> Int -> Int -> SmallArray a
cloneSmallArray Array a
arr Int
0 Int
n

fromList2 :: a -> a -> Array a
fromList2 :: a -> a -> Array a
fromList2 a
x a
y = (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> Array a)
-> (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
arr <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
2 a
x
    SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
arr Int
1 a
y
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
arr

-- | Convert a full tail into an array.
--
-- > fromTail = A.fromListN tailSize . reverse . toList
fromTail :: Int -> L.NonEmpty a -> Array a
fromTail :: Int -> NonEmpty a -> Array a
fromTail Int
size (a
x L.:| [a]
xs) = (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> Array a)
-> (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
arr <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
size a
x
    let loop :: Int -> [a] -> ST s ()
loop Int
_ [] = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        loop Int
i (a
y : [a]
ys) = SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
arr Int
i a
y ST s () -> ST s () -> ST s ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> [a] -> ST s ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
ys
    Int -> [a] -> ST s ()
loop (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) [a]
xs
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
arr

-- | Convert an array into a tail.
--
-- > toTail = L.fromList . reverse . toList
toTail :: Array a -> L.NonEmpty a
toTail :: Array a -> NonEmpty a
toTail = [a] -> NonEmpty a
forall a. [a] -> NonEmpty a
L.fromList ([a] -> NonEmpty a) -> (Array a -> [a]) -> Array a -> NonEmpty a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> a -> [a]) -> [a] -> Array a -> [a]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((a -> [a] -> [a]) -> [a] -> a -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) []