Attempting lens/traversal map multi-update on a subset of its keys

I'm trying to conjure a traversal to update multiple keys of an IntMap as a whole.

To dispell XY: I'm not simply trying to update them, I need the traversal to return to the caller for further composition. Or at least something composable with lenses.

I've tried many variations of the common combinators. I've tried dropping down to a functor-based definition, with a large range of experimentation shifting the foralls' scopes around, with no more success. Building from the ground up again, here's where I'm at:

import Control.Lens
import Control.Lens.Unsound

-- base case: traverse a single fixed element
t1 :: Traversal' (IntMap a) (Maybe a)
t1 = at 0

-- build-up case: traverse a pair of fixed elements
t2 :: Traversal' (IntMap a) (Maybe a)
t2 = at 0 `adjoin` at 1

-- generalizing case: do it with a fold
t3 :: Traversal' (IntMap a) (Maybe a)
t3 = foldr (\e t -> at e `adjoin` t) (at 1) [0]

t1 and t2 work fine; I'd devised t3 to be equivalent to t2, but it fails with the following error:

• Couldn't match type ‘f1’ with ‘f’
  ‘f1’ is a rigid type variable bound by a type expected by the context:
    Traversal' (IntMap a) (Maybe a)
  ‘f’ is a rigid type variable bound by the type signature for:
    t3 :: forall a. Traversal' (IntMap a) (Maybe a)
  Expected type: (Maybe a -> f1 (Maybe a)) -> IntMap a -> f1 (IntMap a)
  Actual type: (Maybe a -> f (Maybe a)) -> IntMap a -> f (IntMap a)
• In the second argument of ‘adjoin’, namely ‘t’   
  In the expression: at x `adjoin` t
  In the first argument of ‘foldr’, namely ‘(\ x t -> at x `adjoin` t)’

I suppose this is some rank-2 trickery that's still a bit over my head. Is there any way to make this work?

I aimed for a final signature of

ats :: Foldable l => l Int -> Traversal' (IntMap a) (Maybe a)

…assuming unique keys, of course. Which I dreamed could be implemented just almost like t3.

Solution

Traversal' is a type synonym for a type containing forall, which makes it a second class citizen in the type system: we can't instantiate a type variable with such a type.

In particular, here we are trying to do so with foldr :: (a -> b -> b) -> b -> [a] -> b, we can't instantiate b = Traversal' _ _, because Traversal' contains a forall.

One work around is to wrap Traversal' in a newtype, ReifiedTraversal. Wrap (using the Traversal constructor) before passing at 1 to foldr; inside foldr, unwrap to use adjoin, and rewrap; unwrap at the end.

t3 :: Traversal' (IntMap a) (Maybe a)
t3 = runTraversal (foldr (\e t -> Traversal (at e `adjoin` runTraversal t)) (Traversal (at 1)) [0])

A traversal is a function Applicative f => (t -> f t) -> (s -> f s). You have a function f :: Maybe a -> f (Maybe a) and you want to apply it to some entries in IntMap a.

It's a bit of a puzzle to do with Applicative (there is a more natural solution using Monad), but requires less expertise than composing traversals as first-class values:

import Control.Applicative
import Data.IntMap (IntMap)
import qualified Data.IntMap as M

-- [Int] -> Traversal' (IntMap a) (Maybe a)
traverseAtKeys :: Applicative f => [Int] -> (Maybe a -> f (Maybe a)) -> IntMap a -> f (IntMap a)
traverseAtKeys keys f m =
  let go i k = liftA2 (insertMaybe i) (f (M.lookup i m)) k
      insertMaybe i Nothing = M.delete i
      insertMaybe i (Just v) = M.insert i v
  in foldr go (pure m) keys