{-# LANGUAGE CPP #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleContexts #-}
#ifdef TRUSTWORTHY
{-# LANGUAGE Trustworthy #-}
#endif

#include "lens-common.h"

-------------------------------------------------------------------------------
-- |
-- Module      :  Data.Vector.Generic.Lens
-- Copyright   :  (C) 2012-2016 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  non-portable
--
-- This module provides lenses and traversals for working with generic
-- vectors.
-------------------------------------------------------------------------------
module Data.Vector.Generic.Lens
  ( toVectorOf
  -- * Isomorphisms
  , forced
  , vector
  , asStream
  , asStreamR
  , cloned
  , converted
  -- * Lenses
  , sliced
  -- * Traversal of individual indices
  , ordinals
  , vectorIx
  , vectorTraverse
  ) where

import Prelude ()

import Control.Lens.Type
import Control.Lens.Lens
import Control.Lens.Getter
import Control.Lens.Fold
import Control.Lens.Iso
import Control.Lens.Indexed
import Control.Lens.Setter
import Control.Lens.Traversal
import Control.Lens.Internal.List (ordinalNub)
import Control.Lens.Internal.Prelude
import qualified Data.Vector.Generic as V
import Data.Vector.Generic (Vector)
import Data.Vector.Generic.New (New)

#if MIN_VERSION_vector(0,11,0)
import Data.Vector.Fusion.Bundle (Bundle)
#else
import Data.Vector.Fusion.Stream (Stream)
#endif

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

-- | @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 v a -> f (v a)
f v a
v = 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) f (v a) -> (v a -> v a) -> f (v a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \ 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)
{-# INLINE sliced #-}

-- | Similar to 'toListOf', but returning a 'Vector'.
--
-- >>> (toVectorOf both (8,15) :: Vector.Vector Int) == Vector.fromList [8,15]
-- True
toVectorOf :: Vector v a => Getting (Endo [a]) s a -> s -> v a
toVectorOf :: Getting (Endo [a]) s a -> s -> v a
toVectorOf Getting (Endo [a]) s a
l s
s = [a] -> v a
forall (v :: * -> *) a. Vector v a => [a] -> v 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 #-}

-- | 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] ^. from 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 #-}

#if MIN_VERSION_vector(0,11,0)
-- | 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)
#else
-- | Convert a 'Vector' to a finite 'Stream' (or back.)
asStream :: (Vector v a, Vector v b) => Iso (v a) (v b) (Stream a) (Stream b)
#endif
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 #-}

#if MIN_VERSION_vector(0,11,0)
-- | 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)
#else
-- | Convert a 'Vector' to a finite 'Stream' from right to left (or
-- back.)
asStreamR :: (Vector v a, Vector v b) => Iso (v a) (v b) (Stream a) (Stream b)
#endif
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 => Iso' (v a) (v a)
forced :: Iso' (v a) (v a)
forced = (v a -> v a) -> Iso' (v a) (v a)
forall a. (a -> a) -> Iso' a a
involuted v a -> v a
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 :: Vector v a => [Int] -> IndexedTraversal' Int (v a) a
ordinals :: [Int] -> IndexedTraversal' Int (v a) a
ordinals [Int]
is p a (f a)
f 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
fmap (v a
v v a -> [(Int, a)] -> v a
forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
V.//) (f [(Int, a)] -> f (v a)) -> f [(Int, a)] -> f (v 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 (v a
v v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v 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 (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 a -> f a
f v a
v
  | 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 = a -> f a
f (v a
v v a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
V.! Int
i) f a -> (a -> v a) -> f (v a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \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)]
  | Bool
otherwise                = 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 :: (V.Vector v a, V.Vector w b) => IndexedTraversal Int (v a) (w b) a b
vectorTraverse :: IndexedTraversal Int (v a) (w b) a b
vectorTraverse p a (f b)
f v a
v = Int -> [b] -> w b
forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
V.fromListN (v a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v) ([b] -> w b) -> f [b] -> f (w b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> p a (f b) -> [a] -> f [b]
forall (f :: * -> *) a b.
Traversable f =>
IndexedTraversal Int (f a) (f b) a b
traversed p a (f b)
f (v a -> [a]
forall (v :: * -> *) a. Vector v a => v a -> [a]
V.toList v a
v)
{-# INLINE [0] vectorTraverse #-}

{-# RULES
"vectorTraverse -> mapped" vectorTraverse  = sets V.map         :: (V.Vector v a, V.Vector v b) => ASetter (v a) (v b) a b;
"vectorTraverse -> imapped" vectorTraverse = isets V.imap       :: (V.Vector v a, V.Vector v b) => AnIndexedSetter Int (v a) (v b) a b;
"vectorTraverse -> foldr"  vectorTraverse  = foldring V.foldr   :: V.Vector v a => Getting (Endo r) (v a) a;
"vectorTraverse -> ifoldr" vectorTraverse  = ifoldring V.ifoldr :: V.Vector v a => IndexedGetting Int (Endo r) (v a) a;
 #-}

-- | 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