{-# 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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
Index (Vector a)
i Bool -> Bool -> Bool
&& Int
Index (Vector a)
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
v = (a -> f a
IxValue (Vector a) -> f (IxValue (Vector a))
f (a -> f a) -> a -> f a
forall a b. (a -> b) -> a -> b
$! Vector a
v Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.! Int
Index (Vector a)
i) f a -> (a -> Vector a) -> f (Vector a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> Vector a
v Vector a -> [(Int, a)] -> Vector a
forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Int
Index (Vector a)
i, a
a)]
    | Bool
otherwise                     = Vector a -> f (Vector a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector a
v
  {-# INLINE ix #-}

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

instance Each (Vector a) (Vector b) a b where
  each :: (a -> f b) -> Vector a -> f (Vector b)
each = (a -> f b) -> Vector a -> f (Vector b)
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' :: p (Unwrapped (Vector a)) (f (Unwrapped (Vector a)))
-> p (Vector a) (f (Vector a))
_Wrapped' = (Vector a -> [a])
-> ([a] -> Vector a) -> Iso (Vector a) (Vector a) [a] [a]
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Vector a -> [a]
forall a. Vector a -> [a]
V.toList [a] -> Vector a
forall a. [a] -> Vector a
V.fromList
  {-# INLINE _Wrapped' #-}

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

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

instance Reversing (Vector a) where
  reversing :: Vector a -> Vector a
reversing = Vector a -> Vector a
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 :: 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 (Int -> Int -> Vector a -> Vector a
forall a. Int -> Int -> Vector a -> Vector a
V.slice Int
i Int
n Vector a
v) f (Vector a) -> (Vector a -> Vector a) -> f (Vector a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \ Vector a
v0 -> Vector a
v Vector a -> [(Int, a)] -> Vector a
forall a. Vector a -> [(Int, a)] -> Vector a
V.// [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
i..Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] (Vector a -> [a]
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 :: Getting (Endo [a]) s a -> s -> Vector a
toVectorOf Getting (Endo [a]) s a
l s
s = [a] -> Vector a
forall a. [a] -> Vector a
V.fromList (Getting (Endo [a]) s a -> s -> [a]
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 :: p (Vector a) (f (Vector b)) -> p [a] (f [b])
vector = ([a] -> Vector a)
-> (Vector b -> [b]) -> Iso [a] [b] (Vector a) (Vector b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso [a] -> Vector a
forall a. [a] -> Vector a
V.fromList Vector b -> [b]
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 :: p (Vector a) (f (Vector b)) -> p (Vector a) (f (Vector b))
forced = (Vector a -> Vector a)
-> (Vector b -> Vector b)
-> Iso (Vector a) (Vector b) (Vector a) (Vector b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Vector a -> Vector a
forall a. Vector a -> Vector a
V.force Vector b -> Vector b
forall a. Vector a -> Vector a
V.force
{-# INLINE forced #-}

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