module Geomancy.Tree
  ( Tree(..)

  , apply
  , applyWith
  , mapAccum

  , node_
  , leaf_
  , collect_

  , annotateMap
  , annotateWith
  ) where

import Data.Tree (Tree(..))
import Data.Foldable (toList)

-- * Merging annotations

{- |
  Distribute annotations down the tree without changing the type.
-}
{-# INLINEABLE apply #-}
apply :: Semigroup ann => Tree (ann, a) -> Tree (ann, a)
apply :: forall ann a. Semigroup ann => Tree (ann, a) -> Tree (ann, a)
apply (Node (ann
rootAnn, a
root) [Tree (ann, a)]
rootChildren) =
  forall a. a -> [Tree a] -> Tree a
Node
    (ann
rootAnn, a
root)
    (forall a b. (a -> b) -> [a] -> [b]
map (forall ann acc a.
(ann -> acc -> acc) -> acc -> Tree (ann, a) -> Tree (acc, a)
applyWith forall a. Semigroup a => a -> a -> a
(<>) ann
rootAnn) [Tree (ann, a)]
rootChildren)

{- |
  Distribute accumulator down the tree using the accumulator function.
-}
{-# INLINEABLE applyWith #-}
applyWith
  :: (ann -> acc -> acc)
  -> acc
  -> Tree (ann, a)
  -> Tree (acc, a)
applyWith :: forall ann acc a.
(ann -> acc -> acc) -> acc -> Tree (ann, a) -> Tree (acc, a)
applyWith ann -> acc -> acc
f = forall t a b. (t -> a -> (t, b)) -> t -> Tree a -> Tree b
mapAccum forall {b}. acc -> (ann, b) -> (acc, (acc, b))
next
  where
    next :: acc -> (ann, b) -> (acc, (acc, b))
next acc
acc (ann
ann, b
item) =
      let
        acc' :: acc
acc' = ann -> acc -> acc
f ann
ann acc
acc
      in
        (acc
acc', (acc
acc', b
item))

{- |
  Transform a tree by combining branch-independent accumulator with node contents.
-}
{-# INLINEABLE mapAccum #-}
mapAccum
  :: (t -> a -> (t, b))
  -> t
  -> Tree a
  -> Tree b
mapAccum :: forall t a b. (t -> a -> (t, b)) -> t -> Tree a -> Tree b
mapAccum t -> a -> (t, b)
f t
acc (Node a
item [Tree a]
children) =
  forall a. a -> [Tree a] -> Tree a
Node
    b
nextNode
    (forall a b. (a -> b) -> [a] -> [b]
map (forall t a b. (t -> a -> (t, b)) -> t -> Tree a -> Tree b
mapAccum t -> a -> (t, b)
f t
nextAcc) [Tree a]
children)
  where
    (t
nextAcc, b
nextNode) = t -> a -> (t, b)
f t
acc a
item

-- ** Shortcuts for monoidal annotation and Maybe-wrapped items

{-# INLINEABLE node_ #-}
node_ :: ann -> [Tree (ann, Maybe a)] -> Tree (ann, Maybe a)
node_ :: forall ann a. ann -> [Tree (ann, Maybe a)] -> Tree (ann, Maybe a)
node_ ann
ann = forall a. a -> [Tree a] -> Tree a
Node (ann
ann, forall a. Maybe a
Nothing)

{-# INLINEABLE leaf_ #-}
leaf_ :: Monoid ann => a -> Tree (ann, Maybe a)
leaf_ :: forall ann a. Monoid ann => a -> Tree (ann, Maybe a)
leaf_ a
x = forall a. a -> [Tree a] -> Tree a
Node (forall a. Monoid a => a
mempty, forall a. a -> Maybe a
Just a
x) []

collect_ :: Monoid ann => Tree (ann, Maybe a) -> [(ann, a)]
collect_ :: forall ann a. Monoid ann => Tree (ann, Maybe a) -> [(ann, a)]
collect_ Tree (ann, Maybe a)
root = do
  (ann
ann, Just a
item) <- forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall a b. (a -> b) -> a -> b
$ forall ann a. Semigroup ann => Tree (ann, a) -> Tree (ann, a)
apply Tree (ann, Maybe a)
root
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (ann
ann, a
item)

-- * Adding annotations

{- |
  Annotate nodes with bottom-up monoidal summary.
-}
annotateMap
  :: Monoid ann
  => (a -> ann)
  -> Tree a
  -> Tree (ann, a)
annotateMap :: forall ann a. Monoid ann => (a -> ann) -> Tree a -> Tree (ann, a)
annotateMap a -> ann
f =
  forall a ann.
(a -> ann) -> (a -> [ann] -> ann) -> Tree a -> Tree (ann, a)
annotateWith a -> ann
f (\a
x [ann]
anns -> a -> ann
f a
x forall a. Semigroup a => a -> a -> a
<> forall a. Monoid a => [a] -> a
mconcat [ann]
anns)

{- |
  Annotate the nodes with bottom-up summary.
-}
annotateWith
  :: (a -> ann)
  -> (a -> [ann] -> ann)
  -> Tree a
  -> Tree (ann, a)
annotateWith :: forall a ann.
(a -> ann) -> (a -> [ann] -> ann) -> Tree a -> Tree (ann, a)
annotateWith a -> ann
leaf a -> [ann] -> ann
node = Tree a -> Tree (ann, a)
go
  where
    go :: Tree a -> Tree (ann, a)
go (Node a
x [Tree a]
ts) =
      case [Tree a]
ts of
        [] ->
          forall a. a -> [Tree a] -> Tree a
Node (a -> ann
leaf a
x, a
x) []
        [Tree a]
_ ->
          let
            inner :: [Tree (ann, a)]
inner = forall a b. (a -> b) -> [a] -> [b]
map Tree a -> Tree (ann, a)
go [Tree a]
ts
          in
            forall a. a -> [Tree a] -> Tree a
Node
              ( a -> [ann] -> ann
node
                  a
x
                  (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
rootLabel) [Tree (ann, a)]
inner)
              , a
x
              )
              [Tree (ann, a)]
inner