-- 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 Data.Label ((:->), get, modify, lens)
import qualified Data.Label as Label
import qualified Data.Map as Map
import qualified Data.Set as Set

import qualified Util.Seq as Seq


-- * 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
(#) = (Lens b c -> Lens a b -> Lens a c)
-> Lens a b -> Lens b c -> Lens a c
forall a b c. (a -> b -> c) -> b -> a -> c
flip Lens b c -> Lens a b -> Lens a c
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
(#$) = (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
(#=) = (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
(%=) = (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
(<#>) = (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> f a -> f b)
-> (Lens a b -> a -> b) -> Lens a b -> f a -> f b
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Lens a b -> a -> b
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 = (Map k a -> Maybe a)
-> ((Maybe a -> Maybe a) -> Map k a -> Map k a)
-> Map k a :-> Maybe a
forall f a. (f -> a) -> ((a -> a) -> f -> f) -> f :-> a
lens (k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k) (\Maybe a -> Maybe a
modify -> (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
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 = (Set k -> Bool)
-> ((Bool -> Bool) -> Set k -> Set k) -> Set k :-> Bool
forall f a. (f -> a) -> ((a -> a) -> f -> f) -> f :-> a
lens (k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member k
k) (((Bool -> Bool) -> Set k -> Set k) -> Set k :-> Bool)
-> ((Bool -> Bool) -> Set k -> Set k) -> Set k :-> Bool
forall a b. (a -> b) -> a -> b
$ \Bool -> Bool
modify Set k
s ->
    if Bool -> Bool
modify (k -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member k
k Set k
s) then k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.insert k
k Set k
s else k -> Set k -> Set k
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = ([a] -> Maybe a)
-> ((Maybe a -> Maybe a) -> [a] -> [a]) -> [a] :-> Maybe a
forall f a. (f -> a) -> ((a -> a) -> f -> f) -> f :-> a
lens (\[a]
xs -> [a] -> Maybe a
forall a. [a] -> Maybe a
Seq.head (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) [a]
xs))
        (\Maybe a -> Maybe a
modify [a]
xs -> Int -> (Maybe a -> Maybe a) -> [a] -> [a]
forall {a}. Int -> (Maybe a -> Maybe a) -> [a] -> [a]
at ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) Maybe a -> Maybe a
modify [a]
xs)
    | Bool
otherwise = ([a] -> Maybe a)
-> ((Maybe a -> Maybe a) -> [a] -> [a]) -> [a] :-> Maybe a
forall f a. (f -> a) -> ((a -> a) -> f -> f) -> f :-> a
lens ([a] -> Maybe a
forall a. [a] -> Maybe a
Seq.head ([a] -> Maybe a) -> ([a] -> [a]) -> [a] -> Maybe a
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
i) (Int -> (Maybe a -> Maybe a) -> [a] -> [a]
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 ([a] -> Int -> Maybe a
forall a. [a] -> Int -> Maybe a
Seq.at [a]
xs Int
i) of
        Maybe a
Nothing -> Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
Seq.remove_at Int
i [a]
xs
        Just a
x -> Int -> (a -> a) -> [a] -> [a]
forall a. Int -> (a -> a) -> [a] -> [a]
Seq.modify_at Int
i (a -> a -> a
forall a b. a -> b -> a
const a
x) [a]
xs