content: {-# LANGUAGE NoImplicitPrelude #-}
  
  module Data.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
  import Relude.Extra.Group (groupBy)
  
  mkTreeFromPaths :: Ord a => [[a]] -> Forest a
  mkTreeFromPaths paths = uncurry mkNode <$> Map.assocs groups
    where
      groups = fmap tail <$> groupBy head (mapMaybe nonEmpty paths)
      mkNode label children =
        Node label $ mkTreeFromPaths $ toList children
  
  annotatePathsWith :: (NonEmpty a -> ann) -> Tree a -> Tree (a, ann)
  annotatePathsWith f = go []
    where
      go ancestors (Node rel children) =
        let path = rel :| ancestors
         in Node (rel, f $ NE.reverse path) $ fmap (go $ toList path) 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 f = go
    where
      go (Node parent children) =
        case fmap go children of
          [Node child grandChildren]
            | Just new <- f parent child -> Node new grandChildren
          xs -> Node parent xs
  