How to use Lenses to traverse and assign to some (but not all) elements in a Map

Question

I've been trying to use lenses and containers together with some success but i have hit the limitations of my understanding trying to use filtering traversals of Data.Map - i can change individual instances in the map or traverse all instances but i can't work out how to act on some identifiable partition (ie keys in a range).

Essentially i'm trying to do something similar with maps to what Gabriel Gonzalez excellent lenses tutorial does with lists [1]

Here's a working skeleton of my code with the traverseSome function which i don't know how to write commented out. Any help gratefully received!

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveGeneric   #-}
{-# LANGUAGE RankNTypes      #-}

import Control.Lens
import Control.Monad.State
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

type CharSet = Set.Set Char
type MapOfSets = Map.Map Int CharSet

data DB = DB { _mos  :: MapOfSets } deriving (Show, Eq)

makeLenses ''DB

initDB :: DB
initDB =  DB { _mos  = Map.fromList (zip [1..5] (repeat Set.empty)) }

add2Map :: Int -> CharSet -> State DB ()
add2Map i cs = mos.ix i %= (Set.union cs)

traverseAll :: Traversal' DB CharSet
traverseAll = mos.traversed

add2MapsAll :: CharSet -> State DB ()
add2MapsAll cs = traverseAll %= (Set.union cs)

--        <problematic part>          
{-
traverseSome :: [Int] -> Int -> Traversal' DB MapOfSets
traverseSome ids i = _

add2MapsSome :: [Int] -> CharSet -> State DB ()
add2MapsSome ids cs = mos.(traverseSome ids 2) %= (Set.union cs)
-}         
--        </problematic part>

main :: IO ()
main = do
  let db = initDB
  let bar = Set.fromList ['a'..'g'] :: CharSet
  let baz = Set.fromList ['f'..'m'] :: CharSet
  let quux = Set.fromList ['n'..'z'] :: CharSet

  let db2 = execState (add2Map 5 bar) db
  let db3 = execState (add2MapsAll baz) db
  -- let db4 = execState (add2MapsSome [1,3] quux) db

  print db2
  print db3
  -- print db4

[1] http://www.haskellforall.com/2013/05/program-imperatively-using-haskell.html

Answer 1

I'm assuming you mean

traverseSome :: [Int] -> Traversal' DB CharSet

Here's a more general version

keys :: Ord k => [k] -> IndexedTraversal' k (Map.Map k a) a
keys ks f m = go ks <&> \m' -> foldr (uncurry M.insert) m m'
  where
    go []     = pure []
    go (i:is) = case Map.lookup i m of
                  Just a  -> (:) . (,) i <$> indexed f i a <*> go is
                  Nothing -> go is

which is very similar to ordinals from Data.Vector.Lens (my version doesn't nub duplicates, so make sure the list doesn't have duplicates). go goes through the list of indices and looks them up in the map, adding the index as it goes. The foldr bit goes through the list of edited elements and inserts them back into the original map.

You can write your's as

traverseSome :: [Int] -> IndexedTraversal' Int DB CharSet
traverseSome is = mos . keys is

add2MapsSome :: [Int] -> CharSet -> State DB ()
add2MapsSome is cs = traverseSome is %= Set.union cs

---

If you did want

traverseSome :: [Int] -> Lens' DB MapOfSets

this can be written as (note you shouldn't add new keys to the Map or you'll break the lens laws)

submap :: Ord k => [k] -> Lens' (Map.Map k a) (Map.Map k a)
submap ks f m = f (Map.fromList as) <&> (<> m)
  where as = Maybe.mapMaybe (\i -> (,) i <$> Map.lookup i m) ks

which could be used to write keys (but would be less efficient because you make an intermediate Map):

keys :: Ord k => [k] -> IndexedTraversal' k (Map k a) a
keys ks = submap ks . itraversed

---

edit: version without intermediate lists:

keys :: Ord k => [k] -> IndexedTraversal' k (Map.Map k a) a
keys ks f m = go ks
  where
    go []     = pure m
    go (i:is) =
      case Map.lookup i m of
        Just a  -> Map.insert i <$> indexed f i a <*> go is
        Nothing -> go is