haskellfunctortraversablefoldable

Changing indices within a Haskell tree


(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.


Solution

  • 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:

    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.