{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UndecidableInstances  #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.Strict.Vector.Lens
  ( toVectorOf
  , vector
  , forced
  , sliced
  , ordinals
  ) where

import           Control.Lens
import           Control.Lens.Internal.List (ordinalNub)
import           Data.Vector.Generic.Lens (vectorTraverse)
import           Data.Monoid (Endo)

import qualified Data.Strict.Vector as V

import           Data.Strict.Vector (Vector)


#if !MIN_VERSION_lens(5,0,0)
instance FunctorWithIndex Int Vector where
  imap = V.imap
  {-# INLINE imap #-}

instance FoldableWithIndex Int Vector where
  ifoldr = V.ifoldr
  {-# INLINE ifoldr #-}
  ifoldl = V.ifoldl . flip
  {-# INLINE ifoldl #-}
  ifoldr' = V.ifoldr'
  {-# INLINE ifoldr' #-}
  ifoldl' = V.ifoldl' . flip
  {-# INLINE ifoldl' #-}

instance TraversableWithIndex Int Vector where
  itraverse f v =
    let !n = V.length v in V.fromListN n <$> itraverse f (V.toList v)
  {-# INLINE itraverse #-}
#endif

type instance Index (Vector a) = Int
type instance IxValue (Vector a) = a
instance Ixed (Vector a) where
  -- This is slightly different from lens' definition to make our ixTest work.
  -- Unlike Sequence, since the element is stored inside a primitive array the
  -- only way to get it is via a primop, so we have to force-apply ($!) the
  -- extraction. Pattern matching on a case-expr (e.g. the result of V.!?) is
  -- ineffective because that still has to call the primop, so we would still
  -- have a thunk after pattern matching.
  ix :: Index (Vector a) -> Traversal' (Vector a) (IxValue (Vector a))
ix Index (Vector a)
i IxValue (Vector a) -> f (IxValue (Vector a))
f Vector a
v
    | Int
0 forall a. Ord a => a -> a -> Bool
<= Index (Vector a)
i Bool -> Bool -> Bool
&& Index (Vector a)
i forall a. Ord a => a -> a -> Bool
< forall a. Vector a -> Int
V.length Vector a
v = (IxValue (Vector a) -> f (IxValue (Vector a))
f forall a b. (a -> b) -> a -> b
$! Vector a
v forall a. Vector a -> Int -> a
V.! Index (Vector a)
i) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> Vector a
v forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Index (Vector a)
i, a
a)]
    | Bool
otherwise                     = forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector a
v
  {-# INLINE ix #-}

instance AsEmpty (Vector a) where
  _Empty :: Prism' (Vector a) ()
_Empty = forall a. a -> (a -> Bool) -> Prism' a ()
nearly forall a. Vector a
V.empty forall a. Vector a -> Bool
V.null
  {-# INLINE _Empty #-}

instance Each (Vector a) (Vector b) a b where
  each :: Traversal (Vector a) (Vector b) a b
each = forall (v :: * -> *) a (w :: * -> *) b.
(Vector v a, Vector w b) =>
IndexedTraversal Int (v a) (w b) a b
vectorTraverse
  {-# INLINE each #-}

instance (t ~ Vector a') => Rewrapped (Vector a) t
instance Wrapped (Vector a) where
  type Unwrapped (Vector a) = [a]
  _Wrapped' :: Iso' (Vector a) (Unwrapped (Vector a))
_Wrapped' = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall a. Vector a -> [a]
V.toList forall a. [a] -> Vector a
V.fromList
  {-# INLINE _Wrapped' #-}

instance Cons (Vector a) (Vector b) a b where
  _Cons :: Prism (Vector a) (Vector b) (a, Vector a) (b, Vector b)
_Cons = forall b t s a. (b -> t) -> (s -> Either t a) -> Prism s t a b
prism (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. a -> Vector a -> Vector a
V.cons) forall a b. (a -> b) -> a -> b
$ \Vector a
v ->
    if forall a. Vector a -> Bool
V.null Vector a
v
    then forall a b. a -> Either a b
Left forall a. Vector a
V.empty
    else forall a b. b -> Either a b
Right (forall a. Vector a -> a
V.unsafeHead Vector a
v, forall a. Vector a -> Vector a
V.unsafeTail Vector a
v)
  {-# INLINE _Cons #-}

instance Snoc (Vector a) (Vector b) a b where
  _Snoc :: Prism (Vector a) (Vector b) (Vector a, a) (Vector b, b)
_Snoc = forall b t s a. (b -> t) -> (s -> Either t a) -> Prism s t a b
prism (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. Vector a -> a -> Vector a
V.snoc) forall a b. (a -> b) -> a -> b
$ \Vector a
v -> if forall a. Vector a -> Bool
V.null Vector a
v
    then forall a b. a -> Either a b
Left forall a. Vector a
V.empty
    else forall a b. b -> Either a b
Right (forall a. Vector a -> Vector a
V.unsafeInit Vector a
v, forall a. Vector a -> a
V.unsafeLast Vector a
v)
  {-# INLINE _Snoc #-}

instance Reversing (Vector a) where
  reversing :: Vector a -> Vector a
reversing = forall a. Vector a -> Vector a
V.reverse

-- | Analogous to 'Data.Vector.Lens.sliced'.
sliced :: Int -- ^ @i@ starting index
       -> Int -- ^ @n@ length
       -> Lens' (Vector a) (Vector a)
sliced :: forall a. Int -> Int -> Lens' (Vector a) (Vector a)
sliced Int
i Int
n Vector a -> f (Vector a)
f Vector a
v = Vector a -> f (Vector a)
f (forall a. Int -> Int -> Vector a -> Vector a
V.slice Int
i Int
n Vector a
v) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \ Vector a
v0 -> Vector a
v forall a. Vector a -> [(Int, a)] -> Vector a
V.// forall a b. [a] -> [b] -> [(a, b)]
zip [Int
i..Int
iforall a. Num a => a -> a -> a
+Int
nforall a. Num a => a -> a -> a
-Int
1] (forall a. Vector a -> [a]
V.toList Vector a
v0)
{-# INLINE sliced #-}

-- | Analogous to 'Data.Vector.Lens.toVectorOf'.
toVectorOf :: Getting (Endo [a]) s a -> s -> Vector a
toVectorOf :: forall a s. Getting (Endo [a]) s a -> s -> Vector a
toVectorOf Getting (Endo [a]) s a
l s
s = forall a. [a] -> Vector a
V.fromList (forall a s. Getting (Endo [a]) s a -> s -> [a]
toListOf Getting (Endo [a]) s a
l s
s)
{-# INLINE toVectorOf #-}

-- | Analogous to 'Data.Vector.Lens.vector'.
vector :: Iso [a] [b] (Vector a) (Vector b)
vector :: forall a b. Iso [a] [b] (Vector a) (Vector b)
vector = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall a. [a] -> Vector a
V.fromList forall a. Vector a -> [a]
V.toList
{-# INLINE vector #-}

-- | Analogous to 'Data.Vector.Lens.forced'.
forced :: Iso (Vector a) (Vector b) (Vector a) (Vector b)
forced :: forall a b. Iso (Vector a) (Vector b) (Vector a) (Vector b)
forced = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall a. Vector a -> Vector a
V.force forall a. Vector a -> Vector a
V.force
{-# INLINE forced #-}

-- | Analogous to 'Data.Vector.Lens.ordinals'.
ordinals :: [Int] -> IndexedTraversal' Int (Vector a) a
ordinals :: forall a. [Int] -> IndexedTraversal' Int (Vector a) a
ordinals [Int]
is p a (f a)
f Vector a
v = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Vector a
v forall a. Vector a -> [(Int, a)] -> Vector a
V.//) forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Int
i -> (,) Int
i forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p a (f a)
f Int
i (Vector a
v forall a. Vector a -> Int -> a
V.! Int
i)) forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
ordinalNub (forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v) [Int]
is
{-# INLINE ordinals #-}