How to use lenses to look up a value in a map, increase it or set it to a default value

Question

While working on a state called AppState I want keep track of the number of, say, instances. These instances have distinct ids of type InstanceId.

Therefore my state look likes this

import           Control.Lens

data AppState = AppState
  { -- ...
  , _instanceCounter :: Map InstanceId Integer
  }

makeLenses ''AppState

The function to keep track of counts should yield 1 when no instance with given id has been counted before and n + 1 otherwise:

import Data.Map as Map
import Data.Map (Map)

countInstances :: InstanceId -> State AppState Integer
countInstances instanceId = do
    instanceCounter %= incOrSetToOne
    fromMaybe (error "This cannot logically happen.")
              <$> use (instanceCounter . at instanceId)
  where
    incOrSetToOne :: Map InstanceId Integer -> Map InstanceId Integer
    incOrSetToOne m = case Map.lookup instanceId m of
      Just c  -> Map.insert instanceId (c + 1) m
      Nothing -> Map.insert instanceId 1 m

While the above code works, there is hopefully a way to improve it. What I don't like:

• I have to evoke the map instanceCounter twice (first for setting, then for getting the value)
• I use fromMaybe where always Just is expected (so I might as well use fromJust)
• I don't use lenses for the lookup and insertion in incOrSetToOne. The reason is that at does not allow to handle the case where lookup yields Nothing but instead fmaps over Maybe.

Suggestions for improvement?

Answer 1

The way to do this using lens is:

 countInstances :: InstanceId -> State AppState Integer
 countInstances instanceId = instanceCounter . at instanceId . non 0 <+= 1

The key here is to use non

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

This allows us to treat missing elements from the instanceCounter Map as 0