Changing indices within a Haskell tree

Question

(Sorry for the long context description, but I couldn't find a simpler way to explain my problem) Consider the following types:

import Data.Array

data UnitDir = Xp | Xm | Yp | Ym | Zp | Zm
    deriving (Show, Eq, Ord, Enum, Bounded, Ix)

type Neighborhood a = Array UnitDir (Tree a)

data Tree a = Empty | Leaf a | Internal a (Neighborhood a)
    deriving (Eq, Show)

Clearly, Tree can be defined as an instance of Functor as follows:

instance Functor Tree where
    fmap _ Empty           = Empty
    fmap f (Leaf x)        = Leaf (f x)
    fmap f (Internal x ts) = Internal (f x) $ fmap (fmap f) ts

I would like to define a function that traverses an instance of Tree by permuting the indices of the Array UnitDir (Tree a) (so it's a permutation on the 6 possible values of UnitDir).

A possible implementation would be this one:

type Permutation = Array UnitDir UnitDir

applyPermutation :: Permutation -> Tree a -> Tree a
applyPermutation _ Empty = Empty
applyPermutation _ (Leaf x) = Leaf x
applyPermutation f (Internal x ts) = Internal x (applyPermutation' ts)
    where applyPermutation' ts = ixmap (Xp, Zm) (f !) (applyPermutation f <$> ts)

My question is the following: Is there a natural Haskell construction to "traverse" the tree while reindexing the children?

Functor does not work, since I use it to change the content of the tree, not its indexing scheme. It seems I would need two instances of Functor, one to change the content and the other to change the array indices.

I thought that Traversable would be the right choice, but none of the signatures of the provided functions matches that of applyPermutation.

Thanks in advance for any help.

Answer 1

Functor does not work, since I use it to change the content of the tree, not its indexing scheme. It seems I would need two instances of Functor, one to change the content and the other to change the array indices.

Your intuition here is spot on: a functor that acted on the Neighborhood a field would do what you need, and it is correct to call such a thing "functor". Here is one possible refactoring of applyPermutation:

{-# LANGUAGE LambdaCase #-}

-- I prefer case syntax for this sort of definition; with it, there is less stuff
-- that needs to be repeated. LambdaCase is the icing on the cake: it frees me
-- me from naming the Tree a argument -- without it I would be forced to write
-- mapOverNeighborhoods f t = case t of {- etc. -}
mapOverNeighborhoods :: (Neighborhood a -> Neighborhood a) -> Tree a -> Tree a
mapOverNeighborhoods f = \case 
    Empty -> Empty
    Leaf x -> Leaf x
    Internal x ts -> Internal x (f (mapOverNeighborhoods f <$> ts))

applyPermutation :: Permutation -> Tree a -> Tree a
applyPermutation perm = mapOverNeighborhoods applyPermutation'
    where applyPermutation' = ixmap (Xp, Zm) (perm !)

(You might prefer to go even further and use a mapping that takes UnitDirection -> UnitDirection directly, rather than Neighborhood a -> Neighborhood a. I didn't do that primarily to make it mirror the rest of this answer more closely, but also because it arguably makes for a more honest interface -- rearranging indices in an Array is not quite as straightforward as applying an arbitrary function to the indices.)

There are two limitations of this attempt to define another functor:

• We already have a Functor instance, as you point out. It wouldn't be sensible to replace just for this use case, and defining a newtype for it would be too annoying.

• Even if that wasn't the case, mapOverNeighborhoods can't be made into a Functor instance, as fmap takes arbitrary a -> b functions, and changing the type of the neighborhoods is not an option.

These two concerns are addressed by optics libraries such as lens (if you end up using optics for just this one thing in your code base, though, you might prefer microlens for a smaller dependency footprint).

{-# LANGUAGE TemplateHaskell #-} -- makeLenses needs this.
{-# LANGUAGE DeriveFunctor #-} -- For the sake of convenience.
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}

-- Record fields on sum types are nasty; these, however, are only here for the
-- sake of automatically generating optics with makeLenses, so it's okay.
data Tree a
    = Empty 
    | Leaf { _value :: a } 
    | Internal { _value :: a, _neighborhood :: Neighborhood a }
    deriving (Eq, Show, Functor, Foldable, Traversable)
makeLenses ''Tree

applyPermutation :: Permutation -> Tree a -> Tree a
applyPermutation perm = over neighborhood applyPermutation'
    where applyPermutation' = ixmap (Xp, Zm) (perm !)

over (infix spelling: %~) is literally an fmap which allows choosing the targets. We do that by passing it an appropriate optic (in this case, neighborhood, which is a Traversal that targets all neighborhoods in a tree -- over neighborhood can be read as "map over all neighborhoods"). Note that the fact that we can't change the type of the neighborhood is not a problem (and also, in other circumstances, it would be possible to have type-changing optics).

On a final note, the type of neighborhoods is Traversal' (Tree a) (Neighborhood a). If we expand the Traversal' type synonym, we get:

GHCi> :t neighborhood
neighborhood
  :: Applicative f =>
     (Neighborhood a -> f (Neighborhood a)) -> Tree a -> f (Tree a)

While going into the reasons why it is like that would make this answer too long, it is worth noting that this is a lot like the signature of traverse for Tree...

GHCi> :set -XTypeApplications
GHCi> :t traverse @Tree
traverse @Tree
  :: Applicative f => (a -> f b) -> Tree a -> f (Tree b)

... except that it acts on the neighborhoods rather than on the values (cf. the parallel between fmap and mapOverNeighborhoods). In fact, if you were to adequately implement the traverse analogue with that type, you would be able to use it instead of the one automatically generated by makeLenses.

Answer 2

For completeness, I coded a small variant based on catamorphisms, exploiting recursion-schemes.

{-# LANGUAGE LambdaCase, DeriveFunctor, KindSignatures, TypeFamilies, 
    DeriveFoldable, DeriveTraversable, TemplateHaskell #-}

import Data.Functor.Foldable
import Data.Functor.Foldable.TH

import Data.Array
data UnitDir = Xp | Xm | Yp | Ym | Zp | Zm
    deriving (Show, Eq, Ord, Enum, Bounded, Ix)

type Neighborhood a = Array UnitDir (Tree a)

data Tree a = Empty | Leaf a | Internal a (Neighborhood a)
    deriving (Eq, Show, Functor)

-- Use TH to automatically define a base functor for Tree,
-- enabling recursion-schemes
makeBaseFunctor ''Tree

The wanted mapping function is then:

mapOverNeighborhoods :: (Neighborhood a -> Neighborhood a) -> Tree a -> Tree a
mapOverNeighborhoods f = cata $ \case
   EmptyF -> Empty
   LeafF x -> Leaf x
   InternalF x nb -> Internal x (f nb)

Roughly, cata does all the recursion for us. It provides to its function argument (\case ..., above) a value of type TreeF a (Tree a) which is essentially the same as a plain Tree a except that the very first "layer" uses different constructors, ending with an extra F. All such constructors have their inner trees already pre-processed by cata: above, we can assume that all the trees inside the nb array already had f applied recursively. What we need to do is to handle the first "layer", transforming the F constructors into the regular ones, and applying f to this first "layer" as well.