Traversal to combine multiple map operations into single ADT

Given a record consisting of multiple maps, how can I write a traversal (or prism, or Lens' TestLens (Maybe Interim)) that allows me to group together lookups?

First off, my current attempts.


data TestLens = TL
    { _foo1 :: Map.Map Text Int
    , _foo2 :: Map.Map Text Bool
    , _foo3 :: Map.Map Text Text
    } deriving Show
 
tl = TL (Map.fromList [("a", 5), ("b", 6), ("c", 1), ("d", 3)])
         (Map.fromList [("b", True), ("c", False), ("d", True)])
         (Map.fromList [("c", "foo"), ("d", "bar")])
 
makeLenses ''TestLens
 
data Interim = Interim Int Bool Text deriving Show
data Interim2 = Interim2 Int Bool deriving Show

getOnePart s l k = s ^. l . at k
 
interim s k = Interim <$> getOnePart s foo1 k <*> getOnePart s foo2 k <*> getOnePart s foo3 k
interim2 s k = Interim2 <$> getOnePart s foo1 k <*> getOnePart s foo2 k
doTestStuff = tl ^.. folding (\s -> mapMaybe (interim s) (Map.keys $ s ^. foo1))

The intended behaviour is that interim (as it stands, it's a mishmash of lens and..not lens) combines at over multiple Maps:

interim tl "a" = Nothing
interim tl "c" = Just (Interim 1 False "foo")

and then I can fold over all possible keys to get the complete list of Interims.

What I'd like to be able to do is build an indexed traversal (rather than an unindexed fold) over all possible Interims, but so far I've had no luck in the combo of itraversed I need here..I suspect because I flip between map and lens:

itraverseInterim2s = ...

> tl ^@.. itraverseInterim2s
[("b", Interim2 6 True), ("c", Interim2 1 False), ("d", Interim2 3 True)]
-- and if we assume there exists _1 :: Lens' Interim2 Int
> tl & itraverseInterim2s . _1 %~ (+5)
TL (Map.fromList [("a", 5), ("b", 11), ("c", 6), ("d", 8)])
         (Map.fromList [("b", True), ("c", False), ("d", True)])
         (Map.fromList [("c", "foo"), ("d", "bar")])

I can't equally work out if last behaviour is better solved by making a Lens' TestLens (Maybe Interim2), a k -> Prism' TestLens Interim2 (I think only one of these satisfies lens laws), or by having individual elements traversed with itraverseInterim2s . index k.

Obviously for every InterimX ADT I want to be able to extract from the combination of fooX maps I'll have to write minor boilerplate but that bit's fine.

Solution

Have you considered writing something like:

fanoutTraversal :: Traversal' s a -> Traversal' s b -> Traversal' s (a,b)
fanoutTraversal t1 t2 fab s =
  maybe (pure s) (fmap update . fab) mv
  where
    mv = liftA2 (,) (s ^? t1) (s ^? t2)
    update (c,d) = s & t1 .~ c & t2 .~ d

With this function, you can write interim as:

interim :: Text -> Traversal' TestLens Interim
interim k = (((foo1 . ix k) `fanoutTraversal` (foo2 . ix k)) `fanoutTraversal` (foo3 . ix k)) . interimIso
  where
    interimIso = iso (\((a,b),c) -> Interim a b c) (\(Interim a b c) -> ((a,b),c))

Things would need to change a little if you want to use at instead of ix or to use IndexedTraversal instead of Traversal, but the idea is hopefully sound.

If your goal is to traverse through all the Interims in the TestLens, it may be easier to first convert TestLens to Map.Map Text Interim and then traverse that map:

import Control.Lens hiding ((<.>))
import Data.Functor.Apply (Apply(..)) -- could just as well use Map.intersectionWith

manyInterim :: Traversal' TestLens Interim
manyInterim = manyInterim' . traverse

-- Let's use this version of Interim so that we have record access
data Interim = Interim
  { i1 :: Int
  , i2 :: Bool
  , i3 :: Text
  } deriving Show

manyInterim' :: Lens' TestLens (Map.Map Text Interim)
manyInterim' = lens sa sbt
  where
    sa TL{..} = Interim <$> _foo1 <.> _foo2 <.> _foo3
    sbt TL{..} interimMap = TL
      { _foo1 = Map.union (i1 <$> interimMap) _foo1
      , _foo2 = Map.union (i2 <$> interimMap) _foo2
      , _foo3 = Map.union (i3 <$> interimMap) _foo3
      }