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 :: Tree (ann, a) -> Tree (ann, a)
apply (Node (ann
rootAnn, a
root) Forest (ann, a)
rootChildren) =
  (ann, a) -> Forest (ann, a) -> Tree (ann, a)
forall a. a -> Forest a -> Tree a
Node
    (ann
rootAnn, a
root)
    ((Tree (ann, a) -> Tree (ann, a))
-> Forest (ann, a) -> Forest (ann, a)
forall a b. (a -> b) -> [a] -> [b]
map ((ann -> ann -> ann) -> ann -> Tree (ann, a) -> Tree (ann, a)
forall ann acc a.
(ann -> acc -> acc) -> acc -> Tree (ann, a) -> Tree (acc, a)
applyWith ann -> ann -> ann
forall a. Semigroup a => a -> a -> a
(<>) ann
rootAnn) Forest (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 :: (ann -> acc -> acc) -> acc -> Tree (ann, a) -> Tree (acc, a)
applyWith ann -> acc -> acc
f = (acc -> (ann, a) -> (acc, (acc, a)))
-> acc -> Tree (ann, a) -> Tree (acc, a)
forall t a b. (t -> a -> (t, b)) -> t -> Tree a -> Tree b
mapAccum acc -> (ann, a) -> (acc, (acc, a))
forall b. acc -> (ann, b) -> (acc, (acc, b))
next
  where
    -- nextAcc = f ann acc
    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 :: (t -> a -> (t, b)) -> t -> Tree a -> Tree b
mapAccum t -> a -> (t, b)
f t
acc (Node a
item Forest a
children) =
  b -> Forest b -> Tree b
forall a. a -> Forest a -> Tree a
Node
    b
nextNode
    ((Tree a -> Tree b) -> Forest a -> Forest b
forall a b. (a -> b) -> [a] -> [b]
map ((t -> a -> (t, b)) -> t -> Tree a -> Tree b
forall t a b. (t -> a -> (t, b)) -> t -> Tree a -> Tree b
mapAccum t -> a -> (t, b)
f t
nextAcc) Forest 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_ :: ann -> [Tree (ann, Maybe a)] -> Tree (ann, Maybe a)
node_ ann
ann = (ann, Maybe a) -> [Tree (ann, Maybe a)] -> Tree (ann, Maybe a)
forall a. a -> Forest a -> Tree a
Node (ann
ann, Maybe a
forall a. Maybe a
Nothing)

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

collect_ :: Monoid ann => Tree (ann, Maybe a) -> [(ann, a)]
collect_ :: Tree (ann, Maybe a) -> [(ann, a)]
collect_ Tree (ann, Maybe a)
root = do
  (ann
ann, Just a
item) <- Tree (ann, Maybe a) -> [(ann, Maybe a)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Tree (ann, Maybe a) -> [(ann, Maybe a)])
-> Tree (ann, Maybe a) -> [(ann, Maybe a)]
forall a b. (a -> b) -> a -> b
$ Tree (ann, Maybe a) -> Tree (ann, Maybe a)
forall ann a. Semigroup ann => Tree (ann, a) -> Tree (ann, a)
apply Tree (ann, Maybe a)
root
  (ann, a) -> [(ann, a)]
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 :: (a -> ann) -> Tree a -> Tree (ann, a)
annotateMap a -> ann
f =
  (a -> ann) -> (a -> [ann] -> ann) -> Tree a -> Tree (ann, a)
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 ann -> ann -> ann
forall a. Semigroup a => a -> a -> a
<> [ann] -> ann
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 :: (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 Forest a
ts) =
      case Forest a
ts of
        [] ->
          (ann, a) -> Forest (ann, a) -> Tree (ann, a)
forall a. a -> Forest a -> Tree a
Node (a -> ann
leaf a
x, a
x) []
        Forest a
_ ->
          let
            inner :: Forest (ann, a)
inner = (Tree a -> Tree (ann, a)) -> Forest a -> Forest (ann, a)
forall a b. (a -> b) -> [a] -> [b]
map Tree a -> Tree (ann, a)
go Forest a
ts
          in
            (ann, a) -> Forest (ann, a) -> Tree (ann, a)
forall a. a -> Forest a -> Tree a
Node
              ( a -> [ann] -> ann
node
                  a
x
                  ((Tree (ann, a) -> ann) -> Forest (ann, a) -> [ann]
forall a b. (a -> b) -> [a] -> [b]
map ((ann, a) -> ann
forall a b. (a, b) -> a
fst ((ann, a) -> ann)
-> (Tree (ann, a) -> (ann, a)) -> Tree (ann, a) -> ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree (ann, a) -> (ann, a)
forall a. Tree a -> a
rootLabel) Forest (ann, a)
inner)
              , a
x
              )
              Forest (ann, a)
inner