Lenses, the State monad, and Maps with known keys

Question

here is a puzzle that I keep on bumping into and that, I believe, no previous SO question has been addressing: How can I best use the lens library to set or get values within a State monad managing a nested data structure that involves Maps when I know for a fact that certain keys are present in the maps involved?

Here is the puzzle

{-# LANGUAGE TemplateHaskell, DerivingVia #-}

import Control.Monad.State
import Control.Monad.Except
import Control.Lens
import Data.Maybe
import Control.Monad
import Data.Map

type M = StateT World (ExceptT String Identity)

data World = World
  { _users      :: Map UserId User
  , _otherStuff :: Int
  }

type UserId = Int

data User = User
  { _balance     :: Balance
  , _moreStuff   :: Int
  }

newtype Balance = Balance Int
  deriving (Eq, Ord, Num) via Int

makeLenses 'World
makeLenses 'User

deleteUser :: UserId -> M ()
deleteUser uid = do
   user <- use $ users . at uid
   unless (isJust user) (throwError "unknown user")

   -- from here on we know the users exists.

   -- Question: how should the following lens look like?
   balance <- use $ users . ix uid . balance
   when (balance < 0) (throwError "you first have to settle your debt")
   when (balance > 0) (throwError "you first have to withdraw the remaining balance")

   users . at uid .= Nothing

Attempt 1: using ix

The snippet above is using ix.

   balance <- use $ users . ix uid . balance

This yields a Traversal, so it may focus on multiple elements or none at all. In the context of use this means we need a Monoid and a Semigroup instance. In fact, this is what GHC has to say:

    • No instance for (Monoid Balance) arising from a use of ‘ix’
    • In the first argument of ‘(.)’, namely ‘ix uid’
      In the second argument of ‘(.)’, namely ‘ix uid . balance’
      In the second argument of ‘($)’, namely ‘users . ix uid . balance’
   |
45 |    balance <- use $ users . ix uid . balance

There is no good way to implement <> for Balance. I could just implement addition, or use error, because, in fact, this function will never be called. But is this the cleanest way to do this?

Attempt 2: using at

Another option seems to be using at.

   balance <- use $ users . at uid . balance

This yields a Lens that focuses on a Maybe User. This means, the follow-up lens balance has the wrong type.

    • Couldn't match type ‘User’
                     with ‘Maybe (IxValue (Map UserId User))’
      Expected type: (User -> Const Balance User)
                     -> Map UserId User -> Const Balance (Map UserId User)
        Actual type: (Maybe (IxValue (Map UserId User))
                      -> Const Balance (Maybe (IxValue (Map UserId User))))
                     -> Map UserId User -> Const Balance (Map UserId User)
    • In the first argument of ‘(.)’, namely ‘at uid’
      In the second argument of ‘(.)’, namely ‘at uid . balance’
      In the second argument of ‘($)’, namely ‘users . at uid . balance’

Attempt 3: working with at's Maybe

Let's try to work with that Maybe

   balance <- use $ users . at uid . _Just . balance

This time, we have a Prism, which needs to deal with the situation when it has to work with Nothing. So we are back at requiring a Monoid.

    • No instance for (Monoid Balance) arising from a use of ‘_Just’
    • In the first argument of ‘(.)’, namely ‘_Just’
      In the second argument of ‘(.)’, namely ‘_Just . balance’
      In the second argument of ‘(.)’, namely ‘at uid . _Just . balance’

Attempt 3b: working with at's Maybe

Let's try another way to work with that Maybe

   balance <- use $ users . at uid . non undefined . balance

From the documentation:

If v is an element of a type a, and a' is a sans the element v, then non v is an isomorphism from Maybe a' to a.

We can use undefined, error or any "empty" User we want, it does not matter, since this case is never triggered, given the user id is present in the map.

For this to work we need Eq for User, which is fair enough. And it compiles and seems to work; that is, for reading. For writing, there is an, initially, unexpected twist:

topUp :: UserId -> Balance -> M ()
topUp uid b = do
   user <- use $ users . at uid
   unless (isJust user) (throwError "unknown user")

   users . at uid . non undefined . balance += b

Running this blows up

experiment-exe: Prelude.undefined
CallStack (from HasCallStack):
  error, called at libraries/base/GHC/Err.hs:79:14 in base:GHC.Err
  undefined, called at app/Main.hs:55:25 in main:Main

The explanation is that when writing, we use the optics "right-to-left", and in this direction non is injecting the User we provide as its parameter. Replacing undefined with a "empty" User obscures this mistake. It always replaces the existing user with the empty user, effectively loosing the user's initial balance when trying to top-up.

Conclusion

So, I have found options to make this work for reading, but none seems to be convincing. And I could not figure this out for writing.

What is your recommendation? How should that lens be constructed?

Edit: moved solution to self answer.

Answer 1

If you are sure that the key is present then you can use fromJust to turn the Maybe User into a User:

balance <- use $ users . at uid . to fromJust . balance

Although as a design issue I'd suggest replacing use $ users . at uid ... with functions that throw a meaningful error:

getUser :: UserId -> M User

And for handy lens accessing:

getsUser :: UserId -> Getter User a -> M a

Then just call one of those every time you want to look up a user. That way you don't have to have a separate check at the head of your function.

Answer 2

The most direct facilities lens has for doing this are unsafe operations that treat a traversal that you "know" will target only one element as if it were a lens. As you're probably aware, there's an operator for this that serves as a variation of ^. or ^?:

s <- get
let user = s ^?! users . at uid

but for view (within a Reader) or use (within a State), there don't seem to be any built-in variations. You can write your own using the unsafeSingular function in Control.Lens.Traversal though:

use1 :: MonadState s m => Traversal' s a -> m a
use1 = use . unsafeSingular

view1 :: MonadReader s m => Traversal' s a -> m a
view1 = view . unsafeSingular

after which:

balance <- use1 $ users . ix uid . balance

should work.

If you would rather make the optic itself unsafe rather using a safe optic unsafely, you can use unsafeSingular directly to modify the optic. For example:

balance <- use $ users . unsafeSingular (ix uid) . balance