module Data.TagTree.PathTree
  ( mkTreeFromPaths,
    annotatePathsWith,
    foldSingleParentsWith,
  )
where

import qualified Data.List.NonEmpty as NE
import qualified Data.Map as Map
import Data.Tree (Forest, Tree (Node))
import Relude.Extra.Group (groupBy)

mkTreeFromPaths :: Ord a => [[a]] -> Forest a
mkTreeFromPaths :: forall a. Ord a => [[a]] -> Forest a
mkTreeFromPaths [[a]]
paths = forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall {a} {t :: * -> *}.
(Ord a, Foldable t) =>
a -> t [a] -> Tree a
mkNode forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k a. Map k a -> [(k, a)]
Map.assocs Map a (NonEmpty [a])
groups
  where
    groups :: Map a (NonEmpty [a])
groups = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *) a. IsNonEmpty f a [a] "tail" => f a -> [a]
tail forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) t a.
(Foldable f, DynamicMap t, Val t ~ NonEmpty a, Monoid t) =>
(a -> Key t) -> f a -> t
groupBy forall (f :: * -> *) a. IsNonEmpty f a a "head" => f a -> a
head (forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [[a]]
paths)
    mkNode :: a -> t [a] -> Tree a
mkNode a
label t [a]
children =
      forall a. a -> [Tree a] -> Tree a
Node a
label forall a b. (a -> b) -> a -> b
$ forall a. Ord a => [[a]] -> Forest a
mkTreeFromPaths forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t [a]
children

annotatePathsWith :: (NonEmpty a -> ann) -> Tree a -> Tree (a, ann)
annotatePathsWith :: forall a ann. (NonEmpty a -> ann) -> Tree a -> Tree (a, ann)
annotatePathsWith NonEmpty a -> ann
f = [a] -> Tree a -> Tree (a, ann)
go []
  where
    go :: [a] -> Tree a -> Tree (a, ann)
go [a]
ancestors (Node a
rel [Tree a]
children) =
      let path :: NonEmpty a
path = a
rel forall a. a -> [a] -> NonEmpty a
:| [a]
ancestors
       in forall a. a -> [Tree a] -> Tree a
Node (a
rel, NonEmpty a -> ann
f forall a b. (a -> b) -> a -> b
$ forall a. NonEmpty a -> NonEmpty a
NE.reverse NonEmpty a
path) forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([a] -> Tree a -> Tree (a, ann)
go forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty a
path) [Tree a]
children

-- | Fold nodes with one child using the given function
--
-- The function is called with the parent and the only child. If a Just value is
-- returned, folding happens with that value, otherwise there is no effect.
foldSingleParentsWith :: (a -> a -> Maybe a) -> Tree a -> Tree a
foldSingleParentsWith :: forall a. (a -> a -> Maybe a) -> Tree a -> Tree a
foldSingleParentsWith a -> a -> Maybe a
f = Tree a -> Tree a
go
  where
    go :: Tree a -> Tree a
go (Node a
parent [Tree a]
children) =
      case forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Tree a -> Tree a
go [Tree a]
children of
        [Node a
child [Tree a]
grandChildren]
          | Just a
new <- a -> a -> Maybe a
f a
parent a
child -> forall a. a -> [Tree a] -> Tree a
Node a
new [Tree a]
grandChildren
        [Tree a]
xs -> forall a. a -> [Tree a] -> Tree a
Node a
parent [Tree a]
xs