{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.Vec.Pull.Optics (
    -- * Indexing
    ix,
    _Cons,
    _head,
    _tail,
    -- * Conversions
    _Vec,
    ) where

import Data.Fin    (Fin (..))
import Data.Nat    (Nat (..))
import Optics.Core ((<&>))

import qualified Data.Type.Nat as N
import qualified Optics.Core   as L

import Data.Vec.Pull

-- $setup
-- >>> :set -XScopedTypeVariables
-- >>> import Data.Fin (Fin (..))
-- >>> import Data.Vec.Pull
-- >>> import Optics.Core (over, view, set, (%))
-- >>> import qualified Data.Vec.Lazy as L
-- >>> import qualified Data.Vec.Lazy.Optics as L

type LensLikeVL f s t a b = (a -> f b) -> s -> f t
type LensLikeVL' f s a = LensLikeVL f s s a a

-------------------------------------------------------------------------------
-- Indexing
-------------------------------------------------------------------------------

-- | Index lens.
--
-- 'a' ::: 'x' ::: 'c' ::: VNil
-- >>> view (L._Pull % ix (FS FZ)) ('a' L.::: 'b' L.::: 'c' L.::: L.VNil)
-- 'b'
--
-- >>> set (L._Pull % ix (FS FZ)) 'x' ('a' L.::: 'b' L.::: 'c' L.::: L.VNil)
-- 'a' ::: 'x' ::: 'c' ::: VNil
--
ix :: Fin n -> L.Lens' (Vec n a) a
ix :: forall (n :: Nat) a. Fin n -> Lens' (Vec n a) a
ix Fin n
i = LensVL (Vec n a) (Vec n a) a a -> Lens (Vec n a) (Vec n a) a a
forall s t a b. LensVL s t a b -> Lens s t a b
L.lensVL (Fin n -> LensLikeVL' f (Vec n a) a
forall (f :: * -> *) (n :: Nat) a.
Functor f =>
Fin n -> LensLikeVL' f (Vec n a) a
ixVL Fin n
i)
{-# INLINE ix #-}

ixVL :: Functor f => Fin n -> LensLikeVL' f (Vec n a) a
ixVL :: forall (f :: * -> *) (n :: Nat) a.
Functor f =>
Fin n -> LensLikeVL' f (Vec n a) a
ixVL Fin n
i a -> f a
f (Vec Fin n -> a
v) = a -> f a
f (Fin n -> a
v Fin n
i) f a -> (a -> Vec n a) -> f (Vec n a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> (Fin n -> a) -> Vec n a
forall (n :: Nat) a. (Fin n -> a) -> Vec n a
Vec ((Fin n -> a) -> Vec n a) -> (Fin n -> a) -> Vec n a
forall a b. (a -> b) -> a -> b
$ \Fin n
j ->
    if Fin n
i Fin n -> Fin n -> Bool
forall a. Eq a => a -> a -> Bool
== Fin n
j
    then a
a
    else Fin n -> a
v Fin n
j
{-# INLINE ixVL #-}

-- | Match on non-empty 'Vec'.
--
-- /Note:/ @lens@ 'L._Cons' is a 'L.Prism'.
-- In fact, @'Vec' n a@ cannot have an instance of 'L.Cons' as types don't match.
--
_Cons :: L.Iso (Vec ('S n) a) (Vec ('S n) b) (a, Vec n a) (b, Vec n b)
_Cons :: forall (n :: Nat) a b.
Iso (Vec ('S n) a) (Vec ('S n) b) (a, Vec n a) (b, Vec n b)
_Cons = (Vec ('S n) a -> (a, Vec n a))
-> ((b, Vec n b) -> Vec ('S n) b)
-> Iso (Vec ('S n) a) (Vec ('S n) b) (a, Vec n a) (b, Vec n b)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
L.iso (\(Vec Fin ('S n) -> a
v) -> (Fin ('S n) -> a
v Fin ('S n)
forall (n1 :: Nat). Fin ('S n1)
FZ, (Fin n -> a) -> Vec n a
forall (n :: Nat) a. (Fin n -> a) -> Vec n a
Vec (Fin ('S n) -> a
v (Fin ('S n) -> a) -> (Fin n -> Fin ('S n)) -> Fin n -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fin n -> Fin ('S n)
forall (n1 :: Nat). Fin n1 -> Fin ('S n1)
FS))) (\(b
x, Vec n b
xs) -> b -> Vec n b -> Vec ('S n) b
forall a (n :: Nat). a -> Vec n a -> Vec ('S n) a
cons b
x Vec n b
xs)

-- | Head lens. /Note:/ @lens@ 'L._head' is a 'L.Traversal''.
--
-- >>> view (L._Pull % _head) ('a' L.::: 'b' L.::: 'c' L.::: L.VNil)
-- 'a'
--
-- >>> set (L._Pull % _head) 'x' ('a' L.::: 'b' L.::: 'c' L.::: L.VNil)
-- 'x' ::: 'b' ::: 'c' ::: VNil
--
_head :: L.Lens' (Vec ('S n) a) a
_head :: forall (n :: Nat) a. Lens' (Vec ('S n) a) a
_head = LensVL (Vec ('S n) a) (Vec ('S n) a) a a
-> Lens (Vec ('S n) a) (Vec ('S n) a) a a
forall s t a b. LensVL s t a b -> Lens s t a b
L.lensVL LensLikeVL' f (Vec ('S n) a) a
LensVL (Vec ('S n) a) (Vec ('S n) a) a a
forall (f :: * -> *) (n :: Nat) a.
Functor f =>
LensLikeVL' f (Vec ('S n) a) a
headVL
{-# INLINE _head #-}

headVL :: Functor f => LensLikeVL' f (Vec ('S n) a) a
headVL :: forall (f :: * -> *) (n :: Nat) a.
Functor f =>
LensLikeVL' f (Vec ('S n) a) a
headVL a -> f a
f (Vec Fin ('S n) -> a
v) = a -> f a
f (Fin ('S n) -> a
v Fin ('S n)
forall (n1 :: Nat). Fin ('S n1)
FZ) f a -> (a -> Vec ('S n) a) -> f (Vec ('S n) a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> (Fin ('S n) -> a) -> Vec ('S n) a
forall (n :: Nat) a. (Fin n -> a) -> Vec n a
Vec ((Fin ('S n) -> a) -> Vec ('S n) a)
-> (Fin ('S n) -> a) -> Vec ('S n) a
forall a b. (a -> b) -> a -> b
$ \Fin ('S n)
j -> case Fin ('S n)
j of
    Fin ('S n)
FZ -> a
a
    Fin ('S n)
_   -> Fin ('S n) -> a
v Fin ('S n)
j
{-# INLINE headVL #-}

-- | Tail lens.
_tail :: L.Lens' (Vec ('S n) a) (Vec n a)
_tail :: forall (n :: Nat) a. Lens' (Vec ('S n) a) (Vec n a)
_tail = LensVL (Vec ('S n) a) (Vec ('S n) a) (Vec n a) (Vec n a)
-> Lens (Vec ('S n) a) (Vec ('S n) a) (Vec n a) (Vec n a)
forall s t a b. LensVL s t a b -> Lens s t a b
L.lensVL LensLikeVL' f (Vec ('S n) a) (Vec n a)
LensVL (Vec ('S n) a) (Vec ('S n) a) (Vec n a) (Vec n a)
forall (f :: * -> *) (n :: Nat) a.
Functor f =>
LensLikeVL' f (Vec ('S n) a) (Vec n a)
tailVL
{-# INLINE _tail #-}

tailVL :: Functor f => LensLikeVL' f (Vec ('S n) a) (Vec n a)
tailVL :: forall (f :: * -> *) (n :: Nat) a.
Functor f =>
LensLikeVL' f (Vec ('S n) a) (Vec n a)
tailVL Vec n a -> f (Vec n a)
f (Vec Fin ('S n) -> a
v) = Vec n a -> f (Vec n a)
f ((Fin n -> a) -> Vec n a
forall (n :: Nat) a. (Fin n -> a) -> Vec n a
Vec (Fin ('S n) -> a
v (Fin ('S n) -> a) -> (Fin n -> Fin ('S n)) -> Fin n -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fin n -> Fin ('S n)
forall (n1 :: Nat). Fin n1 -> Fin ('S n1)
FS)) f (Vec n a) -> (Vec n a -> Vec ('S n) a) -> f (Vec ('S n) a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Vec n a
xs -> a -> Vec n a -> Vec ('S n) a
forall a (n :: Nat). a -> Vec n a -> Vec ('S n) a
cons (Fin ('S n) -> a
v Fin ('S n)
forall (n1 :: Nat). Fin ('S n1)
FZ) Vec n a
xs
{-# INLINE tailVL #-}

-------------------------------------------------------------------------------
-- Conversions
-------------------------------------------------------------------------------

-- | Prism from list.
_Vec :: N.SNatI n => L.Prism' [a] (Vec n a)
_Vec :: forall (n :: Nat) a. SNatI n => Prism' [a] (Vec n a)
_Vec = (Vec n a -> [a])
-> ([a] -> Maybe (Vec n a)) -> Prism [a] [a] (Vec n a) (Vec n a)
forall b s a. (b -> s) -> (s -> Maybe a) -> Prism s s a b
L.prism' Vec n a -> [a]
forall (n :: Nat) a. SNatI n => Vec n a -> [a]
toList [a] -> Maybe (Vec n a)
forall (n :: Nat) a. SNatI n => [a] -> Maybe (Vec n a)
fromList