Implementing polymorphic 'deep' function for traversals and folds

I'm using lens together with xml-lens. I'd like to make the following function more polymorphic, so that it also works for Folds and not only Traversals:

-- | Traverse a plated structure recursively, trying to match a fold at each level. Don't recurse
-- if the fold matches.
deep :: forall s a. Plated s => Traversal' s a -> Traversal' s a
deep f = let go :: Traversal' s a; go = cloneTraversal $ failing f (plate . go) in go

This function works like the deep function from hxt. Is it possible to make it more polymorphic in the way I want?

Solution

This one is rather tricky given the current publicly exposed API.

I took the liberty of expanding the type of deepOf to also support indexed folds and indexed traversals along the way as it was easier than not doing so, and makes the implementation be the full one we'd want to export from lens, anyways.

Let's import the parts of lens that we don't normally show users.

{-# LANGUAGE RankNTypes #-}

import Control.Applicative
import Control.Lens
import Control.Lens.Internal.Bazaar
import Control.Lens.Internal.Context
import Control.Lens.Traversal
import Control.Monad.State
import Data.Profunctor.Rep
import Data.Profunctor.Unsafe

We'll need a few internal combinators that we don't expose from Control.Lens.Traversal that are used to manipulate a Traversal/Fold as a BazaarT and collapse the answer back out.

pins :: (Bizarre p w, Corepresentable p) => w a b t -> [Corep p a]
pins = getConst #. bazaar (cotabulate $ \ra -> Const [ra])
{-# INLINE pins #-}

unsafeOuts :: (Bizarre p w, Corepresentable p) => w a b t -> [b] -> t
unsafeOuts = evalState `rmap` bazaar (cotabulate (\_ -> state (unconsWithDefault fakeVal)))
  where fakeVal = error "unsafePartsOf': not enough elements were supplied"
{-# INLINE unsafeOuts #-}

unconsWithDefault :: a -> [a] -> (a,[a])
unconsWithDefault d []     = (d,[])
unconsWithDefault _ (x:xs) = (x,xs)
{-# INLINE unconsWithDefault #-}

Now that we have that, we build a proper version of deep.

-- |
-- @
-- 'deep' :: 'Plated' s => 'Fold' s a                 -> 'Fold' s a
-- 'deep' :: 'Plated' s => 'Traversal' s s a b        -> 'Traversal' s s a b
-- 'deep' :: 'Plated' s => 'IndexedFold' i s a        -> 'IndexedFold' i s a
-- 'deep' :: 'Plated' s => 'IndexedTraversal' s s a b -> 'Traversal' i s s a b
-- @
deep :: (Plated s, Conjoined p, Applicative f) => Traversing p f s s a b -> Over p f s s a b
deep = deepOf plate

-- |
-- @
-- 'deepOf' :: 'Fold s s'       -> 'Fold' s a                 -> 'Fold' s a
-- 'deepOf' :: 'Traversal' s s' -> 'Traversal' s s a b        -> 'Traversal' s s a b
-- 'deepOf' :: 'Fold s s'       -> 'IndexedFold' i s a        -> 'IndexedFold' i s a
-- 'deepOf' :: 'Traversal' s s' -> 'IndexedTraversal' s s a b -> 'Traversal' i s s a b
-- @
deepOf :: (Plated s, Conjoined p, Applicative f) => LensLike' f s s -> Traversing p f s s a b -> Over p f s s a b
deepOf r l pafb s = case pins b of
  [] -> r (deep l pafb) s
  xs -> unsafeOuts b <$> traverse (corep pafb) xs
  where b = l sell s

The guts of deepOf are very similar to the existing guts of failing, which you were rightly trying to use as your workhorse.

failing :: (Conjoined p, Applicative f) => Traversing p f s t a b -> Traversing p f s t a b -> Over p f s t a b
failing l r pafb s = case pins b of
  [] -> runBazaarT (r sell s) pafb
  xs -> unsafeOuts b <$> traverse (corep pafb) xs
  where b = l sell s

The only thing different is the [] case, where instead of falling over, what we do is run the whole nested Traversal.

I've merely typechecked this and not actually executed it, but it looks right to me.

Feel free to put in an issue on http://github.com/ekmett/lens/issues to add these combinators (or some bikeshedded renaming of them), they probably belong in the core lens API, lest code like this fall on users, while it is trivial to implement within the library itself.

This is the kind of code we try to write once, so that end users do not have to.