{-# 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 Data.Vector.Fusion.Bundle (Bundle)
import qualified Data.Vector.Generic as V
import Data.Vector.Generic (Vector)
import Data.Vector.Generic.New (New)

-- $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.
(HasCallStack, 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 #-}

-- | 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 => 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.
(HasCallStack, 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.
(HasCallStack, 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