-- Copyright 2013 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

{-# LANGUAGE TypeOperators, MultiParamTypeClasses #-}
-- | Re-export a minimal version of the fclabels API.
--
-- I'd like to be able to write @(a #= x) . (b #= y)@ without parentheses,
-- but since (.) already binds the strongest at 9, I can't make (#=) stronger.
-- Besides, I already want (#=) to bind loosely so I can write @x#y #= 1+2@.
-- I would need a version of (.) at 0, but that's too much trouble.
module Util.Lens (
    Lens, lens
    , (#)
    -- * operators
    , (#$), (#=), (%=)
    , (<#>)

    -- * data
    , map, set, list
) where
import           Prelude hiding ((.), map)
import           Control.Category ((.))
import qualified Data.Label as Label
import           Data.Label (get, lens, modify, (:->))
import qualified Data.Map as Map
import qualified Data.Set as Set

import qualified Util.Lists as Lists


-- * lens operators

type Lens a b = a :-> b

-- | Compose lenses.
(#) :: Lens a b -> Lens b c -> Lens a c
# :: forall a b c. Lens a b -> Lens b c -> Lens a c
(#) = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
(.)
infixr 9 #

-- | Get: @bval = a#b $# record@
(#$) :: Lens f a -> f -> a
#$ :: forall f a. Lens f a -> f -> a
(#$) = forall f a. Lens f a -> f -> a
get
infixr 1 #$

-- | Set: @a#b #= 42 record@
(#=) :: Lens f a -> a -> f -> f
#= :: forall f a. Lens f a -> a -> f -> f
(#=) = forall f a. Lens f a -> a -> f -> f
Label.set
infix 1 #=

-- | Modify: @a#b %= (+1) record@
(%=) :: Lens f a -> (a -> a) -> f -> f
%= :: forall f a. Lens f a -> (a -> a) -> f -> f
(%=) = forall f a. Lens f a -> (a -> a) -> f -> f
modify
infix 1 %=

-- | Use like @a#b <#> State.get@.
(<#>) :: Functor f => Lens a b -> f a -> f b
<#> :: forall (f :: * -> *) a b. Functor f => Lens a b -> f a -> f b
(<#>) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall f a. Lens f a -> f -> a
get
infixl 4 <#> -- same as <$>

-- * data

map :: Ord k => k -> Lens (Map.Map k a) (Maybe a)
map :: forall k a. Ord k => k -> Lens (Map k a) (Maybe a)
map k
k = forall f a. (f -> a) -> ((a -> a) -> f -> f) -> f :-> a
lens (forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k) (\Maybe a -> Maybe a
modify -> forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter Maybe a -> Maybe a
modify k
k)

set :: Ord k => k -> Lens (Set.Set k) Bool
set :: forall k. Ord k => k -> Lens (Set k) Bool
set k
k = forall f a. (f -> a) -> ((a -> a) -> f -> f) -> f :-> a
lens (forall a. Ord a => a -> Set a -> Bool
Set.member k
k) forall a b. (a -> b) -> a -> b
$ \Bool -> Bool
modify Set k
s ->
    if Bool -> Bool
modify (forall a. Ord a => a -> Set a -> Bool
Set.member k
k Set k
s) then forall a. Ord a => a -> Set a -> Set a
Set.insert k
k Set k
s else forall a. Ord a => a -> Set a -> Set a
Set.delete k
k Set k
s

-- | A negative index counts from the end of the list, and a too-high index
-- will be ignored.
list :: Int -> Lens [a] (Maybe a)
list :: forall a. Int -> Lens [a] (Maybe a)
list Int
i
    | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = forall f a. (f -> a) -> ((a -> a) -> f -> f) -> f :-> a
lens (\[a]
xs -> forall a. [a] -> Maybe a
Lists.head (forall a. Int -> [a] -> [a]
drop (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs forall a. Num a => a -> a -> a
+ Int
i) [a]
xs))
        (\Maybe a -> Maybe a
modify [a]
xs -> forall {a}. Int -> (Maybe a -> Maybe a) -> [a] -> [a]
at (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs forall a. Num a => a -> a -> a
+ Int
i) Maybe a -> Maybe a
modify [a]
xs)
    | Bool
otherwise = forall f a. (f -> a) -> ((a -> a) -> f -> f) -> f :-> a
lens (forall a. [a] -> Maybe a
Lists.head forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Int -> [a] -> [a]
drop Int
i) (forall {a}. Int -> (Maybe a -> Maybe a) -> [a] -> [a]
at Int
i)
    where
    at :: Int -> (Maybe a -> Maybe a) -> [a] -> [a]
at Int
i Maybe a -> Maybe a
modify [a]
xs = case Maybe a -> Maybe a
modify (forall a. [a] -> Int -> Maybe a
Lists.at [a]
xs Int
i) of
        Maybe a
Nothing -> forall a. Int -> [a] -> [a]
Lists.removeAt Int
i [a]
xs
        Just a
x -> forall a. Int -> (a -> a) -> [a] -> [a]
Lists.modifyAt Int
i (forall a b. a -> b -> a
const a
x) [a]
xs