{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeFamilies          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.RAVec.Tree.DF.Lens (
    -- * Indexing
    ix,
    ) where

import Control.Applicative ((<$>))
import Data.Wrd            (Wrd (..))
import Prelude ()

import qualified Control.Lens  as L
import qualified Data.Type.Nat as N

import Data.RAVec.Tree.DF

-- $setup
-- >>> import Control.Lens ((^.), (&), (.~), (^?), (#))
-- >>> import Prelude
-- >>> import Data.RAVec.Tree.DF
-- >>> import Data.Wrd (Wrd (..))

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

-- | Index lens.
--
-- >>> let tree = Node (Node (Leaf 'a') (Leaf 'b')) (Node (Leaf 'c') (Leaf 'd'))
-- >>> tree & ix (W1 $ W0 WE) .~ 'z'
-- Node (Node (Leaf 'a') (Leaf 'b')) (Node (Leaf 'z') (Leaf 'd'))
--
ix :: Wrd n -> L.Lens' (Tree n a) a
ix :: forall (n :: Nat) a. Wrd n -> Lens' (Tree n a) a
ix Wrd n
WE      a -> f a
f (Leaf a
x)   = a -> Tree n a
a -> Tree 'Z a
forall a. a -> Tree 'Z a
Leaf (a -> Tree n a) -> f a -> f (Tree n a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
x
ix (W0 Wrd n1
is) a -> f a
f (Node Tree n1 a
x Tree n1 a
y) = (Tree n1 a -> Tree n1 a -> Tree ('S n1) a
forall (n :: Nat) a. Tree n a -> Tree n a -> Tree ('S n) a
`Node` Tree n1 a
y) (Tree n1 a -> Tree n a) -> f (Tree n1 a) -> f (Tree n a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Wrd n1 -> Lens' (Tree n1 a) a
forall (n :: Nat) a. Wrd n -> Lens' (Tree n a) a
ix Wrd n1
is a -> f a
f Tree n1 a
x
ix (W1 Wrd n1
is) a -> f a
f (Node Tree n1 a
x Tree n1 a
y) = (Tree n1 a
x Tree n1 a -> Tree n1 a -> Tree ('S n1) a
forall (n :: Nat) a. Tree n a -> Tree n a -> Tree ('S n) a
`Node`) (Tree n1 a -> Tree n a) -> f (Tree n1 a) -> f (Tree n a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Wrd n1 -> Lens' (Tree n1 a) a
forall (n :: Nat) a. Wrd n -> Lens' (Tree n a) a
ix Wrd n1
is a -> f a
f Tree n1 a
y

-------------------------------------------------------------------------------
-- Instances
-------------------------------------------------------------------------------

instance N.SNatI n => L.Each (Tree n a) (Tree n b) a b

type instance L.Index (Tree n a)   = Wrd n
type instance L.IxValue (Tree n a) = a

instance L.Ixed (Tree n a) where
    ix :: Index (Tree n a) -> Traversal' (Tree n a) (IxValue (Tree n a))
ix Index (Tree n a)
i = Wrd n -> Lens' (Tree n a) a
forall (n :: Nat) a. Wrd n -> Lens' (Tree n a) a
ix Wrd n
Index (Tree n a)
i