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

annotatePathsWith :: (NonEmpty a -> ann) -> Tree a -> Tree (a, ann)
annotatePathsWith :: (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 Forest a
children) =
      let path :: NonEmpty a
path = a
rel a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
ancestors
       in (a, ann) -> Forest (a, ann) -> Tree (a, ann)
forall a. a -> Forest a -> Tree a
Node (a
rel, NonEmpty a -> ann
f (NonEmpty a -> ann) -> NonEmpty a -> ann
forall a b. (a -> b) -> a -> b
$ NonEmpty a -> NonEmpty a
forall a. NonEmpty a -> NonEmpty a
NE.reverse NonEmpty a
path) (Forest (a, ann) -> Tree (a, ann))
-> Forest (a, ann) -> Tree (a, ann)
forall a b. (a -> b) -> a -> b
$ (Tree a -> Tree (a, ann)) -> Forest a -> Forest (a, ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([a] -> Tree a -> Tree (a, ann)
go ([a] -> Tree a -> Tree (a, ann)) -> [a] -> Tree a -> Tree (a, ann)
forall a b. (a -> b) -> a -> b
$ NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty a
path) Forest 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 :: (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 Forest a
children) =
      case (Tree a -> Tree a) -> Forest a -> Forest a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Tree a -> Tree a
go Forest a
children of
        [Node a
child Forest a
grandChildren]
          | Just a
new <- a -> a -> Maybe a
f a
parent a
child -> a -> Forest a -> Tree a
forall a. a -> Forest a -> Tree a
Node a
new Forest a
grandChildren
        Forest a
xs -> a -> Forest a -> Tree a
forall a. a -> Forest a -> Tree a
Node a
parent Forest a
xs