{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module: Optics.At
-- Description: Optics for 'Data.Map.Map' and 'Data.Set.Set'-like containers.
--
-- This module provides optics for 'Data.Map.Map' and 'Data.Set.Set'-like
-- containers, including an 'AffineTraversal' to traverse a key in a map or an
-- element of a sequence:
--
-- >>> preview (ix 1) ['a','b','c']
-- Just 'b'
--
-- a 'Lens' to get, set or delete a key in a map:
--
-- >>> set (at 0) (Just 'b') (Map.fromList [(0, 'a')])
-- fromList [(0,'b')]
--
-- and a 'Lens' to insert or remove an element of a set:
--
-- >>> IntSet.fromList [1,2,3,4] & contains 3 .~ False
-- fromList [1,2,4]
--
-- This module includes the core definitions from "Optics.At.Core" along with
-- extra (orphan) instances.
--
module Optics.At
  (
    -- * Type families
    Index
  , IxValue

    -- * Ixed
  , Ixed(..)
  , ixAt

    -- * At
  , At(..)
  , at'
  , sans

  -- * Contains
  , Contains(..)
  ) where

import qualified Data.ByteString as StrictB
import qualified Data.ByteString.Lazy as LazyB
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as HashMap
import Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
import Data.Hashable (Hashable)
import Data.Int (Int64)
import qualified Data.Text as StrictT
import qualified Data.Text.Lazy as LazyT
import qualified Data.Vector as Vector hiding (indexed)
import qualified Data.Vector.Primitive as Prim
import qualified Data.Vector.Storable as Storable
import qualified Data.Vector.Unboxed as Unboxed hiding (indexed)
import Data.Word (Word8)

import Optics.Core

type instance Index (HashSet a) = a
type instance Index (HashMap k a) = k
type instance Index (Vector.Vector a) = Int
type instance Index (Prim.Vector a) = Int
type instance Index (Storable.Vector a) = Int
type instance Index (Unboxed.Vector a) = Int
type instance Index StrictT.Text = Int
type instance Index LazyT.Text = Int64
type instance Index StrictB.ByteString = Int
type instance Index LazyB.ByteString = Int64

-- Contains

instance (Eq a, Hashable a) => Contains (HashSet a) where
  contains :: Index (HashSet a) -> Lens' (HashSet a) Bool
contains Index (HashSet a)
k = LensVL (HashSet a) (HashSet a) Bool Bool -> Lens' (HashSet a) Bool
forall s t a b. LensVL s t a b -> Lens s t a b
lensVL (LensVL (HashSet a) (HashSet a) Bool Bool
 -> Lens' (HashSet a) Bool)
-> LensVL (HashSet a) (HashSet a) Bool Bool
-> Lens' (HashSet a) Bool
forall a b. (a -> b) -> a -> b
$ \Bool -> f Bool
f HashSet a
s -> Bool -> f Bool
f (a -> HashSet a -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
HashSet.member a
Index (HashSet a)
k HashSet a
s) f Bool -> (Bool -> HashSet a) -> f (HashSet a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Bool
b ->
    if Bool
b then a -> HashSet a -> HashSet a
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.insert a
Index (HashSet a)
k HashSet a
s else a -> HashSet a -> HashSet a
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.delete a
Index (HashSet a)
k HashSet a
s
  {-# INLINE contains #-}

-- Ixed

type instance IxValue (HashMap k a) = a
-- Default implementation uses HashMap.alterF
instance (Eq k, Hashable k) => Ixed (HashMap k a)

type instance IxValue (HashSet k) = ()
instance (Eq k, Hashable k) => Ixed (HashSet k) where
  ix :: Index (HashSet k)
-> Optic'
     (IxKind (HashSet k)) NoIx (HashSet k) (IxValue (HashSet k))
ix Index (HashSet k)
k = AffineTraversalVL (HashSet k) (HashSet k) () ()
-> AffineTraversal (HashSet k) (HashSet k) () ()
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL (HashSet k) (HashSet k) () ()
 -> AffineTraversal (HashSet k) (HashSet k) () ())
-> AffineTraversalVL (HashSet k) (HashSet k) () ()
-> AffineTraversal (HashSet k) (HashSet k) () ()
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point () -> f ()
f HashSet k
m ->
    if k -> HashSet k -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
HashSet.member k
Index (HashSet k)
k HashSet k
m
    then () -> f ()
f () f () -> (() -> HashSet k) -> f (HashSet k)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \() -> k -> HashSet k -> HashSet k
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.insert k
Index (HashSet k)
k HashSet k
m
    else HashSet k -> f (HashSet k)
forall r. r -> f r
point HashSet k
m
  {-# INLINE ix #-}

type instance IxValue (Vector.Vector a) = a
instance Ixed (Vector.Vector a) where
  ix :: Index (Vector a)
-> Optic' (IxKind (Vector a)) NoIx (Vector a) (IxValue (Vector a))
ix Index (Vector a)
i = AffineTraversalVL (Vector a) (Vector a) a a
-> AffineTraversal (Vector a) (Vector a) a a
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL (Vector a) (Vector a) a a
 -> AffineTraversal (Vector a) (Vector a) a a)
-> AffineTraversalVL (Vector a) (Vector a) a a
-> AffineTraversal (Vector a) (Vector a) a a
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point a -> f a
f Vector a
v ->
    if Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
Index (Vector a)
i Bool -> Bool -> Bool
&& Int
Index (Vector a)
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector a -> Int
forall a. Vector a -> Int
Vector.length Vector a
v
    then a -> f a
f (Vector a
v Vector a -> Int -> a
forall a. Vector a -> Int -> a
Vector.! Int
Index (Vector a)
i) f a -> (a -> Vector a) -> f (Vector a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> Vector a
v Vector a -> [(Int, a)] -> Vector a
forall a. Vector a -> [(Int, a)] -> Vector a
Vector.// [(Int
Index (Vector a)
i, a
a)]
    else Vector a -> f (Vector a)
forall r. r -> f r
point Vector a
v
  {-# INLINE ix #-}

type instance IxValue (Prim.Vector a) = a
instance Prim.Prim a => Ixed (Prim.Vector a) where
  ix :: Index (Vector a)
-> Optic' (IxKind (Vector a)) NoIx (Vector a) (IxValue (Vector a))
ix Index (Vector a)
i = AffineTraversalVL (Vector a) (Vector a) a a
-> AffineTraversal (Vector a) (Vector a) a a
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL (Vector a) (Vector a) a a
 -> AffineTraversal (Vector a) (Vector a) a a)
-> AffineTraversalVL (Vector a) (Vector a) a a
-> AffineTraversal (Vector a) (Vector a) a a
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point a -> f a
f Vector a
v ->
    if Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
Index (Vector a)
i Bool -> Bool -> Bool
&& Int
Index (Vector a)
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector a -> Int
forall a. Prim a => Vector a -> Int
Prim.length Vector a
v
    then a -> f a
f (Vector a
v Vector a -> Int -> a
forall a. Prim a => Vector a -> Int -> a
Prim.! Int
Index (Vector a)
i) f a -> (a -> Vector a) -> f (Vector a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> Vector a
v Vector a -> [(Int, a)] -> Vector a
forall a. Prim a => Vector a -> [(Int, a)] -> Vector a
Prim.// [(Int
Index (Vector a)
i, a
a)]
    else Vector a -> f (Vector a)
forall r. r -> f r
point Vector a
v
  {-# INLINE ix #-}

type instance IxValue (Storable.Vector a) = a
instance Storable.Storable a => Ixed (Storable.Vector a) where
  ix :: Index (Vector a)
-> Optic' (IxKind (Vector a)) NoIx (Vector a) (IxValue (Vector a))
ix Index (Vector a)
i = AffineTraversalVL (Vector a) (Vector a) a a
-> AffineTraversal (Vector a) (Vector a) a a
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL (Vector a) (Vector a) a a
 -> AffineTraversal (Vector a) (Vector a) a a)
-> AffineTraversalVL (Vector a) (Vector a) a a
-> AffineTraversal (Vector a) (Vector a) a a
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point a -> f a
f Vector a
v ->
    if Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
Index (Vector a)
i Bool -> Bool -> Bool
&& Int
Index (Vector a)
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector a -> Int
forall a. Storable a => Vector a -> Int
Storable.length Vector a
v
    then a -> f a
f (Vector a
v Vector a -> Int -> a
forall a. Storable a => Vector a -> Int -> a
Storable.! Int
Index (Vector a)
i) f a -> (a -> Vector a) -> f (Vector a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> Vector a
v Vector a -> [(Int, a)] -> Vector a
forall a. Storable a => Vector a -> [(Int, a)] -> Vector a
Storable.// [(Int
Index (Vector a)
i, a
a)]
    else Vector a -> f (Vector a)
forall r. r -> f r
point Vector a
v
  {-# INLINE ix #-}

type instance IxValue (Unboxed.Vector a) = a
instance Unboxed.Unbox a => Ixed (Unboxed.Vector a) where
  ix :: Index (Vector a)
-> Optic' (IxKind (Vector a)) NoIx (Vector a) (IxValue (Vector a))
ix Index (Vector a)
i = AffineTraversalVL (Vector a) (Vector a) a a
-> AffineTraversal (Vector a) (Vector a) a a
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL (Vector a) (Vector a) a a
 -> AffineTraversal (Vector a) (Vector a) a a)
-> AffineTraversalVL (Vector a) (Vector a) a a
-> AffineTraversal (Vector a) (Vector a) a a
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point a -> f a
f Vector a
v ->
    if Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
Index (Vector a)
i Bool -> Bool -> Bool
&& Int
Index (Vector a)
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Vector a -> Int
forall a. Unbox a => Vector a -> Int
Unboxed.length Vector a
v
    then a -> f a
f (Vector a
v Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
Unboxed.! Int
Index (Vector a)
i) f a -> (a -> Vector a) -> f (Vector a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
a -> Vector a
v Vector a -> [(Int, a)] -> Vector a
forall a. Unbox a => Vector a -> [(Int, a)] -> Vector a
Unboxed.// [(Int
Index (Vector a)
i, a
a)]
    else Vector a -> f (Vector a)
forall r. r -> f r
point Vector a
v
  {-# INLINE ix #-}

type instance IxValue StrictT.Text = Char
instance Ixed StrictT.Text where
  ix :: Index Text -> Optic' (IxKind Text) NoIx Text (IxValue Text)
ix Index Text
e = AffineTraversalVL Text Text Char Char
-> AffineTraversal Text Text Char Char
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL Text Text Char Char
 -> AffineTraversal Text Text Char Char)
-> AffineTraversalVL Text Text Char Char
-> AffineTraversal Text Text Char Char
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point Char -> f Char
f Text
s ->
    case Int -> Text -> (Text, Text)
StrictT.splitAt Int
Index Text
e Text
s of
      (Text
l, Text
mr) -> case Text -> Maybe (Char, Text)
StrictT.uncons Text
mr of
        Maybe (Char, Text)
Nothing      -> Text -> f Text
forall r. r -> f r
point Text
s
        Just (Char
c, Text
xs) -> Char -> f Char
f Char
c f Char -> (Char -> Text) -> f Text
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Char
d -> [Text] -> Text
StrictT.concat [Text
l, Char -> Text
StrictT.singleton Char
d, Text
xs]
  {-# INLINE ix #-}

type instance IxValue LazyT.Text = Char
instance Ixed LazyT.Text where
  ix :: Index Text -> Optic' (IxKind Text) NoIx Text (IxValue Text)
ix Index Text
e = AffineTraversalVL Text Text Char Char
-> AffineTraversal Text Text Char Char
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL Text Text Char Char
 -> AffineTraversal Text Text Char Char)
-> AffineTraversalVL Text Text Char Char
-> AffineTraversal Text Text Char Char
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point Char -> f Char
f Text
s ->
    case Int64 -> Text -> (Text, Text)
LazyT.splitAt Int64
Index Text
e Text
s of
      (Text
l, Text
mr) -> case Text -> Maybe (Char, Text)
LazyT.uncons Text
mr of
        Maybe (Char, Text)
Nothing      -> Text -> f Text
forall r. r -> f r
point Text
s
        Just (Char
c, Text
xs) -> Char -> f Char
f Char
c f Char -> (Char -> Text) -> f Text
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Char
d -> Text -> Text -> Text
LazyT.append Text
l (Char -> Text -> Text
LazyT.cons Char
d Text
xs)
  {-# INLINE ix #-}

type instance IxValue StrictB.ByteString = Word8
instance Ixed StrictB.ByteString where
  ix :: Index ByteString
-> Optic' (IxKind ByteString) NoIx ByteString (IxValue ByteString)
ix Index ByteString
e = AffineTraversalVL ByteString ByteString Word8 Word8
-> AffineTraversal ByteString ByteString Word8 Word8
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL ByteString ByteString Word8 Word8
 -> AffineTraversal ByteString ByteString Word8 Word8)
-> AffineTraversalVL ByteString ByteString Word8 Word8
-> AffineTraversal ByteString ByteString Word8 Word8
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point Word8 -> f Word8
f ByteString
s ->
    case Int -> ByteString -> (ByteString, ByteString)
StrictB.splitAt Int
Index ByteString
e ByteString
s of
      (ByteString
l, ByteString
mr) -> case ByteString -> Maybe (Word8, ByteString)
StrictB.uncons ByteString
mr of
        Maybe (Word8, ByteString)
Nothing      -> ByteString -> f ByteString
forall r. r -> f r
point ByteString
s
        Just (Word8
c, ByteString
xs) -> Word8 -> f Word8
f Word8
c f Word8 -> (Word8 -> ByteString) -> f ByteString
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Word8
d -> [ByteString] -> ByteString
StrictB.concat [ByteString
l, Word8 -> ByteString
StrictB.singleton Word8
d, ByteString
xs]
  {-# INLINE ix #-}

type instance IxValue LazyB.ByteString = Word8
instance Ixed LazyB.ByteString where
  -- TODO: we could be lazier, returning each chunk as it is passed
  ix :: Index ByteString
-> Optic' (IxKind ByteString) NoIx ByteString (IxValue ByteString)
ix Index ByteString
e = AffineTraversalVL ByteString ByteString Word8 Word8
-> AffineTraversal ByteString ByteString Word8 Word8
forall s t a b.
AffineTraversalVL s t a b -> AffineTraversal s t a b
atraversalVL (AffineTraversalVL ByteString ByteString Word8 Word8
 -> AffineTraversal ByteString ByteString Word8 Word8)
-> AffineTraversalVL ByteString ByteString Word8 Word8
-> AffineTraversal ByteString ByteString Word8 Word8
forall a b. (a -> b) -> a -> b
$ \forall r. r -> f r
point Word8 -> f Word8
f ByteString
s ->
    case Int64 -> ByteString -> (ByteString, ByteString)
LazyB.splitAt Int64
Index ByteString
e ByteString
s of
      (ByteString
l, ByteString
mr) -> case ByteString -> Maybe (Word8, ByteString)
LazyB.uncons ByteString
mr of
        Maybe (Word8, ByteString)
Nothing      -> ByteString -> f ByteString
forall r. r -> f r
point ByteString
s
        Just (Word8
c, ByteString
xs) -> Word8 -> f Word8
f Word8
c f Word8 -> (Word8 -> ByteString) -> f ByteString
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Word8
d -> ByteString -> ByteString -> ByteString
LazyB.append ByteString
l (Word8 -> ByteString -> ByteString
LazyB.cons Word8
d ByteString
xs)
  {-# INLINE ix #-}

-- At

instance (Eq k, Hashable k) => At (HashMap k a) where
#if MIN_VERSION_unordered_containers(0,2,10)
  at :: Index (HashMap k a)
-> Lens' (HashMap k a) (Maybe (IxValue (HashMap k a)))
at Index (HashMap k a)
k = LensVL (HashMap k a) (HashMap k a) (Maybe a) (Maybe a)
-> Lens (HashMap k a) (HashMap k a) (Maybe a) (Maybe a)
forall s t a b. LensVL s t a b -> Lens s t a b
lensVL (LensVL (HashMap k a) (HashMap k a) (Maybe a) (Maybe a)
 -> Lens (HashMap k a) (HashMap k a) (Maybe a) (Maybe a))
-> LensVL (HashMap k a) (HashMap k a) (Maybe a) (Maybe a)
-> Lens (HashMap k a) (HashMap k a) (Maybe a) (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Maybe a -> f (Maybe a)
f -> (Maybe a -> f (Maybe a)) -> k -> HashMap k a -> f (HashMap k a)
forall (f :: * -> *) k v.
(Functor f, Eq k, Hashable k) =>
(Maybe v -> f (Maybe v)) -> k -> HashMap k v -> f (HashMap k v)
HashMap.alterF Maybe a -> f (Maybe a)
f k
Index (HashMap k a)
k
#else
  at k = lensVL $ \f m ->
    let mv = HashMap.lookup k m
    in f mv <&> \r -> case r of
      Nothing -> maybe m (const (HashMap.delete k m)) mv
      Just v' -> HashMap.insert k v' m
#endif
  {-# INLINE at #-}

instance (Eq k, Hashable k) => At (HashSet k) where
  at :: Index (HashSet k)
-> Lens' (HashSet k) (Maybe (IxValue (HashSet k)))
at Index (HashSet k)
k = LensVL (HashSet k) (HashSet k) (Maybe ()) (Maybe ())
-> Lens (HashSet k) (HashSet k) (Maybe ()) (Maybe ())
forall s t a b. LensVL s t a b -> Lens s t a b
lensVL (LensVL (HashSet k) (HashSet k) (Maybe ()) (Maybe ())
 -> Lens (HashSet k) (HashSet k) (Maybe ()) (Maybe ()))
-> LensVL (HashSet k) (HashSet k) (Maybe ()) (Maybe ())
-> Lens (HashSet k) (HashSet k) (Maybe ()) (Maybe ())
forall a b. (a -> b) -> a -> b
$ \Maybe () -> f (Maybe ())
f HashSet k
m ->
    let mv :: Maybe ()
mv = if k -> HashSet k -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
HashSet.member k
Index (HashSet k)
k HashSet k
m
             then () -> Maybe ()
forall a. a -> Maybe a
Just ()
             else Maybe ()
forall a. Maybe a
Nothing
    in Maybe () -> f (Maybe ())
f Maybe ()
mv f (Maybe ()) -> (Maybe () -> HashSet k) -> f (HashSet k)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Maybe ()
r -> case Maybe ()
r of
      Maybe ()
Nothing -> HashSet k -> (() -> HashSet k) -> Maybe () -> HashSet k
forall b a. b -> (a -> b) -> Maybe a -> b
maybe HashSet k
m (HashSet k -> () -> HashSet k
forall a b. a -> b -> a
const (k -> HashSet k -> HashSet k
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.delete k
Index (HashSet k)
k HashSet k
m)) Maybe ()
mv
      Just () -> k -> HashSet k -> HashSet k
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.insert k
Index (HashSet k)
k HashSet k
m
  {-# INLINE at #-}

-- $setup
-- >>> import qualified Data.IntSet as IntSet
-- >>> import qualified Data.Map as Map