1{-# LANGUAGE CPP #-} 2{-# LANGUAGE GADTs #-} 3{-# LANGUAGE Rank2Types #-} 4{-# LANGUAGE TypeFamilies #-} 5{-# LANGUAGE FlexibleContexts #-} 6{-# LANGUAGE DefaultSignatures #-} 7{-# LANGUAGE FlexibleInstances #-} 8{-# LANGUAGE ScopedTypeVariables #-} 9{-# LANGUAGE MultiParamTypeClasses #-} 10{-# LANGUAGE UndecidableInstances #-} 11{-# LANGUAGE TypeOperators #-} 12 13#ifdef TRUSTWORTHY 14{-# LANGUAGE Trustworthy #-} 15#endif 16 17#include "lens-common.h" 18 19----------------------------------------------------------------------------- 20-- | 21-- Module : Control.Lens.At 22-- Copyright : (C) 2012-16 Edward Kmett 23-- License : BSD-style (see the file LICENSE) 24-- Maintainer : Edward Kmett <ekmett@gmail.com> 25-- Stability : experimental 26-- Portability : non-portable 27-- 28---------------------------------------------------------------------------- 29module Control.Lens.At 30 ( 31 -- * At 32 At(at) 33 , sans 34 , iat 35 -- * Ixed 36 , Index 37 , IxValue 38 , Ixed(ix) 39 , ixAt 40 , iix 41 -- * Contains 42 , Contains(contains) 43 , icontains 44 ) where 45 46import Prelude () 47 48import Control.Lens.Each 49import Control.Lens.Internal.Prelude 50import Control.Lens.Traversal 51import Control.Lens.Lens 52import Control.Lens.Setter 53import Control.Lens.Type 54import Control.Lens.Indexed 55import Data.Array.IArray as Array 56import Data.Array.Unboxed 57import Data.ByteString as StrictB 58import Data.ByteString.Lazy as LazyB 59import Data.Complex 60import Data.Hashable 61import Data.HashMap.Lazy as HashMap 62import Data.HashSet as HashSet 63import Data.Int 64import Data.IntMap as IntMap 65import Data.IntSet as IntSet 66import Data.Map as Map 67import Data.Set as Set 68import Data.Sequence as Seq 69import Data.Text as StrictT 70import Data.Text.Lazy as LazyT 71import Data.Tree 72import Data.Vector as Vector hiding (indexed) 73import Data.Vector.Primitive as Prim 74import Data.Vector.Storable as Storable 75import Data.Vector.Unboxed as Unboxed hiding (indexed) 76import Data.Word 77 78type family Index (s :: *) :: * 79type instance Index (e -> a) = e 80type instance Index IntSet = Int 81type instance Index (Set a) = a 82type instance Index (HashSet a) = a 83type instance Index [a] = Int 84type instance Index (NonEmpty a) = Int 85type instance Index (Seq a) = Int 86type instance Index (a,b) = Int 87type instance Index (a,b,c) = Int 88type instance Index (a,b,c,d) = Int 89type instance Index (a,b,c,d,e) = Int 90type instance Index (a,b,c,d,e,f) = Int 91type instance Index (a,b,c,d,e,f,g) = Int 92type instance Index (a,b,c,d,e,f,g,h) = Int 93type instance Index (a,b,c,d,e,f,g,h,i) = Int 94type instance Index (IntMap a) = Int 95type instance Index (Map k a) = k 96type instance Index (HashMap k a) = k 97type instance Index (Array.Array i e) = i 98type instance Index (UArray i e) = i 99type instance Index (Vector.Vector a) = Int 100type instance Index (Prim.Vector a) = Int 101type instance Index (Storable.Vector a) = Int 102type instance Index (Unboxed.Vector a) = Int 103type instance Index (Complex a) = Int 104type instance Index (Identity a) = () 105type instance Index (Maybe a) = () 106type instance Index (Tree a) = [Int] 107type instance Index StrictT.Text = Int 108type instance Index LazyT.Text = Int64 109type instance Index StrictB.ByteString = Int 110type instance Index LazyB.ByteString = Int64 111 112-- $setup 113-- >>> :set -XNoOverloadedStrings 114-- >>> import Control.Lens 115-- >>> import Debug.SimpleReflect.Expr 116-- >>> import Debug.SimpleReflect.Vars as Vars hiding (f,g) 117-- >>> let f :: Expr -> Expr; f = Debug.SimpleReflect.Vars.f 118-- >>> let g :: Expr -> Expr; g = Debug.SimpleReflect.Vars.g 119-- >>> let f' :: Int -> Expr -> Expr; f' = Debug.SimpleReflect.Vars.f' 120-- >>> let h :: Int -> Expr; h = Debug.SimpleReflect.Vars.h 121 122-- | 123-- This class provides a simple 'Lens' that lets you view (and modify) 124-- information about whether or not a container contains a given 'Index'. 125class Contains m where 126 -- | 127 -- >>> IntSet.fromList [1,2,3,4] ^. contains 3 128 -- True 129 -- 130 -- >>> IntSet.fromList [1,2,3,4] ^. contains 5 131 -- False 132 -- 133 -- >>> IntSet.fromList [1,2,3,4] & contains 3 .~ False 134 -- fromList [1,2,4] 135 contains :: Index m -> Lens' m Bool 136 137-- | An indexed version of 'contains'. 138-- 139-- >>> IntSet.fromList [1,2,3,4] ^@. icontains 3 140-- (3,True) 141-- 142-- >>> IntSet.fromList [1,2,3,4] ^@. icontains 5 143-- (5,False) 144-- 145-- >>> IntSet.fromList [1,2,3,4] & icontains 3 %@~ \i x -> if odd i then not x else x 146-- fromList [1,2,4] 147-- 148-- >>> IntSet.fromList [1,2,3,4] & icontains 3 %@~ \i x -> if even i then not x else x 149-- fromList [1,2,3,4] 150icontains :: Contains m => Index m -> IndexedLens' (Index m) m Bool 151icontains i f = contains i (indexed f i) 152{-# INLINE icontains #-} 153 154instance Contains IntSet where 155 contains k f s = f (IntSet.member k s) <&> \b -> 156 if b then IntSet.insert k s else IntSet.delete k s 157 {-# INLINE contains #-} 158 159instance Ord a => Contains (Set a) where 160 contains k f s = f (Set.member k s) <&> \b -> 161 if b then Set.insert k s else Set.delete k s 162 {-# INLINE contains #-} 163 164instance (Eq a, Hashable a) => Contains (HashSet a) where 165 contains k f s = f (HashSet.member k s) <&> \b -> 166 if b then HashSet.insert k s else HashSet.delete k s 167 {-# INLINE contains #-} 168 169-- | This provides a common notion of a value at an index that is shared by both 'Ixed' and 'At'. 170type family IxValue (m :: *) :: * 171 172-- | Provides a simple 'Traversal' lets you 'traverse' the value at a given 173-- key in a 'Map' or element at an ordinal position in a list or 'Seq'. 174class Ixed m where 175 -- | 176 -- /NB:/ Setting the value of this 'Traversal' will only set the value in 177 -- 'at' if it is already present. 178 -- 179 -- If you want to be able to insert /missing/ values, you want 'at'. 180 -- 181 -- >>> Seq.fromList [a,b,c,d] & ix 2 %~ f 182 -- fromList [a,b,f c,d] 183 -- 184 -- >>> Seq.fromList [a,b,c,d] & ix 2 .~ e 185 -- fromList [a,b,e,d] 186 -- 187 -- >>> Seq.fromList [a,b,c,d] ^? ix 2 188 -- Just c 189 -- 190 -- >>> Seq.fromList [] ^? ix 2 191 -- Nothing 192 ix :: Index m -> Traversal' m (IxValue m) 193 default ix :: (Applicative f, At m) => Index m -> LensLike' f m (IxValue m) 194 ix = ixAt 195 {-# INLINE ix #-} 196 197-- | An indexed version of 'ix'. 198-- 199-- >>> Seq.fromList [a,b,c,d] & iix 2 %@~ f' 200-- fromList [a,b,f' 2 c,d] 201-- 202-- >>> Seq.fromList [a,b,c,d] & iix 2 .@~ h 203-- fromList [a,b,h 2,d] 204-- 205-- >>> Seq.fromList [a,b,c,d] ^@? iix 2 206-- Just (2,c) 207-- 208-- >>> Seq.fromList [] ^@? iix 2 209-- Nothing 210iix :: Ixed m => Index m -> IndexedTraversal' (Index m) m (IxValue m) 211iix i f = ix i (indexed f i) 212{-# INLINE iix #-} 213 214-- | A definition of 'ix' for types with an 'At' instance. This is the default 215-- if you don't specify a definition for 'ix'. 216ixAt :: At m => Index m -> Traversal' m (IxValue m) 217ixAt i = at i . traverse 218{-# INLINE ixAt #-} 219 220type instance IxValue (e -> a) = a 221instance Eq e => Ixed (e -> a) where 222 ix e p f = p (f e) <&> \a e' -> if e == e' then a else f e' 223 {-# INLINE ix #-} 224 225type instance IxValue (Maybe a) = a 226instance Ixed (Maybe a) where 227 ix () f (Just a) = Just <$> f a 228 ix () _ Nothing = pure Nothing 229 {-# INLINE ix #-} 230 231type instance IxValue [a] = a 232instance Ixed [a] where 233 ix k f xs0 | k < 0 = pure xs0 234 | otherwise = go xs0 k where 235 go [] _ = pure [] 236 go (a:as) 0 = f a <&> (:as) 237 go (a:as) i = (a:) <$> (go as $! i - 1) 238 {-# INLINE ix #-} 239 240type instance IxValue (NonEmpty a) = a 241instance Ixed (NonEmpty a) where 242 ix k f xs0 | k < 0 = pure xs0 243 | otherwise = go xs0 k where 244 go (a:|as) 0 = f a <&> (:|as) 245 go (a:|as) i = (a:|) <$> ix (i - 1) f as 246 {-# INLINE ix #-} 247 248type instance IxValue (Identity a) = a 249instance Ixed (Identity a) where 250 ix () f (Identity a) = Identity <$> f a 251 {-# INLINE ix #-} 252 253type instance IxValue (Tree a) = a 254instance Ixed (Tree a) where 255 ix xs0 f = go xs0 where 256 go [] (Node a as) = f a <&> \a' -> Node a' as 257 go (i:is) t@(Node a as) 258 | i < 0 = pure t 259 | otherwise = Node a <$> ix i (go is) as 260 {-# INLINE ix #-} 261 262type instance IxValue (Seq a) = a 263instance Ixed (Seq a) where 264 ix i f m 265 | 0 <= i && i < Seq.length m = f (Seq.index m i) <&> \a -> Seq.update i a m 266 | otherwise = pure m 267 {-# INLINE ix #-} 268 269type instance IxValue (IntMap a) = a 270instance Ixed (IntMap a) where 271 ix k f m = case IntMap.lookup k m of 272 Just v -> f v <&> \v' -> IntMap.insert k v' m 273 Nothing -> pure m 274 {-# INLINE ix #-} 275 276type instance IxValue (Map k a) = a 277instance Ord k => Ixed (Map k a) where 278 ix k f m = case Map.lookup k m of 279 Just v -> f v <&> \v' -> Map.insert k v' m 280 Nothing -> pure m 281 {-# INLINE ix #-} 282 283type instance IxValue (HashMap k a) = a 284instance (Eq k, Hashable k) => Ixed (HashMap k a) where 285 ix k f m = case HashMap.lookup k m of 286 Just v -> f v <&> \v' -> HashMap.insert k v' m 287 Nothing -> pure m 288 {-# INLINE ix #-} 289 290type instance IxValue (Set k) = () 291instance Ord k => Ixed (Set k) where 292 ix k f m = if Set.member k m 293 then f () <&> \() -> Set.insert k m 294 else pure m 295 {-# INLINE ix #-} 296 297type instance IxValue IntSet = () 298instance Ixed IntSet where 299 ix k f m = if IntSet.member k m 300 then f () <&> \() -> IntSet.insert k m 301 else pure m 302 {-# INLINE ix #-} 303 304type instance IxValue (HashSet k) = () 305instance (Eq k, Hashable k) => Ixed (HashSet k) where 306 ix k f m = if HashSet.member k m 307 then f () <&> \() -> HashSet.insert k m 308 else pure m 309 {-# INLINE ix #-} 310 311type instance IxValue (Array.Array i e) = e 312-- | 313-- @ 314-- arr '!' i ≡ arr 'Control.Lens.Getter.^.' 'ix' i 315-- arr '//' [(i,e)] ≡ 'ix' i 'Control.Lens.Setter..~' e '$' arr 316-- @ 317instance Ix i => Ixed (Array.Array i e) where 318 ix i f arr 319 | inRange (bounds arr) i = f (arr Array.! i) <&> \e -> arr Array.// [(i,e)] 320 | otherwise = pure arr 321 {-# INLINE ix #-} 322 323type instance IxValue (UArray i e) = e 324-- | 325-- @ 326-- arr '!' i ≡ arr 'Control.Lens.Getter.^.' 'ix' i 327-- arr '//' [(i,e)] ≡ 'ix' i 'Control.Lens.Setter..~' e '$' arr 328-- @ 329instance (IArray UArray e, Ix i) => Ixed (UArray i e) where 330 ix i f arr 331 | inRange (bounds arr) i = f (arr Array.! i) <&> \e -> arr Array.// [(i,e)] 332 | otherwise = pure arr 333 {-# INLINE ix #-} 334 335type instance IxValue (Vector.Vector a) = a 336instance Ixed (Vector.Vector a) where 337 ix i f v 338 | 0 <= i && i < Vector.length v = f (v Vector.! i) <&> \a -> v Vector.// [(i, a)] 339 | otherwise = pure v 340 {-# INLINE ix #-} 341 342type instance IxValue (Prim.Vector a) = a 343instance Prim a => Ixed (Prim.Vector a) where 344 ix i f v 345 | 0 <= i && i < Prim.length v = f (v Prim.! i) <&> \a -> v Prim.// [(i, a)] 346 | otherwise = pure v 347 {-# INLINE ix #-} 348 349type instance IxValue (Storable.Vector a) = a 350instance Storable a => Ixed (Storable.Vector a) where 351 ix i f v 352 | 0 <= i && i < Storable.length v = f (v Storable.! i) <&> \a -> v Storable.// [(i, a)] 353 | otherwise = pure v 354 {-# INLINE ix #-} 355 356type instance IxValue (Unboxed.Vector a) = a 357instance Unbox a => Ixed (Unboxed.Vector a) where 358 ix i f v 359 | 0 <= i && i < Unboxed.length v = f (v Unboxed.! i) <&> \a -> v Unboxed.// [(i, a)] 360 | otherwise = pure v 361 {-# INLINE ix #-} 362 363type instance IxValue StrictT.Text = Char 364instance Ixed StrictT.Text where 365 ix e f s = case StrictT.splitAt e s of 366 (l, mr) -> case StrictT.uncons mr of 367 Nothing -> pure s 368 Just (c, xs) -> f c <&> \d -> StrictT.concat [l, StrictT.singleton d, xs] 369 {-# INLINE ix #-} 370 371type instance IxValue LazyT.Text = Char 372instance Ixed LazyT.Text where 373 ix e f s = case LazyT.splitAt e s of 374 (l, mr) -> case LazyT.uncons mr of 375 Nothing -> pure s 376 Just (c, xs) -> f c <&> \d -> LazyT.append l (LazyT.cons d xs) 377 {-# INLINE ix #-} 378 379type instance IxValue StrictB.ByteString = Word8 380instance Ixed StrictB.ByteString where 381 ix e f s = case StrictB.splitAt e s of 382 (l, mr) -> case StrictB.uncons mr of 383 Nothing -> pure s 384 Just (c, xs) -> f c <&> \d -> StrictB.concat [l, StrictB.singleton d, xs] 385 {-# INLINE ix #-} 386 387type instance IxValue LazyB.ByteString = Word8 388instance Ixed LazyB.ByteString where 389 -- TODO: we could be lazier, returning each chunk as it is passed 390 ix e f s = case LazyB.splitAt e s of 391 (l, mr) -> case LazyB.uncons mr of 392 Nothing -> pure s 393 Just (c, xs) -> f c <&> \d -> LazyB.append l (LazyB.cons d xs) 394 {-# INLINE ix #-} 395 396 397 398-- | 'At' provides a 'Lens' that can be used to read, 399-- write or delete the value associated with a key in a 'Map'-like 400-- container on an ad hoc basis. 401-- 402-- An instance of 'At' should satisfy: 403-- 404-- @ 405-- 'ix' k ≡ 'at' k '.' 'traverse' 406-- @ 407class Ixed m => At m where 408 -- | 409 -- >>> Map.fromList [(1,"world")] ^.at 1 410 -- Just "world" 411 -- 412 -- >>> at 1 ?~ "hello" $ Map.empty 413 -- fromList [(1,"hello")] 414 -- 415 -- /Note:/ 'Map'-like containers form a reasonable instance, but not 'Array'-like ones, where 416 -- you cannot satisfy the 'Lens' laws. 417 at :: Index m -> Lens' m (Maybe (IxValue m)) 418 419-- | Delete the value associated with a key in a 'Map'-like container 420-- 421-- @ 422-- 'sans' k = 'at' k .~ Nothing 423-- @ 424sans :: At m => Index m -> m -> m 425sans k m = m & at k .~ Nothing 426{-# INLINE sans #-} 427 428-- | An indexed version of 'at'. 429-- 430-- >>> Map.fromList [(1,"world")] ^@. iat 1 431-- (1,Just "world") 432-- 433-- >>> iat 1 %@~ (\i x -> if odd i then Just "hello" else Nothing) $ Map.empty 434-- fromList [(1,"hello")] 435-- 436-- >>> iat 2 %@~ (\i x -> if odd i then Just "hello" else Nothing) $ Map.empty 437-- fromList [] 438-- 439iat :: At m => Index m -> IndexedLens' (Index m) m (Maybe (IxValue m)) 440iat i f = at i (indexed f i) 441{-# INLINE iat #-} 442 443instance At (Maybe a) where 444 at () f = f 445 {-# INLINE at #-} 446 447instance At (IntMap a) where 448#if MIN_VERSION_containers(0,5,8) 449 at k f = IntMap.alterF f k 450#else 451 at k f m = f mv <&> \r -> case r of 452 Nothing -> maybe m (const (IntMap.delete k m)) mv 453 Just v' -> IntMap.insert k v' m 454 where mv = IntMap.lookup k m 455#endif 456 {-# INLINE at #-} 457 458instance Ord k => At (Map k a) where 459#if MIN_VERSION_containers(0,5,8) 460 at k f = Map.alterF f k 461#else 462 at k f m = f mv <&> \r -> case r of 463 Nothing -> maybe m (const (Map.delete k m)) mv 464 Just v' -> Map.insert k v' m 465 where mv = Map.lookup k m 466#endif 467 {-# INLINE at #-} 468 469instance (Eq k, Hashable k) => At (HashMap k a) where 470 at k f m = f mv <&> \r -> case r of 471 Nothing -> maybe m (const (HashMap.delete k m)) mv 472 Just v' -> HashMap.insert k v' m 473 where mv = HashMap.lookup k m 474 {-# INLINE at #-} 475 476instance At IntSet where 477 at k f m = f mv <&> \r -> case r of 478 Nothing -> maybe m (const (IntSet.delete k m)) mv 479 Just () -> IntSet.insert k m 480 where mv = if IntSet.member k m then Just () else Nothing 481 {-# INLINE at #-} 482 483instance Ord k => At (Set k) where 484 at k f m = f mv <&> \r -> case r of 485 Nothing -> maybe m (const (Set.delete k m)) mv 486 Just () -> Set.insert k m 487 where mv = if Set.member k m then Just () else Nothing 488 {-# INLINE at #-} 489 490instance (Eq k, Hashable k) => At (HashSet k) where 491 at k f m = f mv <&> \r -> case r of 492 Nothing -> maybe m (const (HashSet.delete k m)) mv 493 Just () -> HashSet.insert k m 494 where mv = if HashSet.member k m then Just () else Nothing 495 {-# INLINE at #-} 496 497 498-- | @'ix' :: 'Int' -> 'Traversal'' (a,a) a@ 499type instance IxValue (a,a2) = a 500instance (a~a2) => Ixed (a,a2) where 501 ix = elementOf each 502 503-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a) a@ 504type instance IxValue (a,a2,a3) = a 505instance (a~a2, a~a3) => Ixed (a,a2,a3) where 506 ix = elementOf each 507 508-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a) a@ 509type instance IxValue (a,a2,a3,a4) = a 510instance (a~a2, a~a3, a~a4) => Ixed (a,a2,a3,a4) where 511 ix = elementOf each 512 513-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a) a@ 514type instance IxValue (a,a2,a3,a4,a5) = a 515instance (a~a2, a~a3, a~a4, a~a5) => Ixed (a,a2,a3,a4,a5) where 516 ix = elementOf each 517 518-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a,a) a@ 519type instance IxValue (a,a2,a3,a4,a5,a6) = a 520instance (a~a2, a~a3, a~a4, a~a5, a~a6) => Ixed (a,a2,a3,a4,a5,a6) where 521 ix = elementOf each 522 523-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a,a,a) a@ 524type instance IxValue (a,a2,a3,a4,a5,a6,a7) = a 525instance (a~a2, a~a3, a~a4, a~a5, a~a6, a~a7) => Ixed (a,a2,a3,a4,a5,a6,a7) where 526 ix = elementOf each 527 528-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a,a,a,a) a@ 529type instance IxValue (a,a2,a3,a4,a5,a6,a7,a8) = a 530instance (a~a2, a~a3, a~a4, a~a5, a~a6, a~a7, a~a8) => Ixed (a,a2,a3,a4,a5,a6,a7,a8) where 531 ix = elementOf each 532 533-- | @'ix' :: 'Int' -> 'Traversal'' (a,a,a,a,a,a,a,a,a) a@ 534type instance IxValue (a,a2,a3,a4,a5,a6,a7,a8,a9) = a 535instance (a~a2, a~a3, a~a4, a~a5, a~a6, a~a7, a~a8, a~a9) => Ixed (a,a2,a3,a4,a5,a6,a7,a8,a9) where 536 ix = elementOf each 537