PolyTree.hs

{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}

module Data.PolyTree where

import Control.Lens
import Data.Bifoldable
import Data.Bitraversable
import Data.Functor.Classes
import Data.List.NonEmpty
import qualified Data.Tree as Tree
import Data.Void

data Tree f a b =
  Leaf b
  | Node a (f (Tree f a b))

type Tree' f a =
  Tree f a a

type Tree0 a b =
  Tree [] a b

type Tree1 a b =
  Tree NonEmpty a b

instance Eq1 f => Eq2 (Tree f) where
  liftEq2 _ g (Leaf b1) (Leaf b2) =
    g b1 b2
  liftEq2 _ _ (Leaf _) (Node _ _) =
    False
  liftEq2 _ _ (Node _ _) (Leaf _) =
    False
  liftEq2 f g (Node a1 t1) (Node a2 t2) =
    f a1 a2 && liftEq (liftEq2 f g) t1 t2

instance Ord1 f => Ord2 (Tree f) where
  liftCompare2 _ g (Leaf b1) (Leaf b2) =
    g b1 b2
  liftCompare2 _ _ (Leaf _) (Node _ _) =
    GT
  liftCompare2 _ _ (Node _ _) (Leaf _) =
    LT
  liftCompare2 f g (Node a1 t1) (Node a2 t2) =
    f a1 a2 <> liftCompare (liftCompare2 f g) t1 t2

instance Show1 f => Show2 (Tree f) where
  liftShowsPrec2 _ _ spB _ d (Leaf b) =
    showsUnaryWith spB "Leaf" d b
  liftShowsPrec2 spA slA spB slB d (Node a ts) =
    showsBinaryWith spA (liftShowsPrec (liftShowsPrec2 spA slA spB slB) (liftShowList2 spA slA spB slB)) "Node" d a ts

instance (Eq a, Eq1 f) => Eq1 (Tree f a) where
  liftEq =
    liftEq2 (==)

instance (Ord a, Ord1 f) => Ord1 (Tree f a) where
  liftCompare =
    liftCompare2 compare

instance (Show a, Show1 f) => Show1 (Tree f a) where
  liftShowsPrec =
    liftShowsPrec2 showsPrec showList

instance (Eq a, Eq1 f, Eq b) => Eq (Tree f a b) where
  (==) =
    liftEq (==)

instance (Ord a, Ord1 f, Ord b) => Ord (Tree f a b) where
  compare =
    liftCompare compare

instance (Show a, Show1 f, Show b) => Show (Tree f a b) where
  show x =
    liftShowsPrec showsPrec shows 0 x ""

instance Functor f => Bifunctor (Tree f) where
  bimap _ g (Leaf b) =
    Leaf (g b)
  bimap f g (Node a t) =
    Node (f a) (fmap (bimap f g) t)

instance Functor f => Functor (Tree f a) where
  fmap =
    bimap id

instance Foldable f => Bifoldable (Tree f) where
  bifoldMap _ g (Leaf b) =
    g b
  bifoldMap f g (Node a t) =
    f a <> foldMap (bifoldMap f g) t

instance Foldable f => Foldable (Tree f a) where
  foldMap =
    bifoldMap (pure mempty)

instance Traversable f => Bitraversable (Tree f) where
  bitraverse _ g (Leaf b) =
    Leaf <$> g b
  bitraverse f g (Node a ts) =
    Node <$> f a <*> traverse (bitraverse f g) ts

instance Traversable f => Traversable (Tree f a) where
  traverse =
    bitraverse pure

instance Traversable f => Plated (Tree f a b) where
  plate _ (Leaf b) =
    pure (Leaf b)
  plate f (Node a ts) =
    Node a <$> traverse (plate f) ts

foldTree ::
  (b -> x)
  -> (a -> f (Tree f a b) -> x)
  -> Tree f a b -> x
foldTree l _ (Leaf b) =
  l b
foldTree _ n (Node a t) =
  n a t
  
treeValue ::
  Tree f a b
  -> Either a b
treeValue =
  foldTree Right (pure . Left)

treeChildren ::
  Tree f a b
  -> Either b (f (Tree f a b))
treeChildren =
  foldTree Left (pure Right)

-- | Depth-first search
--
-- >>> dfs (Leaf 1) :: [Either String Int]
-- [Right 1]
--
-- >>> dfs (Node "A" []) :: [Either String Int]
-- [Left "A"]
--
-- >>> dfs (Node "A" [Leaf 1]) :: [Either String Int]
-- [Left "A",Right 1]
--
-- >>> dfs (Node "a" [Node "b" [Leaf 1], Leaf 88, Node "c" [Leaf 2], Leaf 99]) :: [Either String Int]
-- [Left "a",Left "b",Right 1,Right 88,Left "c",Right 2,Right 99]
dfs ::
  (Semigroup (f (Either a b)), Monad f) =>
  Tree f a b
  -> f (Either a b)
dfs (Leaf b) =
  pure (Right b)
dfs (Node a ts) =
  pure (Left a) <> (ts >>= dfs)

-- | Breadth-first search
--
-- >>> bfs (Leaf 1) :: [Either String Int]
-- [Right 1]
--
-- >>> bfs (Node "A" []) :: [Either String Int]
-- [Left "A"]
--
-- >>> bfs (Node "A" [Leaf 1]) :: [Either String Int]
-- [Left "A",Right 1]
--
-- >>> bfs (Node "a" [Node "b" [Leaf 1], Leaf 88, Node "c" [Leaf 2], Leaf 99]) :: [Either String Int]
-- [Left "a",Left "b",Right 88,Left "c",Right 99,Right 1,Right 2]
bfs ::
  (Monoid (f (Either a b)), Monad f, Foldable f) =>
  Tree f a b
  -> f (Either a b)
bfs (Leaf b) =
  pure (Right b)
bfs (Node a ts) =
  let go' xs =
        foldMap (pure . treeValue) xs <>
        foldMap (maybe mempty go' . preview _Right . treeChildren) xs
  in  pure (Left a) <> go' ts

_Leaf ::
  Prism'
    (Tree f a b)
    b
_Leaf =
  prism'
    Leaf
    (\case
        Leaf b ->
          Just b
        _ ->
          Nothing)

_Node ::
  Prism
    (Tree f a b)
    (Tree f' a' b)
    (a, f (Tree f a b))
    (a', f' (Tree f' a' b))
_Node =
  prism
    (uncurry Node)
    (\case
        Node a t ->
          Right (a, t)
        Leaf b ->
          Left (Leaf b))

treeIso ::
  Iso
    (Tree' f a)
    (Tree' f' a')
    (a, Maybe (f (Tree' f a)))
    (a', Maybe (f' (Tree' f' a')))
treeIso =
  iso
    (foldTree (, Nothing) (\a t -> (a, Just t)))
    (\(a, t) -> maybe (Leaf a) (Node a) t)

treeValue' ::
  Lens'
    (Tree' f a)
    a
treeValue' =
  treeIso . _1

treeChildren' ::
  Lens
    (Tree' f a)
    (Tree' f' a)
    (Maybe (f (Tree' f a)))
    (Maybe (f' (Tree' f' a)))
treeChildren' =
  treeIso . _2

baseTree ::
  Iso
    (Tree0 a Void)
    (Tree0 a' Void)
    (Tree.Tree a)
    (Tree.Tree a')
baseTree =
  let mkTree h c =
        let go (Node a t) =
              Tree.Node a (fmap go t)
            go (Leaf b) =
              absurd b
        in  Tree.Node h (fmap go c)
      mkTree0 (Tree.Node h t) = Node h (fmap mkTree0 t)
  in  iso
        (foldTree absurd mkTree)
        (\(Tree.Node h t) -> Node h (fmap mkTree0 t))