-- | This module provides lenses and traversals for working with generic
-- vectors.
module Data.Vector.Generic.Optics
  ( toVectorOf
  -- * Isomorphisms
  , forced
  , vector
  , asStream
  , asStreamR
  , cloned
  , converted
  -- * Lenses
  , sliced
  -- * Traversal of individual indices
  , ordinals
  , vectorIx
  , vectorTraverse
  ) where

import Data.Vector.Fusion.Bundle (Bundle)
import Data.Vector.Generic (Vector)
import qualified Data.Vector.Generic as V
import Data.Vector.Generic.New (New)

import Data.Profunctor.Indexed

import Optics.Core
import Optics.Extra.Internal.Vector
import Optics.Internal.Fold
import Optics.Internal.IxFold
import Optics.Internal.Optic

-- $setup
-- >>> import qualified Data.Vector as Vector

-- | @sliced i n@ provides a 'Lens' that edits the @n@ elements starting at
-- index @i@ from a 'Lens'.
--
-- This is only a valid 'Lens' if you do not change the length of the resulting
-- 'Vector'.
--
-- Attempting to return a longer or shorter vector will result in violations of
-- the 'Lens' laws.
--
-- >>> Vector.fromList [1..10] ^. sliced 2 5 == Vector.fromList [3,4,5,6,7]
-- True
--
-- >>> (Vector.fromList [1..10] & sliced 2 5 % mapped .~ 0) == Vector.fromList [1,2,0,0,0,0,0,8,9,10]
-- True
sliced
  :: Vector v a
  => Int -- ^ @i@ starting index
  -> Int -- ^ @n@ length
  -> Lens' (v a) (v a)
sliced :: Int -> Int -> Lens' (v a) (v a)
sliced Int
i Int
n = LensVL (v a) (v a) (v a) (v a) -> Lens' (v a) (v a)
forall s t a b. LensVL s t a b -> Lens s t a b
lensVL (LensVL (v a) (v a) (v a) (v a) -> Lens' (v a) (v a))
-> LensVL (v a) (v a) (v a) (v a) -> Lens' (v a) (v a)
forall a b. (a -> b) -> a -> b
$ \v a -> f (v a)
f v a
v ->
  (\v a
v0 -> v a
v v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v 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] (v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
V.toList v a
v0)) (v a -> v a) -> f (v a) -> f (v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> v a -> f (v a)
f (Int -> Int -> v a -> v a
forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
V.slice Int
i Int
n v a
v)
{-# INLINE sliced #-}

-- | Similar to 'toListOf', but returning a 'Vector'.
--
-- >>> (toVectorOf each (8,15) :: Vector.Vector Int) == Vector.fromList [8,15]
-- True
toVectorOf
  :: (Is k A_Fold, Vector v a)
  => Optic' k is s a
  -> s
  -> v a
toVectorOf :: Optic' k is s a -> s -> v a
toVectorOf Optic' k is s a
l s
s = [a] -> v a
forall (v :: * -> *) a. Vector v a => [a] -> v a
V.fromList (Optic' k is s a -> s -> [a]
forall k (is :: IxList) s a.
Is k A_Fold =>
Optic' k is s a -> s -> [a]
toListOf Optic' k is s a
l s
s)
{-# INLINE toVectorOf #-}

-- | Convert a list to a 'Vector' (or back.)
--
-- >>> ([1,2,3] ^. vector :: Vector.Vector Int) == Vector.fromList [1,2,3]
-- True
--
-- >>> Vector.fromList [0,8,15] ^. re vector
-- [0,8,15]
vector
  :: (Vector v a, Vector v b)
  => Iso [a] [b] (v a) (v b)
vector :: Iso [a] [b] (v a) (v b)
vector = ([a] -> v a) -> (v b -> [b]) -> Iso [a] [b] (v a) (v b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso [a] -> v a
forall (v :: * -> *) a. Vector v a => [a] -> v a
V.fromList v b -> [b]
forall (v :: * -> *) a. Vector v a => v a -> [a]
V.toList
{-# INLINE vector #-}

-- | Convert a 'Vector' to a finite 'Bundle' (or back.)
asStream
  :: (Vector v a, Vector v b)
  => Iso (v a) (v b) (Bundle v a) (Bundle v b)
asStream :: Iso (v a) (v b) (Bundle v a) (Bundle v b)
asStream = (v a -> Bundle v a)
-> (Bundle v b -> v b) -> Iso (v a) (v b) (Bundle v a) (Bundle v b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso v a -> Bundle v a
forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
V.stream Bundle v b -> v b
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
V.unstream
{-# INLINE asStream #-}

-- | Convert a 'Vector' to a finite 'Bundle' from right to left (or back.)
asStreamR
  :: (Vector v a, Vector v b)
  => Iso (v a) (v b) (Bundle v a) (Bundle v b)
asStreamR :: Iso (v a) (v b) (Bundle v a) (Bundle v b)
asStreamR = (v a -> Bundle v a)
-> (Bundle v b -> v b) -> Iso (v a) (v b) (Bundle v a) (Bundle v b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso v a -> Bundle v a
forall (v :: * -> *) a (u :: * -> *).
Vector v a =>
v a -> Bundle u a
V.streamR Bundle v b -> v b
forall (v :: * -> *) a. Vector v a => Bundle v a -> v a
V.unstreamR
{-# INLINE asStreamR #-}

-- | Convert a 'Vector' back and forth to an initializer that when run produces
-- a copy of the 'Vector'.
cloned :: Vector v a => Iso' (v a) (New v a)
cloned :: Iso' (v a) (New v a)
cloned = (v a -> New v a) -> (New v a -> v a) -> Iso' (v a) (New v a)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso v a -> New v a
forall (v :: * -> *) a. Vector v a => v a -> New v a
V.clone New v a -> v a
forall (v :: * -> *) a. Vector v a => New v a -> v a
V.new
{-# INLINE cloned #-}

-- | Convert a 'Vector' to a version that doesn't retain any extra memory.
forced :: (Vector v a, Vector v b) => Iso (v a) (v b) (v a) (v b)
forced :: Iso (v a) (v b) (v a) (v b)
forced = (v a -> v a) -> (v b -> v b) -> Iso (v a) (v b) (v a) (v b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso v a -> v a
forall (v :: * -> *) a. Vector v a => v a -> v a
V.force v b -> v b
forall (v :: * -> *) a. Vector v a => v a -> v a
V.force
{-# INLINE forced #-}

-- | This 'Traversal' will ignore any duplicates in the supplied list of
-- indices.
--
-- >>> toListOf (ordinals [1,3,2,5,9,10]) $ Vector.fromList [2,4..40]
-- [4,8,6,12,20,22]
ordinals :: forall v a. Vector v a => [Int] -> IxTraversal' Int (v a) a
ordinals :: [Int] -> IxTraversal' Int (v a) a
ordinals [Int]
is = IxTraversalVL Int (v a) (v a) a a -> IxTraversal' Int (v a) a
forall i s t a b. IxTraversalVL i s t a b -> IxTraversal i s t a b
itraversalVL (IxTraversalVL Int (v a) (v a) a a -> IxTraversal' Int (v a) a)
-> IxTraversalVL Int (v a) (v a) a a -> IxTraversal' Int (v a) a
forall a b. (a -> b) -> a -> b
$ \Int -> a -> f a
f v a
v ->
  (v a
v v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
V.//) ([(Int, a)] -> v a) -> f [(Int, a)] -> f (v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f 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
<$> Int -> a -> f a
f Int
i (v a
v v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
V.! Int
i)) (Int -> [Int] -> [Int]
ordinalNub (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v) [Int]
is)
{-# INLINE ordinals #-}

-- | Like 'ix' but polymorphic in the vector type.
vectorIx :: V.Vector v a => Int -> Traversal' (v a) a
vectorIx :: Int -> Traversal' (v a) a
vectorIx Int
i = TraversalVL (v a) (v a) a a -> Traversal' (v a) a
forall s t a b. TraversalVL s t a b -> Traversal s t a b
traversalVL (TraversalVL (v a) (v a) a a -> Traversal' (v a) a)
-> TraversalVL (v a) (v a) a a -> Traversal' (v a) a
forall a b. (a -> b) -> a -> b
$ \a -> f a
f v a
v ->
  if Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v
  then (\a
a -> v a
v v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
V.// [(Int
i, a
a)]) (a -> v a) -> f a -> f (v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f (v a
v v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
V.! Int
i)
  else v a -> f (v a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure v a
v
{-# INLINE vectorIx #-}

-- | Indexed vector traversal for a generic vector.
vectorTraverse
  :: forall v w a b. (V.Vector v a, V.Vector w b)
  => IxTraversal Int (v a) (w b) a b
vectorTraverse :: IxTraversal Int (v a) (w b) a b
vectorTraverse = (forall (p :: * -> * -> * -> *) i.
 Profunctor p =>
 Optic_ A_Traversal p i (Curry (WithIx Int) i) (v a) (w b) a b)
-> IxTraversal Int (v a) (w b) a b
forall k (is :: IxList) s t a b.
(forall (p :: * -> * -> * -> *) i.
 Profunctor p =>
 Optic_ k p i (Curry is i) s t a b)
-> Optic k is s t a b
Optic forall (p :: * -> * -> * -> *) i.
Profunctor p =>
Optic_ A_Traversal p i (Curry (WithIx Int) i) (v a) (w b) a b
forall (p :: * -> * -> * -> *) (v :: * -> *) a (w :: * -> *) b j.
(Traversing p, Vector v a, Vector w b) =>
Optic__ p j (Int -> j) (v a) (w b) a b
vectorTraverse__
{-# INLINE vectorTraverse #-}

-- | Different vector implementations are isomorphic to each other.
converted
  :: (Vector v a, Vector w a, Vector v b, Vector w b)
  => Iso (v a) (v b) (w a) (w b)
converted :: Iso (v a) (v b) (w a) (w b)
converted = (v a -> w a) -> (w b -> v b) -> Iso (v a) (v b) (w a) (w b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso v a -> w a
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert w b -> v b
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert
{-# INLINE converted #-}

----------------------------------------
-- Internal implementations

vectorTraverse__
  :: (Traversing p, V.Vector v a, V.Vector w b)
  => Optic__ p j (Int -> j) (v a) (w b) a b
vectorTraverse__ :: Optic__ p j (Int -> j) (v a) (w b) a b
vectorTraverse__ = (p j a b -> p j (v a) (w b))
-> Optic__ p j (Int -> j) (v a) (w b) a b
-> Optic__ p j (Int -> j) (v a) (w b) a b
forall (p :: * -> * -> * -> *) i a b s t j.
Profunctor p =>
(p i a b -> p i s t) -> (p i a b -> p j s t) -> p i a b -> p j s t
conjoined__ p j a b -> p j (v a) (w b)
forall (p :: * -> * -> * -> *) (v :: * -> *) a (w :: * -> *) b j.
(Traversing p, Vector v a, Vector w b) =>
Optic__ p j j (v a) (w b) a b
vectorTraverseNoIx__ Optic__ p j (Int -> j) (v a) (w b) a b
forall (p :: * -> * -> * -> *) (v :: * -> *) a (w :: * -> *) b j.
(Traversing p, Vector v a, Vector w b) =>
Optic__ p j (Int -> j) (v a) (w b) a b
vectorTraverseIx__
{-# INLINE [0] vectorTraverse__ #-}

vectorTraverseNoIx__
  :: (Traversing p, V.Vector v a, V.Vector w b)
  => Optic__ p j j (v a) (w b) a b
vectorTraverseNoIx__ :: Optic__ p j j (v a) (w b) a b
vectorTraverseNoIx__ = (forall (f :: * -> *).
 Applicative f =>
 (a -> f b) -> v a -> f (w b))
-> Optic__ p j j (v a) (w b) a b
forall (p :: * -> * -> * -> *) a b s t i.
Traversing p =>
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> p i a b -> p i s t
wander ((forall (f :: * -> *).
  Applicative f =>
  (a -> f b) -> v a -> f (w b))
 -> Optic__ p j j (v a) (w b) a b)
-> (forall (f :: * -> *).
    Applicative f =>
    (a -> f b) -> v a -> f (w b))
-> Optic__ p j j (v a) (w b) a b
forall a b. (a -> b) -> a -> b
$ \a -> f b
f v a
v ->
  let !n :: Int
n = v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v in Int -> [b] -> w b
forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
V.fromListN Int
n ([b] -> w b) -> f [b] -> f (w b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> [a] -> f [b]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f (v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
V.toList v a
v)
{-# INLINE vectorTraverseNoIx__ #-}

vectorTraverseIx__
  :: (Traversing p, V.Vector v a, V.Vector w b)
  => Optic__ p j (Int -> j) (v a) (w b) a b
vectorTraverseIx__ :: Optic__ p j (Int -> j) (v a) (w b) a b
vectorTraverseIx__ = (forall (f :: * -> *).
 Applicative f =>
 (Int -> a -> f b) -> v a -> f (w b))
-> Optic__ p j (Int -> j) (v a) (w b) a b
forall (p :: * -> * -> * -> *) i a b s t j.
Traversing p =>
(forall (f :: * -> *).
 Applicative f =>
 (i -> a -> f b) -> s -> f t)
-> p j a b -> p (i -> j) s t
iwander ((forall (f :: * -> *).
  Applicative f =>
  (Int -> a -> f b) -> v a -> f (w b))
 -> Optic__ p j (Int -> j) (v a) (w b) a b)
-> (forall (f :: * -> *).
    Applicative f =>
    (Int -> a -> f b) -> v a -> f (w b))
-> Optic__ p j (Int -> j) (v a) (w b) a b
forall a b. (a -> b) -> a -> b
$ \Int -> a -> f b
f v a
v ->
  let !n :: Int
n = v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v in Int -> [b] -> w b
forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
V.fromListN Int
n ([b] -> w b) -> f [b] -> f (w b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> a -> f b) -> [a] -> f [b]
forall i (t :: * -> *) (f :: * -> *) a b.
(TraversableWithIndex i t, Applicative f) =>
(i -> a -> f b) -> t a -> f (t b)
itraverse Int -> a -> f b
f (v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
V.toList v a
v)
{-# INLINE vectorTraverseIx__ #-}

-- Because vectorTraverse__ inlines late, GHC needs rewrite rules for all cases
-- in order to generate optimal code for each of them. The ones that rewrite
-- traversal into a traversal correspond to an early inline.

{-# RULES

"vectorTraverse__ -> traversed"
  forall (o :: Star g j a b). vectorTraverse__ o = vectorTraverseNoIx__ (reStar o)
    :: (V.Vector v a, V.Vector w b) => Star g (Int -> j) (v a) (w b)

"vectorTraverse__ -> folded"
  forall (o :: Forget r j a b). vectorTraverse__ o = foldring__ V.foldr (reForget o)
    :: (V.Vector v a, V.Vector v b) => Forget r (Int -> j) (v a) (v b)

"vectorTraverse__ -> mapped"
  forall (o :: FunArrow j a b). vectorTraverse__ o = roam V.map (reFunArrow o)
    :: (V.Vector v a, V.Vector v b) => FunArrow (Int -> j) (v a) (v b)

"vectorTraverse__ -> itraversed"
  forall (o :: IxStar g j a b). vectorTraverse__ o = vectorTraverseIx__ o
    :: (V.Vector v a, V.Vector w b) => IxStar g (Int -> j) (v a) (w b)

"vectorTraverse__ -> ifolded"
  forall (o :: IxForget r j a b). vectorTraverse__ o = ifoldring__ V.ifoldr o
    :: (V.Vector v a, V.Vector v b) => IxForget r (Int -> j) (v a) (v b)

"vectorTraverse__ -> imapped"
  forall (o :: IxFunArrow j a b). vectorTraverse__ o = iroam V.imap o
    :: (V.Vector v a, V.Vector v b) => IxFunArrow (Int -> j) (v a) (v b)

#-}