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

import Control.Applicative (Applicative (pure), (<$>))
import Prelude             (Int, Num (..), Ord (..), div, otherwise)

import qualified Control.Lens     as L
import qualified Data.RAList.Tree as Tr

import Data.RAList.NonEmpty

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

ix :: forall f a. Applicative f => Int -> L.LensLike' f (NERAList a) a
ix :: forall (f :: * -> *) a.
Applicative f =>
Int -> LensLike' f (NERAList a) a
ix Int
i0 a -> f a
f (NE NERAList' Leaf a
xs) = NERAList' Leaf a -> NERAList a
forall a. NERAList' Leaf a -> NERAList a
NE (NERAList' Leaf a -> NERAList a)
-> f (NERAList' Leaf a) -> f (NERAList a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> NERAList' Leaf a -> f (NERAList' Leaf a)
forall (t :: * -> *).
TreeIx t =>
Int -> Int -> NERAList' t a -> f (NERAList' t a)
go Int
1 Int
i0 NERAList' Leaf a
xs where
    go :: forall t. TreeIx t => Int -> Int -> NERAList' t a -> f (NERAList' t a)
    go :: forall (t :: * -> *).
TreeIx t =>
Int -> Int -> NERAList' t a -> f (NERAList' t a)
go Int
s Int
i (Last  t a
t)   = t a -> NERAList' t a
forall (f :: * -> *) a. f a -> NERAList' f a
Last (t a -> NERAList' t a) -> f (t a) -> f (NERAList' t a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> (a -> f a) -> t a -> f (t a)
forall (t :: * -> *) (f :: * -> *) a.
(TreeIx t, Applicative f) =>
Int -> Int -> (a -> f a) -> t a -> f (t a)
forall (f :: * -> *) a.
Applicative f =>
Int -> Int -> (a -> f a) -> t a -> f (t a)
treeIx Int
s Int
i a -> f a
f t a
t
    go Int
s Int
i (Cons0   NERAList' (Node t) a
r) = NERAList' (Node t) a -> NERAList' t a
forall (f :: * -> *) a. NERAList' (Node f) a -> NERAList' f a
Cons0 (NERAList' (Node t) a -> NERAList' t a)
-> f (NERAList' (Node t) a) -> f (NERAList' t a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> NERAList' (Node t) a -> f (NERAList' (Node t) a)
forall (t :: * -> *).
TreeIx t =>
Int -> Int -> NERAList' t a -> f (NERAList' t a)
go (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s) Int
i NERAList' (Node t) a
r
    go Int
s Int
i (Cons1 t a
t NERAList' (Node t) a
r)
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
s     = (t a -> NERAList' (Node t) a -> NERAList' t a
forall (f :: * -> *) a.
f a -> NERAList' (Node f) a -> NERAList' f a
`Cons1` NERAList' (Node t) a
r) (t a -> NERAList' t a) -> f (t a) -> f (NERAList' t a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> (a -> f a) -> t a -> f (t a)
forall (t :: * -> *) (f :: * -> *) a.
(TreeIx t, Applicative f) =>
Int -> Int -> (a -> f a) -> t a -> f (t a)
forall (f :: * -> *) a.
Applicative f =>
Int -> Int -> (a -> f a) -> t a -> f (t a)
treeIx Int
s Int
i a -> f a
f t a
t
        | Bool
otherwise = (t a
t t a -> NERAList' (Node t) a -> NERAList' t a
forall (f :: * -> *) a.
f a -> NERAList' (Node f) a -> NERAList' f a
`Cons1`) (NERAList' (Node t) a -> NERAList' t a)
-> f (NERAList' (Node t) a) -> f (NERAList' t a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> NERAList' (Node t) a -> f (NERAList' (Node t) a)
forall (t :: * -> *).
TreeIx t =>
Int -> Int -> NERAList' t a -> f (NERAList' t a)
go (Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
s) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s) NERAList' (Node t) a
r

class TreeIx t where
    treeIx :: Applicative f => Int -> Int -> (a -> f a) -> t a -> f (t a)

instance TreeIx Tr.Leaf where
    treeIx :: forall (f :: * -> *) a.
Applicative f =>
Int -> Int -> (a -> f a) -> Leaf a -> f (Leaf a)
treeIx Int
_ Int
0 a -> f a
f (Tr.Lf a
x) = a -> Leaf a
forall a. a -> Leaf a
Tr.Lf (a -> Leaf a) -> f a -> f (Leaf a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
x
    treeIx Int
_ Int
_ a -> f a
_ Leaf a
leaf   = Leaf a -> f (Leaf a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Leaf a
leaf

instance TreeIx t => TreeIx (Tr.Node t) where
    treeIx :: forall (f :: * -> *) a.
Applicative f =>
Int -> Int -> (a -> f a) -> Node t a -> f (Node t a)
treeIx Int
s Int
i a -> f a
f node :: Node t a
node@(Tr.Nd t a
x t a
y)
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
s2    = (t a -> t a -> Node t a
forall (f :: * -> *) a. f a -> f a -> Node f a
`Tr.Nd` t a
y) (t a -> Node t a) -> f (t a) -> f (Node t a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> (a -> f a) -> t a -> f (t a)
forall (t :: * -> *) (f :: * -> *) a.
(TreeIx t, Applicative f) =>
Int -> Int -> (a -> f a) -> t a -> f (t a)
forall (f :: * -> *) a.
Applicative f =>
Int -> Int -> (a -> f a) -> t a -> f (t a)
treeIx Int
s2 Int
i        a -> f a
f t a
x
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
s     = (t a
x t a -> t a -> Node t a
forall (f :: * -> *) a. f a -> f a -> Node f a
`Tr.Nd`) (t a -> Node t a) -> f (t a) -> f (Node t a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> (a -> f a) -> t a -> f (t a)
forall (t :: * -> *) (f :: * -> *) a.
(TreeIx t, Applicative f) =>
Int -> Int -> (a -> f a) -> t a -> f (t a)
forall (f :: * -> *) a.
Applicative f =>
Int -> Int -> (a -> f a) -> t a -> f (t a)
treeIx Int
s2 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
s2) a -> f a
f t a
x
        | Bool
otherwise = Node t a -> f (Node t a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Node t a
node
      where
        s2 :: Int
s2 = Int
s Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2

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

instance L.Each (NERAList a) (NERAList b) a b

type instance L.Index (NERAList a)   = Int
type instance L.IxValue (NERAList a) = a

instance L.Ixed (NERAList a) where
    ix :: Index (NERAList a)
-> Traversal' (NERAList a) (IxValue (NERAList a))
ix Index (NERAList a)
i = Int -> LensLike' f (NERAList a) a
forall (f :: * -> *) a.
Applicative f =>
Int -> LensLike' f (NERAList a) a
ix Int
Index (NERAList a)
i