{-# 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 :: forall (v :: * -> *) a.
Vector v a =>
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 (forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
Int -> Int -> v a -> v a
V.slice Int
i Int
n v a
v) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \ v a
v0 -> v a
v forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v 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 (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 :: forall (v :: * -> *) a s.
Vector v a =>
Getting (Endo [a]) s a -> s -> v a
toVectorOf Getting (Endo [a]) s a
l s
s = forall (v :: * -> *) a. Vector v a => [a] -> v a
V.fromList (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 :: forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
Iso [a] [b] (v a) (v b)
vector = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall (v :: * -> *) a. Vector v a => [a] -> v a
V.fromList 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 :: forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
Iso (v a) (v b) (Bundle v a) (Bundle v b)
asStream = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall (v :: * -> *) a. Vector v a => v a -> Bundle v a
V.stream 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 :: forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
Iso (v a) (v b) (Bundle v a) (Bundle v b)
asStreamR = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall (v :: * -> *) a (u :: * -> *).
Vector v a =>
v a -> Bundle u a
V.streamR 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 :: forall (v :: * -> *) a. Vector v a => Iso' (v a) (New v a)
cloned = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall (v :: * -> *) a. Vector v a => v a -> New v a
V.clone 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 :: forall (v :: * -> *) a. Vector v a => Iso' (v a) (v a)
forced = forall a. (a -> a) -> Iso' a a
involuted 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 :: forall (v :: * -> *) a.
Vector v a =>
[Int] -> IndexedTraversal' Int (v a) a
ordinals [Int]
is p a (f a)
f v a
v = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (v a
v forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v 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 (v a
v forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
i)) forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
ordinalNub (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 :: forall (v :: * -> *) a. Vector v a => Int -> Traversal' (v a) a
vectorIx Int
i a -> f a
f v a
v
  | Int
0 forall a. Ord a => a -> a -> Bool
<= Int
i Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
< forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v = a -> f a
f (v a
v forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
i) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> v a
v forall (v :: * -> *) a. Vector v a => v a -> [(Int, a)] -> v a
V.// [(Int
i, a
a)]
  | Bool
otherwise                = 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 :: forall (v :: * -> *) a (w :: * -> *) b.
(Vector v a, Vector w b) =>
IndexedTraversal Int (v a) (w b) a b
vectorTraverse p a (f b)
f v a
v = forall (v :: * -> *) a. Vector v a => Int -> [a] -> v a
V.fromListN (forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
v) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a b.
Traversable f =>
IndexedTraversal Int (f a) (f b) a b
traversed p a (f b)
f (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 :: forall (v :: * -> *) a (w :: * -> *) b.
(Vector v a, Vector w a, Vector v b, Vector w b) =>
Iso (v a) (v b) (w a) (w b)
converted = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert