module Chiasma.Ui.Data.Tree where

import Control.Lens (makeClassy)
import Prettyprinter (Pretty (..), nest, vsep)

data Tree f l p =
  Tree {
    forall (f :: * -> *) l p. Tree f l p -> l
_treeData :: l,
    forall (f :: * -> *) l p. Tree f l p -> f (Node f l p)
_forest :: f (Node f l p)
    }

deriving stock instance (Eq l, Eq (Node [] l p)) => Eq (Tree [] l p)
deriving stock instance (Show l, Show (Node [] l p)) => Show (Tree [] l p)

deriving stock instance (Eq l, Eq (Node NonEmpty l p)) => Eq (Tree NonEmpty l p)
deriving stock instance (Show l, Show (Node NonEmpty l p)) => Show (Tree NonEmpty l p)

data Node f l p =
  Sub { forall (f :: * -> *) l p. Node f l p -> Tree f l p
_subTree :: Tree f l p }
  |
  Leaf { forall (f :: * -> *) l p. Node f l p -> p
_leafData :: p }

makeClassy ''Tree
makeClassy ''Node

deriving stock instance (Eq l, Eq p) => Eq (Node [] l p)
deriving stock instance (Show l, Show p) => Show (Node [] l p)

deriving stock instance (Eq l, Eq p) => Eq (Node NonEmpty l p)
deriving stock instance (Show l, Show p) => Show (Node NonEmpty l p)

type LTree l p = Tree [] l p
type LNode l p = Node [] l p

type NTree l p = Tree NonEmpty l p
type NNode l p = Node NonEmpty l p

instance (Foldable f, Pretty l, Pretty p) => Pretty (Tree f l p) where
  pretty :: forall ann. Tree f l p -> Doc ann
pretty (Tree l
l f (Node f l p)
sub) =
    Int -> Doc ann -> Doc ann
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 (Doc ann -> Doc ann) -> Doc ann -> Doc ann
forall a b. (a -> b) -> a -> b
$ [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
vsep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
forall a b. (a -> b) -> a -> b
$ l -> Doc ann
forall ann. l -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty l
l Doc ann -> [Doc ann] -> [Doc ann]
forall a. a -> [a] -> [a]
: (Node f l p -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Node f l p -> Doc ann
pretty (Node f l p -> Doc ann) -> [Node f l p] -> [Doc ann]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Node f l p) -> [Node f l p]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f (Node f l p)
sub)

instance (Foldable f, Pretty l, Pretty p) => Pretty (Node f l p) where
  pretty :: forall ann. Node f l p -> Doc ann
pretty (Sub Tree f l p
tree') = Tree f l p -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. Tree f l p -> Doc ann
pretty Tree f l p
tree'
  pretty (Leaf p
a) = p -> Doc ann
forall ann. p -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty p
a