module Heist.Extra.Splices.Tree (treeSplice) where

import Data.Map.Syntax ((##))
import Data.Tree (Tree (..))
import Heist qualified as H
import Heist.Interpreted qualified as HI
import Heist.Splices qualified as Heist

treeSplice ::
  forall a sortKey.
  (Ord sortKey) =>
  -- | How to sort children
  (NonEmpty a -> [Tree a] -> sortKey) ->
  -- | Input tree
  [Tree a] ->
  -- | How to render a (sub-)tree root
  (NonEmpty a -> [Tree a] -> H.Splices (HI.Splice Identity)) ->
  HI.Splice Identity
treeSplice :: forall a sortKey.
Ord sortKey =>
(NonEmpty a -> [Tree a] -> sortKey)
-> [Tree a]
-> (NonEmpty a -> [Tree a] -> Splices (Splice Identity))
-> Splice Identity
treeSplice =
  [a]
-> (NonEmpty a -> [Tree a] -> sortKey)
-> [Tree a]
-> (NonEmpty a -> [Tree a] -> Splices (Splice Identity))
-> Splice Identity
go []
  where
    go :: [a] -> (NonEmpty a -> [Tree a] -> sortKey) -> [Tree a] -> (NonEmpty a -> [Tree a] -> H.Splices (HI.Splice Identity)) -> HI.Splice Identity
    go :: [a]
-> (NonEmpty a -> [Tree a] -> sortKey)
-> [Tree a]
-> (NonEmpty a -> [Tree a] -> Splices (Splice Identity))
-> Splice Identity
go [a]
pars NonEmpty a -> [Tree a] -> sortKey
sortKey [Tree a]
trees NonEmpty a -> [Tree a] -> Splices (Splice Identity)
childSplice = do
      let extendPars :: a -> NonEmpty a
extendPars a
x = NonEmpty a
-> (NonEmpty a -> NonEmpty a) -> Maybe (NonEmpty a) -> NonEmpty a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (OneItem (NonEmpty a) -> NonEmpty a
forall x. One x => OneItem x -> x
one a
OneItem (NonEmpty a)
x) (NonEmpty a -> NonEmpty a -> NonEmpty a
forall a. Semigroup a => a -> a -> a
<> OneItem (NonEmpty a) -> NonEmpty a
forall x. One x => OneItem x -> x
one a
OneItem (NonEmpty a)
x) (Maybe (NonEmpty a) -> NonEmpty a)
-> Maybe (NonEmpty a) -> NonEmpty a
forall a b. (a -> b) -> a -> b
$ [a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [a]
pars
          nodeKey :: Tree a -> sortKey
nodeKey Tree a
x = NonEmpty a -> [Tree a] -> sortKey
sortKey (a -> NonEmpty a
extendPars (a -> NonEmpty a) -> a -> NonEmpty a
forall a b. (a -> b) -> a -> b
$ Tree a -> a
forall a. Tree a -> a
rootLabel Tree a
x) (Tree a -> [Tree a]
forall a. Tree a -> [Tree a]
subForest Tree a
x)
      ((Tree a -> Splice Identity) -> [Tree a] -> Splice Identity)
-> [Tree a] -> (Tree a -> Splice Identity) -> Splice Identity
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Tree a -> Splice Identity) -> [Tree a] -> Splice Identity
forall b (m :: Type -> Type) (f :: Type -> Type) a.
(Monoid b, Monad m, Foldable f) =>
(a -> m b) -> f a -> m b
foldMapM ((Tree a -> sortKey) -> [Tree a] -> [Tree a]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn Tree a -> sortKey
nodeKey [Tree a]
trees) ((Tree a -> Splice Identity) -> Splice Identity)
-> (Tree a -> Splice Identity) -> Splice Identity
forall a b. (a -> b) -> a -> b
$ \(Node a
lbl [Tree a]
children) -> do
        Splices (Splice Identity) -> Splice Identity
forall (n :: Type -> Type).
Monad n =>
Splices (Splice n) -> Splice n
HI.runChildrenWith (Splices (Splice Identity) -> Splice Identity)
-> Splices (Splice Identity) -> Splice Identity
forall a b. (a -> b) -> a -> b
$ do
          let herePath :: NonEmpty a
herePath = a -> NonEmpty a
extendPars a
lbl
          NonEmpty a -> [Tree a] -> Splices (Splice Identity)
childSplice NonEmpty a
herePath [Tree a]
children
          Text
"has-children" Text -> Splice Identity -> Splices (Splice Identity)
forall k v. k -> v -> MapSyntax k v
## Bool -> Splice Identity
forall (m :: Type -> Type). Monad m => Bool -> Splice m
Heist.ifElseISplice (Bool -> Bool
not (Bool -> Bool) -> ([Tree a] -> Bool) -> [Tree a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree a] -> Bool
forall a. [a] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null ([Tree a] -> Bool) -> [Tree a] -> Bool
forall a b. (a -> b) -> a -> b
$ [Tree a]
children)
          let childKey :: Tree a -> sortKey
childKey Tree a
x = NonEmpty a -> [Tree a] -> sortKey
sortKey (NonEmpty a
herePath NonEmpty a -> NonEmpty a -> NonEmpty a
forall a. Semigroup a => a -> a -> a
<> OneItem (NonEmpty a) -> NonEmpty a
forall x. One x => OneItem x -> x
one (Tree a -> a
forall a. Tree a -> a
rootLabel Tree a
x)) (Tree a -> [Tree a]
forall a. Tree a -> [Tree a]
subForest Tree a
x)
              childrenSorted :: [Tree a]
childrenSorted = (Tree a -> sortKey) -> [Tree a] -> [Tree a]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn Tree a -> sortKey
childKey [Tree a]
children
          Text
"children" Text -> Splice Identity -> Splices (Splice Identity)
forall k v. k -> v -> MapSyntax k v
##
            [a]
-> (NonEmpty a -> [Tree a] -> sortKey)
-> [Tree a]
-> (NonEmpty a -> [Tree a] -> Splices (Splice Identity))
-> Splice Identity
go (NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList NonEmpty a
herePath) NonEmpty a -> [Tree a] -> sortKey
sortKey [Tree a]
childrenSorted NonEmpty a -> [Tree a] -> Splices (Splice Identity)
childSplice