{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}

{- |
= Finite vectors

The @'Vector' a@ type represents a finite vector (or dynamic array) of elements of type @a@.
A 'Vector' is strict in its spine.

The class instances are based on those for lists.

This module should be imported qualified, to avoid name clashes with the 'Prelude'.

== Performance

The worst case running time complexities are given, with /n/ referring the the number of elements in the vector.
A 'Vector' is particularly efficient for applications that require a lot of indexing and updates.
All logarithms are base 16, which means that /O(log n)/ behaves more like /O(1)/ in practice.

For a similar container with efficient concatenation and splitting, but slower indexing and updates,
see [Seq](https://hackage.haskell.org/package/containers/docs/Data-Sequence.html) from the
[containers](https://hackage.haskell.org/package/containers) package.

== Warning

The length of a 'Vector' must not exceed @'maxBound' :: 'Int'@.
Violation of this condition is not detected and if the length limit is exceeded, the behaviour of the vector is undefined.

== Implementation

The implementation of 'Vector' uses array mapped tries. For a good explanation,
see [this blog post](https://hypirion.com/musings/understanding-persistent-vector-pt-1).
-}

module Data.AMT
    ( Vector
    -- * Construction
    , empty, singleton, fromList
    , fromFunction
    , replicate, replicateA
    , unfoldr, unfoldl, iterateN
    , (<|), (|>), (><)
    -- * Deconstruction/Subranges
    , viewl, viewr
    , head, last
    , take
    -- * Indexing
    , lookup, index
    , (!?), (!)
    , update
    , adjust
    -- * Transformations
    , map, mapWithIndex
    , traverseWithIndex
    , indexed
    -- * Folds
    , foldMapWithIndex
    , foldlWithIndex, foldrWithIndex
    , foldlWithIndex', foldrWithIndex'
    -- * Zipping/Unzipping
    , zip, zipWith
    , zip3, zipWith3
    , unzip, unzip3
    -- * To Lists
    , toIndexedList
    ) where

import Control.Applicative (Alternative)
import qualified Control.Applicative as Applicative
import Control.Monad (MonadPlus(..))
#if !(MIN_VERSION_base(4,13,0))
import Control.Monad.Fail (MonadFail(..))
#endif
import Control.Monad.Zip (MonadZip(..))

import Data.Bits
import Data.Foldable (foldl', toList)
import Data.Functor.Classes
import Data.List.NonEmpty (NonEmpty(..), (!!))
import qualified Data.List.NonEmpty as L
import Data.Maybe (fromMaybe)
#if !(MIN_VERSION_base(4,11,0))
import Data.Semigroup (Semigroup((<>)))
#endif
import Data.String (IsString)
import Data.Traversable (mapAccumL)
import GHC.Exts (IsList)
import qualified GHC.Exts as Exts
import Prelude hiding ((!!), head, last, lookup, map, replicate, tail, take, unzip, unzip3, zip, zipWith, zip3, zipWith3)
import Text.Read (Lexeme(Ident), lexP, parens, prec, readPrec)

import Control.DeepSeq (NFData(..))

import qualified Util.Internal.Array as A
import Util.Internal.Indexed (Indexed(..), evalIndexed)

infixr 5 ><
infixr 5 <|
infixl 5 |>

data Tree a
    = Internal !(A.Array (Tree a)) -- never empty
    | Leaf !(A.Array a)

-- | An array mapped trie.
data Vector a
    = Empty
    | Root
        {-# UNPACK #-} !Int  -- ^ size
        {-# UNPACK #-} !Int  -- ^ offset (number of elements in the tree)
        {-# UNPACK #-} !Int  -- ^ height (of the tree)
        !(Tree a)  -- ^ tree
        !(NonEmpty a)  -- ^ tail (reversed)

instance NFData a => NFData (Tree a) where
    rnf :: Tree a -> ()
rnf (Internal Array (Tree a)
v) = Array (Tree a) -> ()
forall a. NFData a => a -> ()
rnf Array (Tree a)
v
    rnf (Leaf Array a
v) = Array a -> ()
forall a. NFData a => a -> ()
rnf Array a
v

errorNegativeLength :: String -> a
errorNegativeLength :: String -> a
errorNegativeLength String
s = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"AMT." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": expected a nonnegative length"

-- | The number of bits used per level.
bits :: Int
bits :: Int
bits = Int
4
{-# INLINE bits #-}

-- | The maximum size of the tail.
tailSize :: Int
tailSize :: Int
tailSize = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
bits

-- | The mask used to extract the index into the array.
mask :: Int
mask :: Int
mask = Int
tailSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

instance Show1 Vector where
    liftShowsPrec :: (Int -> a -> String -> String)
-> ([a] -> String -> String) -> Int -> Vector a -> String -> String
liftShowsPrec Int -> a -> String -> String
sp [a] -> String -> String
sl Int
p Vector a
v = (Int -> [a] -> String -> String)
-> String -> Int -> [a] -> String -> String
forall a.
(Int -> a -> String -> String)
-> String -> Int -> a -> String -> String
showsUnaryWith ((Int -> a -> String -> String)
-> ([a] -> String -> String) -> Int -> [a] -> String -> String
forall (f :: * -> *) a.
Show1 f =>
(Int -> a -> String -> String)
-> ([a] -> String -> String) -> Int -> f a -> String -> String
liftShowsPrec Int -> a -> String -> String
sp [a] -> String -> String
sl) String
"fromList" Int
p (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v)

instance Show a => Show (Vector a) where
    showsPrec :: Int -> Vector a -> String -> String
showsPrec = Int -> Vector a -> String -> String
forall (f :: * -> *) a.
(Show1 f, Show a) =>
Int -> f a -> String -> String
showsPrec1
    {-# INLINE showsPrec #-}

instance Read1 Vector where
    liftReadsPrec :: (Int -> ReadS a) -> ReadS [a] -> Int -> ReadS (Vector a)
liftReadsPrec Int -> ReadS a
rp ReadS [a]
rl = (String -> ReadS (Vector a)) -> Int -> ReadS (Vector a)
forall a. (String -> ReadS a) -> Int -> ReadS a
readsData ((String -> ReadS (Vector a)) -> Int -> ReadS (Vector a))
-> (String -> ReadS (Vector a)) -> Int -> ReadS (Vector a)
forall a b. (a -> b) -> a -> b
$ (Int -> ReadS [a])
-> String -> ([a] -> Vector a) -> String -> ReadS (Vector a)
forall a t.
(Int -> ReadS a) -> String -> (a -> t) -> String -> ReadS t
readsUnaryWith ((Int -> ReadS a) -> ReadS [a] -> Int -> ReadS [a]
forall (f :: * -> *) a.
Read1 f =>
(Int -> ReadS a) -> ReadS [a] -> Int -> ReadS (f a)
liftReadsPrec Int -> ReadS a
rp ReadS [a]
rl) String
"fromList" [a] -> Vector a
forall a. [a] -> Vector a
fromList

instance Read a => Read (Vector a) where
    readPrec :: ReadPrec (Vector a)
readPrec = ReadPrec (Vector a) -> ReadPrec (Vector a)
forall a. ReadPrec a -> ReadPrec a
parens (ReadPrec (Vector a) -> ReadPrec (Vector a))
-> ReadPrec (Vector a) -> ReadPrec (Vector a)
forall a b. (a -> b) -> a -> b
$ Int -> ReadPrec (Vector a) -> ReadPrec (Vector a)
forall a. Int -> ReadPrec a -> ReadPrec a
prec Int
10 (ReadPrec (Vector a) -> ReadPrec (Vector a))
-> ReadPrec (Vector a) -> ReadPrec (Vector a)
forall a b. (a -> b) -> a -> b
$ do
        Ident String
"fromList" <- ReadPrec Lexeme
lexP
        [a]
xs <- ReadPrec [a]
forall a. Read a => ReadPrec a
readPrec
        Vector a -> ReadPrec (Vector a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> Vector a
forall a. [a] -> Vector a
fromList [a]
xs)

instance Eq1 Vector where
    liftEq :: (a -> b -> Bool) -> Vector a -> Vector b -> Bool
liftEq a -> b -> Bool
f Vector a
v1 Vector b
v2 = Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector b -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector b
v2 Bool -> Bool -> Bool
&& (a -> b -> Bool) -> [a] -> [b] -> Bool
forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq a -> b -> Bool
f (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector b
v2)

instance Eq a => Eq (Vector a) where
    == :: Vector a -> Vector a -> Bool
(==) = Vector a -> Vector a -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1

instance Ord1 Vector where
    liftCompare :: (a -> b -> Ordering) -> Vector a -> Vector b -> Ordering
liftCompare a -> b -> Ordering
f Vector a
v1 Vector b
v2 = (a -> b -> Ordering) -> [a] -> [b] -> Ordering
forall (f :: * -> *) a b.
Ord1 f =>
(a -> b -> Ordering) -> f a -> f b -> Ordering
liftCompare a -> b -> Ordering
f (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector b
v2)

instance Ord a => Ord (Vector a) where
    compare :: Vector a -> Vector a -> Ordering
compare = Vector a -> Vector a -> Ordering
forall (f :: * -> *) a. (Ord1 f, Ord a) => f a -> f a -> Ordering
compare1

instance Semigroup (Vector a) where
    <> :: Vector a -> Vector a -> Vector a
(<>) = Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
(><)

instance Monoid (Vector a) where
    mempty :: Vector a
mempty = Vector a
forall a. Vector a
empty

    mappend :: Vector a -> Vector a -> Vector a
mappend = Vector a -> Vector a -> Vector a
forall a. Semigroup a => a -> a -> a
(<>)

instance Foldable Vector where
    foldr :: (a -> b -> b) -> b -> Vector a -> b
foldr a -> b -> b
f b
acc = Vector a -> b
go
      where
        go :: Vector a -> b
go Vector a
Empty = b
acc
        go (Root Int
_ Int
_ Int
_ Tree a
tree NonEmpty a
tail) = Tree a -> b -> b
foldrTree Tree a
tree ((a -> b -> b) -> b -> NonEmpty a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f b
acc (NonEmpty a -> NonEmpty a
forall a. NonEmpty a -> NonEmpty a
L.reverse NonEmpty a
tail))

        foldrTree :: Tree a -> b -> b
foldrTree (Internal Array (Tree a)
v) b
acc' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tree a -> b -> b
foldrTree b
acc' Array (Tree a)
v
        foldrTree (Leaf Array a
v) b
acc' = (a -> b -> b) -> b -> Array a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f b
acc' Array a
v
    {-# INLINE foldr #-}

    null :: Vector a -> Bool
null Vector a
Empty = Bool
True
    null Root{} = Bool
False
    {-# INLINE null #-}

    length :: Vector a -> Int
length Vector a
Empty = Int
0
    length (Root Int
s Int
_ Int
_ Tree a
_ NonEmpty a
_) = Int
s
    {-# INLINE length #-}

instance Functor Vector where
    fmap :: (a -> b) -> Vector a -> Vector b
fmap = (a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
map

instance Traversable Vector where
    traverse :: (a -> f b) -> Vector a -> f (Vector b)
traverse a -> f b
f = Vector a -> f (Vector b)
go
      where
        go :: Vector a -> f (Vector b)
go Vector a
Empty = Vector b -> f (Vector b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector b
forall a. Vector a
empty
        go (Root Int
s Int
offset Int
h Tree a
tree (a
x :| [a]
tail)) =
            Int -> Int -> Int -> Tree b -> NonEmpty b -> Vector b
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
s Int
offset Int
h (Tree b -> NonEmpty b -> Vector b)
-> f (Tree b) -> f (NonEmpty b -> Vector b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree a -> f (Tree b)
traverseTree Tree a
tree f (NonEmpty b -> Vector b) -> f (NonEmpty b) -> f (Vector b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ((b -> [b] -> NonEmpty b) -> [b] -> b -> NonEmpty b
forall a b c. (a -> b -> c) -> b -> a -> c
flip b -> [b] -> NonEmpty b
forall a. a -> [a] -> NonEmpty a
(:|) ([b] -> b -> NonEmpty b) -> f [b] -> f (b -> NonEmpty b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> f [b]
traverseReverse [a]
tail f (b -> NonEmpty b) -> f b -> f (NonEmpty b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> f b
f a
x)

        traverseReverse :: [a] -> f [b]
traverseReverse [] = [b] -> f [b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
        traverseReverse (a
x : [a]
xs) = (b -> [b] -> [b]) -> [b] -> b -> [b]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:) ([b] -> b -> [b]) -> f [b] -> f (b -> [b])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> f [b]
traverseReverse [a]
xs f (b -> [b]) -> f b -> f [b]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> a -> f b
f a
x

        traverseTree :: Tree a -> f (Tree b)
traverseTree (Internal Array (Tree a)
v) = Array (Tree b) -> Tree b
forall a. Array (Tree a) -> Tree a
Internal (Array (Tree b) -> Tree b) -> f (Array (Tree b)) -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Tree a -> f (Tree b)) -> Array (Tree a) -> f (Array (Tree b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Tree a -> f (Tree b)
traverseTree Array (Tree a)
v
        traverseTree (Leaf Array a
v) = Array b -> Tree b
forall a. Array a -> Tree a
Leaf (Array b -> Tree b) -> f (Array b) -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Array a -> f (Array b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Array a
v
    {-# INLINE traverse #-}

instance IsList (Vector a) where
    type Item (Vector a) = a

    fromList :: [Item (Vector a)] -> Vector a
fromList = [Item (Vector a)] -> Vector a
forall a. [a] -> Vector a
fromList

    toList :: Vector a -> [Item (Vector a)]
toList = Vector a -> [Item (Vector a)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

instance a ~ Char => IsString (Vector a) where
    fromString :: String -> Vector a
fromString = String -> Vector a
forall a. [a] -> Vector a
fromList

instance Applicative Vector where
    pure :: a -> Vector a
pure = a -> Vector a
forall a. a -> Vector a
singleton

    Vector (a -> b)
fs <*> :: Vector (a -> b) -> Vector a -> Vector b
<*> Vector a
xs = (Vector b -> (a -> b) -> Vector b)
-> Vector b -> Vector (a -> b) -> Vector b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Vector b
acc a -> b
f -> Vector b
acc Vector b -> Vector b -> Vector b
forall a. Vector a -> Vector a -> Vector a
>< (a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
map a -> b
f Vector a
xs) Vector b
forall a. Vector a
empty Vector (a -> b)
fs

instance Monad Vector where
    Vector a
xs >>= :: Vector a -> (a -> Vector b) -> Vector b
>>= a -> Vector b
f = (Vector b -> a -> Vector b) -> Vector b -> Vector a -> Vector b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Vector b
acc a
x -> Vector b
acc Vector b -> Vector b -> Vector b
forall a. Vector a -> Vector a -> Vector a
>< a -> Vector b
f a
x) Vector b
forall a. Vector a
empty Vector a
xs

instance Alternative Vector where
    empty :: Vector a
empty = Vector a
forall a. Vector a
empty

    <|> :: Vector a -> Vector a -> Vector a
(<|>) = Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
(><)

instance MonadPlus Vector

instance MonadFail Vector where
    fail :: String -> Vector a
fail String
_ = Vector a
forall a. Vector a
empty

instance MonadZip Vector where
    mzip :: Vector a -> Vector b -> Vector (a, b)
mzip = Vector a -> Vector b -> Vector (a, b)
forall a b. Vector a -> Vector b -> Vector (a, b)
zip

    mzipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
mzipWith = (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith

    munzip :: Vector (a, b) -> (Vector a, Vector b)
munzip = Vector (a, b) -> (Vector a, Vector b)
forall a b. Vector (a, b) -> (Vector a, Vector b)
unzip

instance NFData a => NFData (Vector a) where
    rnf :: Vector a -> ()
rnf Vector a
Empty = ()
    rnf (Root Int
_ Int
_ Int
_ Tree a
tree NonEmpty a
tail) = Tree a -> ()
forall a. NFData a => a -> ()
rnf Tree a
tree () -> () -> ()
`seq` NonEmpty a -> ()
forall a. NFData a => a -> ()
rnf NonEmpty a
tail


-- | /O(1)/. The empty vector.
--
-- > empty = fromList []
empty :: Vector a
empty :: Vector a
empty = Vector a
forall a. Vector a
Empty

-- | /O(1)/. A vector with a single element.
--
-- > singleton x = fromList [x]
singleton :: a -> Vector a
singleton :: a -> Vector a
singleton a
x = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
1 Int
0 Int
0 (Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
forall a. Array a
A.empty) (a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [])

-- | /O(n * log n)/. Create a new vector from a list.
fromList :: [a] -> Vector a
fromList :: [a] -> Vector a
fromList = (Vector a -> a -> Vector a) -> Vector a -> [a] -> Vector a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
(|>) Vector a
forall a. Vector a
empty

-- | Create a new vector of the given length from a function.
fromFunction :: Int -> (Int -> a) -> Vector a
fromFunction :: Int -> (Int -> a) -> Vector a
fromFunction Int
n Int -> a
f = if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then String -> Vector a
forall a. String -> a
errorNegativeLength String
"fromFunction" else Int -> Vector a -> Vector a
go Int
0 Vector a
forall a. Vector a
empty
  where
    go :: Int -> Vector a -> Vector a
go Int
i Vector a
acc
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = Int -> Vector a -> Vector a
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Vector a
acc Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
|> Int -> a
f Int
i)
        | Bool
otherwise = Vector a
acc

-- | /O(n * log n)/. @replicate n x@ is a vector consisting of n copies of x.
replicate :: Int -> a -> Vector a
replicate :: Int -> a -> Vector a
replicate Int
n a
x = if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then String -> Vector a
forall a. String -> a
errorNegativeLength String
"replicate" else Int -> Vector a -> Vector a
go Int
0 Vector a
forall a. Vector a
empty
  where
    go :: Int -> Vector a -> Vector a
go Int
i Vector a
acc
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = Int -> Vector a -> Vector a
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Vector a
acc Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
|> a
x)
        | Bool
otherwise = Vector a
acc

-- | @replicateA@ is an 'Applicative' version of 'replicate'.
replicateA :: Applicative f => Int -> f a -> f (Vector a)
replicateA :: Int -> f a -> f (Vector a)
replicateA Int
n f a
x = if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then String -> f (Vector a)
forall a. String -> a
errorNegativeLength String
"replicateA" else Int -> f (Vector a) -> f (Vector a)
go Int
0 (Vector a -> f (Vector a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector a
forall a. Vector a
empty)
  where
    go :: Int -> f (Vector a) -> f (Vector a)
go Int
i f (Vector a)
acc
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = Int -> f (Vector a) -> f (Vector a)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
(|>) (Vector a -> a -> Vector a) -> f (Vector a) -> f (a -> Vector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Vector a)
acc f (a -> Vector a) -> f a -> f (Vector a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> f a
x)
        | Bool
otherwise = f (Vector a)
acc

-- | /O(n * log n)/. Build a vector from left to right by repeatedly applying a function to a seed value.
unfoldr :: (b -> Maybe (a, b)) -> b -> Vector a
unfoldr :: (b -> Maybe (a, b)) -> b -> Vector a
unfoldr b -> Maybe (a, b)
f = Vector a -> b -> Vector a
go Vector a
forall a. Vector a
empty
  where
    go :: Vector a -> b -> Vector a
go Vector a
v b
acc = case b -> Maybe (a, b)
f b
acc of
        Maybe (a, b)
Nothing -> Vector a
v
        Just (a
x, b
acc') -> Vector a -> b -> Vector a
go (Vector a
v Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
|> a
x) b
acc'
{-# INLINE unfoldr #-}

-- | /O(n * log n)/. Build a vector from right to left by repeatedly applying a function to a seed value.
unfoldl :: (b -> Maybe (b, a)) -> b -> Vector a
unfoldl :: (b -> Maybe (b, a)) -> b -> Vector a
unfoldl b -> Maybe (b, a)
f = b -> Vector a
go
  where
    go :: b -> Vector a
go b
acc = case b -> Maybe (b, a)
f b
acc of
        Maybe (b, a)
Nothing -> Vector a
forall a. Vector a
empty
        Just (b
acc', a
x) -> b -> Vector a
go b
acc' Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
|> a
x
{-# INLINE unfoldl #-}

-- | Constructs a vector by repeatedly applying a function to a seed value.
iterateN :: Int -> (a -> a) -> a -> Vector a
iterateN :: Int -> (a -> a) -> a -> Vector a
iterateN Int
n a -> a
f a
x = if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 then String -> Vector a
forall a. String -> a
errorNegativeLength String
"iterateN" else Int -> a -> Vector a -> Vector a
go Int
0 a
x Vector a
forall a. Vector a
empty
  where
    go :: Int -> a -> Vector a -> Vector a
go Int
i a
y Vector a
acc
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = Int -> a -> Vector a -> Vector a
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (a -> a
f a
y) (Vector a
acc Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
|> a
y)
        | Bool
otherwise = Vector a
acc

-- | /O(n * log n)/. Add an element to the left end of the vector.
(<|) :: a -> Vector a -> Vector a
a
x <| :: a -> Vector a -> Vector a
<| Vector a
v = [a] -> Vector a
forall a. [a] -> Vector a
fromList ([a] -> Vector a) -> [a] -> Vector a
forall a b. (a -> b) -> a -> b
$ a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v

-- | /O(n * log n)/. The first element and the vector without the first element or 'Nothing' if the vector is empty.
viewl :: Vector a -> Maybe (a, Vector a)
viewl :: Vector a -> Maybe (a, Vector a)
viewl Vector a
v = case Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v of
    [] -> Maybe (a, Vector a)
forall a. Maybe a
Nothing
    a
x : [a]
xs -> (a, Vector a) -> Maybe (a, Vector a)
forall a. a -> Maybe a
Just (a
x, [a] -> Vector a
forall a. [a] -> Vector a
fromList [a]
xs)

-- | /O(log n)/. Add an element to the right end of the vector.
(|>) :: Vector a -> a -> Vector a
Vector a
Empty |> :: Vector a -> a -> Vector a
|> a
x = a -> Vector a
forall a. a -> Vector a
singleton a
x
Root Int
s Int
offset Int
h Tree a
tree NonEmpty a
tail |> a
x
    | Int
s Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
offset Int
h Tree a
tree (a
x a -> NonEmpty a -> NonEmpty a
forall a. a -> NonEmpty a -> NonEmpty a
L.<| NonEmpty a
tail)
    -- tail is full
    | Int
offset Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
s Int
h (Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty a -> Array a
forall a. Int -> NonEmpty a -> Array a
A.fromTail Int
tailSize NonEmpty a
tail) (a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [])
    | Int
offset Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
s (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Internal (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Tree a -> Tree a -> Array (Tree a)
forall a. a -> a -> Array a
A.fromList2 Tree a
tree (Int -> Tree a
newPath Int
h)) (a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [])
    | Bool
otherwise = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
s Int
h (Int -> Tree a -> Tree a
insertTail (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) Tree a
tree) (a
x a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [])
  where
    -- create a new path from the old tail
    newPath :: Int -> Tree a
newPath Int
0 = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty a -> Array a
forall a. Int -> NonEmpty a -> Array a
A.fromTail Int
tailSize NonEmpty a
tail
    newPath Int
h = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Internal (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Tree a -> Array (Tree a)
forall a. a -> Array a
A.singleton (Int -> Tree a
newPath (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))

    insertTail :: Int -> Tree a -> Tree a
insertTail Int
sh (Internal Array (Tree a)
v)
        | Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
v = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Internal (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> (Tree a -> Tree a) -> Array (Tree a) -> Array (Tree a)
forall a. Int -> (a -> a) -> Array a -> Array a
A.adjust Int
idx (Int -> Tree a -> Tree a
insertTail (Int
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bits)) Array (Tree a)
v
        | Bool
otherwise = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Internal (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.snoc Array (Tree a)
v (Int -> Tree a
newPath (Int
sh Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
      where
        idx :: Int
idx = Int
offset Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
sh Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask
    insertTail Int
_ (Leaf Array a
_) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty a -> Array a
forall a. Int -> NonEmpty a -> Array a
A.fromTail Int
tailSize NonEmpty a
tail

-- | /O(log n)/. The vector without the last element and the last element or 'Nothing' if the vector is empty.
viewr :: Vector a -> Maybe (Vector a, a)
viewr :: Vector a -> Maybe (Vector a, a)
viewr Vector a
Empty = Maybe (Vector a, a)
forall a. Maybe a
Nothing
viewr (Root Int
s Int
offset Int
h Tree a
tree (a
x :| [a]
tail))
    | Bool -> Bool
not ([a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
tail) = (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
offset Int
h Tree a
tree ([a] -> NonEmpty a
forall a. [a] -> NonEmpty a
L.fromList [a]
tail), a
x)
    | Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Vector a
forall a. Vector a
Empty, a
x)
    | Int
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
tailSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 = (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
0 Int
0 (Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
forall a. Array a
A.empty) (Tree a -> NonEmpty a
forall a. Tree a -> NonEmpty a
getTail Tree a
tree), a
x)
    | Bool
otherwise = (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Vector a -> Vector a
forall a. Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
tailSize) Int
h (Int -> Tree a -> Tree a
initTree (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) Tree a
tree) (Tree a -> NonEmpty a
forall a. Tree a -> NonEmpty a
getTail Tree a
tree), a
x)
  where
    idx :: Int
idx = Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
tailSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

    initTree :: Int -> Tree a -> Tree a
initTree Int
sh (Internal Array (Tree a)
v) =
        let subIndex :: Int
subIndex = Int
idx Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
sh Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask
            new :: Array (Tree a)
new = Int -> Array (Tree a) -> Array (Tree a)
forall a. Int -> Array a -> Array a
A.take (Int
subIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Array (Tree a)
v
        in Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Internal (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> (Tree a -> Tree a) -> Array (Tree a) -> Array (Tree a)
forall a. Int -> (a -> a) -> Array a -> Array a
A.adjust Int
subIndex (Int -> Tree a -> Tree a
initTree (Int
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bits)) Array (Tree a)
new
    initTree Int
_ (Leaf Array a
v) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
v

    getTail :: Tree a -> NonEmpty a
getTail (Internal Array (Tree a)
v) = Tree a -> NonEmpty a
getTail (Array (Tree a) -> Tree a
forall a. Array a -> a
A.last Array (Tree a)
v)
    getTail (Leaf Array a
v) = Array a -> NonEmpty a
forall a. Array a -> NonEmpty a
A.toTail Array a
v

    normalize :: Vector a -> Vector a
normalize (Root Int
s Int
offset Int
h (Internal Array (Tree a)
v) NonEmpty a
tail)
        | Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
s Int
offset (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
v) NonEmpty a
tail
    normalize Vector a
v = Vector a
v

-- | /O(log n)/. The first element in the vector or 'Nothing' if the vector is empty.
head :: Vector a -> Maybe a
head :: Vector a -> Maybe a
head Vector a
Empty = Maybe a
forall a. Maybe a
Nothing
head (Root Int
_ Int
0 Int
_ Tree a
_ NonEmpty a
tail) = a -> Maybe a
forall a. a -> Maybe a
Just (NonEmpty a -> a
forall a. NonEmpty a -> a
L.last NonEmpty a
tail) -- offset 0, all elements are in the tail
head (Root Int
_ Int
_ Int
_ Tree a
tree NonEmpty a
_) = a -> Maybe a
forall a. a -> Maybe a
Just (Tree a -> a
forall p. Tree p -> p
headTree Tree a
tree)
  where
    headTree :: Tree p -> p
headTree (Internal Array (Tree p)
v) = Tree p -> p
headTree (Array (Tree p) -> Tree p
forall a. Array a -> a
A.head Array (Tree p)
v)
    headTree (Leaf Array p
v) = Array p -> p
forall a. Array a -> a
A.head Array p
v

-- | /O(1)/. The last element in the vector or 'Nothing' if the vector is empty.
last :: Vector a -> Maybe a
last :: Vector a -> Maybe a
last Vector a
Empty = Maybe a
forall a. Maybe a
Nothing
last (Root Int
_ Int
_ Int
_ Tree a
_ (a
x :| [a]
_)) = a -> Maybe a
forall a. a -> Maybe a
Just a
x

-- | /O(log n)/. Take the first n elements of the vector or the vector if n is larger than the length of the vector.
-- Returns the empty vector if n is negative.
take :: Int -> Vector a -> Vector a
take :: Int -> Vector a -> Vector a
take Int
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
take Int
n root :: Vector a
root@(Root Int
s Int
offset Int
h Tree a
tree NonEmpty a
tail)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Vector a
forall a. Vector a
Empty
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
s = Vector a
root
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
offset = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
n Int
offset Int
h Tree a
tree ([a] -> NonEmpty a
forall a. [a] -> NonEmpty a
L.fromList ([a] -> NonEmpty a) -> [a] -> NonEmpty a
forall a b. (a -> b) -> a -> b
$ Int -> NonEmpty a -> [a]
forall a. Int -> NonEmpty a -> [a]
L.drop (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) NonEmpty a
tail)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
tailSize = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
n Int
0 Int
0 (Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
forall a. Array a
A.empty) (Int -> Tree a -> NonEmpty a
getTail (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) Tree a
tree)
    | Bool
otherwise =
        let sh :: Int
sh = Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h
        in Vector a -> Vector a
forall a. Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
n ((Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int -> Int
forall a. Bits a => a -> a
complement Int
mask) Int
h (Int -> Tree a -> Tree a
takeTree Int
sh Tree a
tree) (Int -> Tree a -> NonEmpty a
getTail Int
sh Tree a
tree)  -- n - 1 because if 'n .&. mask == 0', we need to subtract tailSize
  where
    idx :: Int
idx = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 -- index of the last element in the new vector

    idx' :: Int
idx' = Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
tailSize

    takeTree :: Int -> Tree a -> Tree a
takeTree Int
sh (Internal Array (Tree a)
v) =
        let subIndex :: Int
subIndex = Int
idx' Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
sh Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask
            new :: Array (Tree a)
new = Int -> Array (Tree a) -> Array (Tree a)
forall a. Int -> Array a -> Array a
A.take (Int
subIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Array (Tree a)
v
        in Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Internal (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> (Tree a -> Tree a) -> Array (Tree a) -> Array (Tree a)
forall a. Int -> (a -> a) -> Array a -> Array a
A.adjust Int
subIndex (Int -> Tree a -> Tree a
takeTree (Int
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bits)) Array (Tree a)
new
    takeTree Int
_ (Leaf Array a
v) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
v

    getTail :: Int -> Tree a -> NonEmpty a
getTail Int
sh (Internal Array (Tree a)
v) = Int -> Tree a -> NonEmpty a
getTail (Int
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bits) (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array a -> a
A.index (Int
idx Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
sh Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask) Array (Tree a)
v)
    getTail Int
_ (Leaf Array a
v) = Array a -> NonEmpty a
forall a. Array a -> NonEmpty a
A.toTail (Array a -> NonEmpty a) -> Array a -> NonEmpty a
forall a b. (a -> b) -> a -> b
$ Int -> Array a -> Array a
forall a. Int -> Array a -> Array a
A.take (Int
idx Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Array a
v

    normalize :: Vector a -> Vector a
normalize (Root Int
s Int
offset Int
h (Internal Array (Tree a)
v) NonEmpty a
tail)
        | Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
s Int
offset (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
v) NonEmpty a
tail
    normalize Vector a
v = Vector a
v

-- | /O(log n)/. The element at the index or 'Nothing' if the index is out of range.
lookup :: Int -> Vector a -> Maybe a
lookup :: Int -> Vector a -> Maybe a
lookup Int
_ Vector a
Empty = Maybe a
forall a. Maybe a
Nothing
lookup Int
i (Root Int
s Int
offset Int
h Tree a
tree NonEmpty a
tail)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
s = Maybe a
forall a. Maybe a
Nothing  -- index out of range
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
offset = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Int -> Tree a -> a
lookupTree (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) Tree a
tree
    | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ NonEmpty a
tail NonEmpty a -> Int -> a
forall a. NonEmpty a -> Int -> a
!! (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    lookupTree :: Int -> Tree a -> a
lookupTree Int
sh (Internal Array (Tree a)
v) = Int -> Tree a -> a
lookupTree (Int
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bits) (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array a -> a
A.index (Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
sh Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask) Array (Tree a)
v)
    lookupTree Int
_ (Leaf Array a
v) = Int -> Array a -> a
forall a. Int -> Array a -> a
A.index (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask) Array a
v

-- | /O(log n)/. The element at the index. Calls 'error' if the index is out of range.
index :: Int -> Vector a -> a
index :: Int -> Vector a -> a
index Int
i = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe (String -> a
forall a. HasCallStack => String -> a
error String
"AMT.index: index out of range") (Maybe a -> a) -> (Vector a -> Maybe a) -> Vector a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector a -> Maybe a
forall a. Int -> Vector a -> Maybe a
lookup Int
i

-- | /O(log n)/. Flipped version of 'lookup'.
(!?) :: Vector a -> Int -> Maybe a
!? :: Vector a -> Int -> Maybe a
(!?) = (Int -> Vector a -> Maybe a) -> Vector a -> Int -> Maybe a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Vector a -> Maybe a
forall a. Int -> Vector a -> Maybe a
lookup
{-# INLINE (!?) #-}

-- | /O(log n)/. Flipped version of 'index'.
(!) :: Vector a -> Int -> a
(!) = (Int -> Vector a -> a) -> Vector a -> Int -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Vector a -> a
forall a. Int -> Vector a -> a
index
{-# INLINE (!) #-}

-- | /O(log n)/. Update the element at the index with a new element.
-- Returns the original vector if the index is out of range.
update :: Int -> a -> Vector a -> Vector a
update :: Int -> a -> Vector a -> Vector a
update Int
i a
x = Int -> (a -> a) -> Vector a -> Vector a
forall a. Int -> (a -> a) -> Vector a -> Vector a
adjust Int
i (a -> a -> a
forall a b. a -> b -> a
const a
x)
{-# INLINE update #-}

-- | /O(log n)/. Adjust the element at the index by applying the function to it.
-- Returns the original vector if the index is out of range.
adjust :: Int -> (a -> a) -> Vector a -> Vector a
adjust :: Int -> (a -> a) -> Vector a -> Vector a
adjust Int
_ a -> a
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
adjust Int
i a -> a
f root :: Vector a
root@(Root Int
s Int
offset Int
h Tree a
tree NonEmpty a
tail)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
s = Vector a
root  -- index out of range
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
offset = Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
s Int
offset Int
h (Int -> Tree a -> Tree a
adjustTree (Int
bits Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) Tree a
tree) NonEmpty a
tail
    | Bool
otherwise = let ([a]
l, a
x : [a]
r) = Int -> NonEmpty a -> ([a], [a])
forall a. Int -> NonEmpty a -> ([a], [a])
L.splitAt (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) NonEmpty a
tail in Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
s Int
offset Int
h Tree a
tree ([a] -> NonEmpty a
forall a. [a] -> NonEmpty a
L.fromList ([a] -> NonEmpty a) -> [a] -> NonEmpty a
forall a b. (a -> b) -> a -> b
$ [a]
l [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ (a -> a
f a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
r))
  where
    adjustTree :: Int -> Tree a -> Tree a
adjustTree Int
sh (Internal Array (Tree a)
v) =
        let idx :: Int
idx = Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
sh Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask
        in Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Internal (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> (Tree a -> Tree a) -> Array (Tree a) -> Array (Tree a)
forall a. Int -> (a -> a) -> Array a -> Array a
A.adjust Int
idx (Int -> Tree a -> Tree a
adjustTree (Int
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bits)) Array (Tree a)
v
    adjustTree Int
_ (Leaf Array a
v) =
        let idx :: Int
idx = Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
mask
        in Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> (a -> a) -> Array a -> Array a
forall a. Int -> (a -> a) -> Array a -> Array a
A.adjust Int
idx a -> a
f Array a
v

-- | /O(m * log n)/. Concatenate two vectors.
(><) :: Vector a -> Vector a -> Vector a
Vector a
Empty >< :: Vector a -> Vector a -> Vector a
>< Vector a
v = Vector a
v
Vector a
v >< Vector a
Empty = Vector a
v
Vector a
v1 >< Vector a
v2 = (Vector a -> a -> Vector a) -> Vector a -> Vector a -> Vector a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
(|>) Vector a
v1 Vector a
v2
{-# INLINE (><) #-}

-- | /O(n)/. Map a function over the vector.
map :: (a -> b) -> Vector a -> Vector b
map :: (a -> b) -> Vector a -> Vector b
map a -> b
_ Vector a
Empty = Vector b
forall a. Vector a
Empty
map a -> b
f (Root Int
s Int
offset Int
h Tree a
tree NonEmpty a
tail) = Int -> Int -> Int -> Tree b -> NonEmpty b -> Vector b
forall a. Int -> Int -> Int -> Tree a -> NonEmpty a -> Vector a
Root Int
s Int
offset Int
h (Tree a -> Tree b
mapTree Tree a
tree) ((a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f NonEmpty a
tail)
  where
    mapTree :: Tree a -> Tree b
mapTree (Internal Array (Tree a)
v) = Array (Tree b) -> Tree b
forall a. Array (Tree a) -> Tree a
Internal ((Tree a -> Tree b) -> Array (Tree a) -> Array (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Tree a -> Tree b
mapTree Array (Tree a)
v)
    mapTree (Leaf Array a
v) = Array b -> Tree b
forall a. Array a -> Tree a
Leaf ((a -> b) -> Array a -> Array b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Array a
v)

-- | /O(n)/. Map a function that has access to the index of an element over the vector.
mapWithIndex :: (Int -> a -> b) -> Vector a -> Vector b
mapWithIndex :: (Int -> a -> b) -> Vector a -> Vector b
mapWithIndex Int -> a -> b
f = (Int, Vector b) -> Vector b
forall a b. (a, b) -> b
snd ((Int, Vector b) -> Vector b)
-> (Vector a -> (Int, Vector b)) -> Vector a -> Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> a -> (Int, b)) -> Int -> Vector a -> (Int, Vector b)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL (\Int
i a
x -> Int
i Int -> (Int, b) -> (Int, b)
`seq` (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int -> a -> b
f Int
i a
x)) Int
0
{-# INLINE mapWithIndex #-}

-- | /O(n)/. Fold the values in the vector, using the given monoid.
foldMapWithIndex :: Monoid m => (Int -> a -> m) -> Vector a -> m
foldMapWithIndex :: (Int -> a -> m) -> Vector a -> m
foldMapWithIndex Int -> a -> m
f = (Int -> a -> m -> m) -> m -> Vector a -> m
forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
foldrWithIndex (\Int
i -> m -> m -> m
forall a. Monoid a => a -> a -> a
mappend (m -> m -> m) -> (a -> m) -> a -> m -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a -> m
f Int
i) m
forall a. Monoid a => a
mempty
{-# INLINE foldMapWithIndex #-}

-- | /O(n)/. Fold using the given left-associative function that has access to the index of an element.
foldlWithIndex :: (b -> Int -> a -> b) -> b -> Vector a -> b
foldlWithIndex :: (b -> Int -> a -> b) -> b -> Vector a -> b
foldlWithIndex b -> Int -> a -> b
f b
acc Vector a
v = ((Int -> b) -> a -> Int -> b) -> (Int -> b) -> Vector a -> Int -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Int -> b) -> a -> Int -> b
f' (b -> Int -> b
forall a b. a -> b -> a
const b
acc) Vector a
v (Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    f' :: (Int -> b) -> a -> Int -> b
f' Int -> b
g a
x Int
i = Int
i Int -> b -> b
`seq` b -> Int -> a -> b
f (Int -> b
g (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Int
i a
x
{-# INLINE foldlWithIndex #-}

-- | /O(n)/. Fold using the given right-associative function that has access to the index of an element.
foldrWithIndex :: (Int -> a -> b -> b) -> b -> Vector a -> b
foldrWithIndex :: (Int -> a -> b -> b) -> b -> Vector a -> b
foldrWithIndex Int -> a -> b -> b
f b
acc Vector a
v = (a -> (Int -> b) -> Int -> b) -> (Int -> b) -> Vector a -> Int -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (Int -> b) -> Int -> b
f' (b -> Int -> b
forall a b. a -> b -> a
const b
acc) Vector a
v Int
0
  where
    f' :: a -> (Int -> b) -> Int -> b
f' a
x Int -> b
g Int
i = Int
i Int -> b -> b
`seq` Int -> a -> b -> b
f Int
i a
x (Int -> b
g (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
{-# INLINE foldrWithIndex #-}

-- | /O(n)/. A strict version of 'foldlWithIndex'.
-- Each application of the function is evaluated before using the result in the next application.
foldlWithIndex' :: (b -> Int -> a -> b) -> b -> Vector a -> b
foldlWithIndex' :: (b -> Int -> a -> b) -> b -> Vector a -> b
foldlWithIndex' b -> Int -> a -> b
f b
acc Vector a
v = (Int -> a -> (b -> b) -> b -> b) -> (b -> b) -> Vector a -> b -> b
forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
foldrWithIndex Int -> a -> (b -> b) -> b -> b
f' b -> b
forall a. a -> a
id Vector a
v b
acc
  where
    f' :: Int -> a -> (b -> b) -> b -> b
f' Int
i a
x b -> b
k b
z = b -> b
k (b -> b) -> b -> b
forall a b. (a -> b) -> a -> b
$! b -> Int -> a -> b
f b
z Int
i a
x
{-# INLINE foldlWithIndex' #-}

-- | /O(n)/. A strict version of 'foldrWithIndex'.
-- Each application of the function is evaluated before using the result in the next application.
foldrWithIndex' :: (Int -> a -> b -> b) -> b -> Vector a -> b
foldrWithIndex' :: (Int -> a -> b -> b) -> b -> Vector a -> b
foldrWithIndex' Int -> a -> b -> b
f b
acc Vector a
v = ((b -> b) -> Int -> a -> b -> b) -> (b -> b) -> Vector a -> b -> b
forall b a. (b -> Int -> a -> b) -> b -> Vector a -> b
foldlWithIndex (b -> b) -> Int -> a -> b -> b
f' b -> b
forall a. a -> a
id Vector a
v b
acc
  where
    f' :: (b -> b) -> Int -> a -> b -> b
f' b -> b
k Int
i a
x b
z = b -> b
k (b -> b) -> b -> b
forall a b. (a -> b) -> a -> b
$! Int -> a -> b -> b
f Int
i a
x b
z
{-# INLINE foldrWithIndex' #-}

-- | /O(n)/. Traverse the vector with a function that has access to the index of an element.
traverseWithIndex :: Applicative f => (Int -> a -> f b) -> Vector a -> f (Vector b)
traverseWithIndex :: (Int -> a -> f b) -> Vector a -> f (Vector b)
traverseWithIndex Int -> a -> f b
f Vector a
v = Indexed f (Vector b) -> Int -> f (Vector b)
forall (f :: * -> *) a. Indexed f a -> Int -> f a
evalIndexed ((a -> Indexed f b) -> Vector a -> Indexed f (Vector b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Int -> (f b, Int)) -> Indexed f b
forall (f :: * -> *) a. (Int -> (f a, Int)) -> Indexed f a
Indexed ((Int -> (f b, Int)) -> Indexed f b)
-> (a -> Int -> (f b, Int)) -> a -> Indexed f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int -> (f b, Int)
f') Vector a
v) Int
0
  where
    f' :: a -> Int -> (f b, Int)
f' a
x Int
i = Int
i Int -> (f b, Int) -> (f b, Int)
`seq` (Int -> a -> f b
f Int
i a
x, Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE traverseWithIndex #-}

-- | /O(n)/. Pair each element in the vector with its index.
indexed :: Vector a -> Vector (Int, a)
indexed :: Vector a -> Vector (Int, a)
indexed = (Int -> a -> (Int, a)) -> Vector a -> Vector (Int, a)
forall a b. (Int -> a -> b) -> Vector a -> Vector b
mapWithIndex (,)
{-# INLINE indexed #-}

-- | /O(n)/. Takes two vectors and returns a vector of corresponding pairs.
zip :: Vector a -> Vector b -> Vector (a, b)
zip :: Vector a -> Vector b -> Vector (a, b)
zip = (a -> b -> (a, b)) -> Vector a -> Vector b -> Vector (a, b)
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith (,)
{-# INLINE zip #-}

-- | /O(n)/. A generalized 'zip' zipping with a function.
zipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith a -> b -> c
f Vector a
v1 Vector b
v2
    | Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector b -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector b
v2 = ([a], Vector c) -> Vector c
forall a b. (a, b) -> b
snd (([a], Vector c) -> Vector c) -> ([a], Vector c) -> Vector c
forall a b. (a -> b) -> a -> b
$ ([a] -> b -> ([a], c)) -> [a] -> Vector b -> ([a], Vector c)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL [a] -> b -> ([a], c)
f' (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) Vector b
v2
    | Bool
otherwise = (b -> a -> c) -> Vector b -> Vector a -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith ((a -> b -> c) -> b -> a -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> c
f) Vector b
v2 Vector a
v1
  where
    f' :: [a] -> b -> ([a], c)
f' [] b
_ = String -> ([a], c)
forall a. HasCallStack => String -> a
error String
"unreachable"
    f' (a
x : [a]
xs) b
y = ([a]
xs, a -> b -> c
f a
x b
y)

-- | /O(n)/. Takes three vectors and returns a vector of corresponding triples.
zip3 :: Vector a -> Vector b -> Vector c -> Vector (a, b, c)
zip3 :: Vector a -> Vector b -> Vector c -> Vector (a, b, c)
zip3 = (a -> b -> c -> (a, b, c))
-> Vector a -> Vector b -> Vector c -> Vector (a, b, c)
forall a b c d.
(a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
zipWith3 (,,)
{-# INLINE zip3 #-}

-- | /O(n)/. A generalized 'zip3' zipping with a function.
zipWith3 :: (a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
zipWith3 :: (a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
zipWith3 a -> b -> c -> d
f Vector a
v1 Vector b
v2 Vector c
v3 = ((c -> d) -> c -> d) -> Vector (c -> d) -> Vector c -> Vector d
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith (c -> d) -> c -> d
forall a b. (a -> b) -> a -> b
($) ((a -> b -> c -> d) -> Vector a -> Vector b -> Vector (c -> d)
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith a -> b -> c -> d
f Vector a
v1 Vector b
v2) Vector c
v3

-- | /O(n)/. Transforms a vector of pairs into a vector of first components and a vector of second components.
unzip :: Vector (a, b) -> (Vector a, Vector b)
unzip :: Vector (a, b) -> (Vector a, Vector b)
unzip Vector (a, b)
v = (((a, b) -> a) -> Vector (a, b) -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
map (a, b) -> a
forall a b. (a, b) -> a
fst Vector (a, b)
v, ((a, b) -> b) -> Vector (a, b) -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
map (a, b) -> b
forall a b. (a, b) -> b
snd Vector (a, b)
v)

-- | /O(n)/. Takes a vector of triples and returns three vectors, analogous to 'unzip'.
unzip3 :: Vector (a, b, c) -> (Vector a, Vector b, Vector c)
unzip3 :: Vector (a, b, c) -> (Vector a, Vector b, Vector c)
unzip3 Vector (a, b, c)
v = (((a, b, c) -> a) -> Vector (a, b, c) -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
map (a, b, c) -> a
forall a b c. (a, b, c) -> a
fst3 Vector (a, b, c)
v, ((a, b, c) -> b) -> Vector (a, b, c) -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
map (a, b, c) -> b
forall a b c. (a, b, c) -> b
snd3 Vector (a, b, c)
v, ((a, b, c) -> c) -> Vector (a, b, c) -> Vector c
forall a b. (a -> b) -> Vector a -> Vector b
map (a, b, c) -> c
forall a b c. (a, b, c) -> c
trd3 Vector (a, b, c)
v)
  where
    fst3 :: (a, b, c) -> a
fst3 (a
x, b
_, c
_) = a
x
    snd3 :: (a, b, c) -> b
snd3 (a
_, b
y, c
_) = b
y
    trd3 :: (a, b, c) -> c
trd3 (a
_, b
_, c
z) = c
z

-- | /O(n)/. Create a list of index-value pairs from the vector.
toIndexedList :: Vector a -> [(Int, a)]
toIndexedList :: Vector a -> [(Int, a)]
toIndexedList = (Int -> a -> [(Int, a)] -> [(Int, a)])
-> [(Int, a)] -> Vector a -> [(Int, a)]
forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
foldrWithIndex (((Int, a) -> [(Int, a)] -> [(Int, a)])
-> Int -> a -> [(Int, a)] -> [(Int, a)]
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (:)) []