Diverse.hs

{-# LANGUAGE RankNTypes, StandaloneDeriving, TemplateHaskell, UndecidableInstances #-}

module Data.Tree.Diverse
    ( Node(..), _Node
    , Children(..), overChildren
    , Ann(..), ann, val
    ) where

import qualified Control.Lens as Lens
import Data.Functor.Const (Const(..))
import Data.Functor.Identity (Identity(..))
import GHC.Generics (Generic)

import Prelude

newtype Node f expr = Node (f (expr f))
    deriving Generic
deriving instance Eq (f (expr f)) => Eq (Node f expr)
deriving instance Ord (f (expr f)) => Ord (Node f expr)
deriving instance Show (f (expr f)) => Show (Node f expr)

Lens.makePrisms ''Node

class Children expr where
    children ::
        Applicative f =>
        (forall sub. Children sub => Node n sub -> f (Node m sub)) ->
        expr n -> f (expr m)

overChildren ::
    Children expr =>
    (forall sub. Children sub => Node n sub -> Node m sub) ->
    expr n -> expr m
overChildren f = runIdentity . children (Identity . f)

instance Children (Const val) where
    children _ (Const x) = pure (Const x)

-- Annotate tree nodes
data Ann a v = Ann
    { _ann :: a
    , _val :: v
    } deriving (Eq, Ord, Show, Functor, Foldable, Traversable, Generic)

Lens.makeLenses ''Ann