3
\$\begingroup\$

Considering the lack of general purpose library to manipulate trees (not to mention zipper ones), I may end up packaging this piece of code, so I would like it to be reviewed first.

Any comment is welcome, notably on the following points:

  • the use of Map within the Tree data type (e.g. versus HashMap or bare lists)
  • the semantic of the Foldable instance of ZTree
  • the implementation of the Comonad instance of ZTree
  • types/functions naming
  • opportunities of optimization/generalization/code reuse

EDIT: re-using Cofree to define Tree.

{-# LANGUAGE DeriveFunctor      #-}
{-# LANGUAGE StandaloneDeriving #-}
import           Control.Comonad -- from comonad
import           Control.Comonad.Cofree -- from free

import           Data.Containers -- from mono-traversable
import           Data.Map               (Map) -- from containers
import           Data.Maybe

import           Lens.Micro -- from microlens

import           Prelude                hiding (lookup)

-- * Generic tree

type Tree e = Cofree (Map e)

rootEdges :: (Ord e) => Tree e a -> [e]
rootEdges (_ :< c) = keys c

removeEdge :: (Ord e) => e -> Tree e a -> Tree e a
removeEdge e (a :< c) = a :< deleteMap e c

attach :: (Ord e) => e -> Tree e a -> Tree e a -> Tree e a
attach e t (a :< c) = a :< insertMap e t c

prune :: (Ord e) => Int -> Tree e a -> Tree e a
prune 0 (a :< _) = a :< mempty
prune n (a :< c) = a :< fmap (prune (n-1)) c

-- * Zipper tree

-- | A zipper tree is a 'Tree' with a cursor
data ZTree e a = ZRoot (Tree e a)
               | ZNode (Tree e a) e (ZTree e a)
deriving instance Functor (ZTree e)

-- | Only fold the 'descendants' of the currently focused node
instance Foldable (ZTree e) where
  foldMap f ztree = foldMap f $ ztree ^. underlyingTree

instance Traversable (ZTree e) where
  traverse f (ZRoot tree) = ZRoot <$> traverse f tree
  traverse f (ZNode tree e parent) = ZNode <$> traverse f tree <*> pure e <*> traverse f parent

instance (Ord e) => Comonad (ZTree e) where
  extract ztree = extract $ ztree ^. underlyingTree
  duplicate w@(ZRoot _) = ZRoot $ duplicate' w
  duplicate w@(ZNode _ e parent) = ZNode (duplicate' w) e (duplicate parent)

duplicate' :: (Ord e) => ZTree e a -> Tree e (ZTree e a)
duplicate' w = w :< mapFromList (map (\e -> (e, duplicate' $ forward e w)) $ nextEdges w)

instance (Eq a, Ord e) => Eq (ZTree e a) where
  a == b = extract a == extract b

instance (Ord a, Ord e) => Ord (ZTree e a) where
  compare a b = extract a `compare` extract b

-- | ZTree constructor
zTree :: Tree e a -> ZTree e a
zTree = ZRoot

underlyingTree :: Lens' (ZTree e a) (Tree e a)
underlyingTree f (ZRoot tree) = ZRoot <$> f tree
underlyingTree f (ZNode tree e parent) = (\tree' -> ZNode tree' e parent) <$> f tree

nextEdges :: (Ord e) => ZTree e a -> [e]
nextEdges = rootEdges . (^. underlyingTree)

removeEdgeZ :: (Ord e) => e -> ZTree e a -> ZTree e a
removeEdgeZ e = over underlyingTree (removeEdge e)

attachZ :: (Ord e) => e -> Tree e a -> ZTree e a -> ZTree e a
attachZ e t = over underlyingTree $ attach e t

pruneZ :: (Ord e) => Int -> ZTree e a -> ZTree e a
pruneZ n = over underlyingTree $ prune n

forwardMay :: (Ord e) => e -> ZTree e a -> Maybe (ZTree e a)
forwardMay e w = do
  child <- lookup e children
  return $ ZNode child e $ removeEdgeZ e w
  where _ :< children = w ^. underlyingTree

forward :: (Ord e) => e -> ZTree e a -> ZTree e a
forward e ztree = fromMaybe ztree $ forwardMay e ztree

backwardMay :: (Ord e) => ZTree e a -> Maybe (ZTree e a)
backwardMay (ZRoot _) = Nothing
backwardMay (ZNode tree e parent) = Just $ attachZ e tree parent

backward :: (Ord e) => ZTree e a -> ZTree e a
backward ztree = fromMaybe ztree $ backwardMay ztree
\$\endgroup\$
  • 1
    \$\begingroup\$ If you want to package any code as a library, you should always include proper documentation (and tests). \$\endgroup\$ – Zeta May 21 '16 at 14:38
  • \$\begingroup\$ @Zeta It goes without saying that if this piece of code ever makes it to hackage, it won't be as is. But thanks for the reminder :) . \$\endgroup\$ – koral May 21 '16 at 15:01
  • \$\begingroup\$ Fair enough. I didn't dive into your code yet, but I'm somewhat surprised by your statement of "lack of general purpose library". Isn't zippers what you want? Or do you mean Tree as in a graph? \$\endgroup\$ – Zeta May 21 '16 at 15:39
  • 1
    \$\begingroup\$ type Tree e = Cofree (Map e) \$\endgroup\$ – Gurkenglas May 21 '16 at 19:45
  • \$\begingroup\$ @Gurkenglas I've just updated the code to re-use Cofree, this is really elegant, thank you. \$\endgroup\$ – koral May 22 '16 at 10:11

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.