Composing lenses with `at` and `ix`

Question

Let's say I have some fairly simple data type Person with a couple of fields, and a type that holds a collection of Persons.

data Person = Person { _name :: String, _age  :: Int }

data ProgramState = PS { _dict :: IntMap Person }

makeLenses ''Person
makeLenses ''ProgramState

I want to create a lens that allows me to access individual people by looking up their key

person :: Int -> Lens' ProgramState Person

It seems my two options for doing this are to use at or ix to index into the dictionary

-- Option 1, using 'at'
person :: Int -> Lens' ProgramState (Maybe Person)
person key = dict . at key

-- Option 2, using 'ix'
person :: Int -> Traversal' ProgramState Person
person key = dict . ix key

but neither of these options lets me do what I want, which is to have a Lens' that accesses a Person rather than a Maybe Person. Option 1 doesn't compose nicely with other lenses, and option 2 means that I have to give up my getters.

I understand why ix and at are written like this. The key might not exist in the dict, so if you want a Lens' which enables both getters and setters, it must access a Maybe a. The alternative is to accept a Traversal' which gives access to 0 or 1 values, but that means giving up your getters. But in my case, I know that the element I want will always be present, so I don't need to worry about missing keys.

Is there a way to write what I want to write - or should I be rethinking the structure of my program?

Answer 1

You probably want to use at together with the non isomorphism. You can specify a default map entry with it to get rid of the Maybe of the lookup.

non :: Eq a => a -> Iso' (Maybe a) a

person key = dict . at key . non defaultEntry

-- can get and set just like plain lenses
someProgramState & dict . at someKey . non defaultEntry .~ somePerson

You can look at more examples in the docs.

Answer 2

Based on András Kovács answer I ended up defining an unsafeFromJust lens that witnesses the 'isomorphism' I require to compose these lenses

import Data.Maybe (fromJust)

unsafeFromJust :: Lens' (Maybe a) a
unsafeFromJust = lens fromJust setJust
 where
  setJust (Just _) b = Just a
  setJust Nothing  _ = error "setJust: Nothing"

An alternative definition is

unsafeFromJust :: Lens' (Maybe a) a
unsafeFromJust = anon (error "unsafeFromJust: Nothing") (\_ -> False)

but I felt that wasn't as clear as the first form. I didn't use non as that requires an Eq instance that is unnecessary in this case.

I can now write

person :: Lens' ProgramState Person
person key = dict . at key . unsafeFromJust

Answer 3

Actually, @Chris Taylor's answer is incorrect. You can see this with the following commands (in GHCi):

>view (at 0 . unsafeFromJust) (fromList [(0,'b')])
*** Exception: Maybe.fromJust: Nothing
[expected: *** Exception: setJust]

>set (at 1 . unsafeFromJust) 'c' (fromList [(0,'b')])
*** Exception: setJust]
[expected: fromList [(0,'b'),(1,'c')]

The first command still throws an error when looking up a value that doesn't exist, but it doesn't throw the right error. For the second test, I am unable to insert new keys, which doesn't seem to make sense.

Instead, I'm using the following two combinators:

at' :: (Ord a) => a -> Lens' (Map a b) b
at' a = lens r s
  where r m = case lookup a m of
                (Just b) -> b
                Nothing -> error "Could not find key in map!"
        s m b' = insert a b' m

at'' :: (Ord a) => a -> Lens (Map a b) (Map a b) b (Maybe b)
at'' a = lens r s
  where r m = case lookup a m of
                (Just b) -> b
                Nothing -> error "Could not find key in map!"
        s m Nothing = delete a m
        s m (Just b') = insert a b' m

at' is how (at k . unsafeFromJust) should work: attempting to retrieve a non-existent value throws an error, inserting a new value succeeds. at'' is similar: it allows you to read pure values, but you set Maybe values. This allows you to delete keys in the map.

Examples:

> view (at' 0) (fromList [(0,'b')])
'b'

> view (at'' 0) (fromList [(0,'b')])
'b'

> view (at' 1) (fromList [(0,'b')])
*** Exception: Could not find key in map!

> view (at'' 1) (fromList [(0,'b')])
*** Exception: Could not find key in map!

> set (at' 0) 'c' (fromList [(0,'b')])
fromList [(0,'c')]

> set (at'' 0) (Just 'c') (fromList [(0,'b')])
fromList [(0,'c')]

> set (at' 1) 'c' (fromList [(0,'b')])
fromList [(0,'b'),(1,'c')]

> set (at'' 1) (Just 'c') (fromList [(0,'b')])
fromList [(0,'b'),(1,'c')]

> set (at'' 0) Nothing (fromList [(0,'b')])
fromList []