module Data.Tree.Path where

import Data.List qualified as List
import Data.List.NonEmpty qualified as NE
import Data.Tree (Tree (Node))
import Data.Tree qualified as Tree

treeInsertPath :: Eq a => NonEmpty a -> [Tree a] -> [Tree a]
treeInsertPath :: NonEmpty a -> [Tree a] -> [Tree a]
treeInsertPath =
  (NonEmpty a -> NonEmpty ()) -> NonEmpty a -> [Tree a] -> [Tree a]
forall a ord.
(Eq a, Ord ord) =>
(NonEmpty a -> ord) -> NonEmpty a -> [Tree a] -> [Tree a]
treeInsertPathMaintainingOrder NonEmpty a -> NonEmpty ()
forall (f :: Type -> Type) a. Functor f => f a -> f ()
void

-- | Insert a node by path into a tree with descendants that are ordered.
--
-- Insertion will guarantee that descendants continue to be ordered as expected.
--
-- The order of descendents is determined by the given order function, which
-- takes the path to a node and return that node's order. The intention is to
-- lookup the actual order value which exists *outside* of the tree
-- datastructure itself.
treeInsertPathMaintainingOrder :: (Eq a, Ord ord) => (NonEmpty a -> ord) -> NonEmpty a -> [Tree a] -> [Tree a]
treeInsertPathMaintainingOrder :: (NonEmpty a -> ord) -> NonEmpty a -> [Tree a] -> [Tree a]
treeInsertPathMaintainingOrder NonEmpty a -> ord
ordF NonEmpty a
path [Tree a]
t =
  (NonEmpty a -> ord) -> [a] -> [Tree a] -> [a] -> [Tree a]
forall a b.
(Eq a, Ord b) =>
(NonEmpty a -> b) -> [a] -> [Tree a] -> [a] -> [Tree a]
orderedTreeInsertPath NonEmpty a -> ord
ordF (NonEmpty a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList NonEmpty a
path) [Tree a]
t []
  where
    orderedTreeInsertPath :: (Eq a, Ord b) => (NonEmpty a -> b) -> [a] -> [Tree a] -> [a] -> [Tree a]
    orderedTreeInsertPath :: (NonEmpty a -> b) -> [a] -> [Tree a] -> [a] -> [Tree a]
orderedTreeInsertPath NonEmpty a -> b
_ [] [Tree a]
trees [a]
_ =
      [Tree a]
trees
    orderedTreeInsertPath NonEmpty a -> b
pathOrder (a
top : [a]
rest) [Tree a]
trees [a]
ancestors =
      case a -> [Tree a] -> Maybe (Tree a)
forall {t :: Type -> Type} {a}.
(Foldable t, Eq a) =>
a -> t (Tree a) -> Maybe (Tree a)
treeFindChild a
top [Tree a]
trees of
        Maybe (Tree a)
Nothing ->
          let newChild :: Tree a
newChild = a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
top ([Tree a] -> Tree a) -> [Tree a] -> Tree a
forall a b. (a -> b) -> a -> b
$ (NonEmpty a -> b) -> [a] -> [Tree a] -> [a] -> [Tree a]
forall a b.
(Eq a, Ord b) =>
(NonEmpty a -> b) -> [a] -> [Tree a] -> [a] -> [Tree a]
orderedTreeInsertPath NonEmpty a -> b
pathOrder [a]
rest [] (a
top a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ancestors)
           in (NonEmpty a -> b) -> [Tree a] -> [Tree a]
sortChildrenOn NonEmpty a -> b
pathOrder ([Tree a]
trees [Tree a] -> [Tree a] -> [Tree a]
forall a. Semigroup a => a -> a -> a
<> OneItem [Tree a] -> [Tree a]
forall x. One x => OneItem x -> x
one Tree a
OneItem [Tree a]
newChild)
        Just (Node a
_match [Tree a]
grandChildren) ->
          let oneDead :: [Tree a]
oneDead = a -> [Tree a] -> [Tree a]
forall a. Eq a => a -> [Tree a] -> [Tree a]
treeDeleteChild a
top [Tree a]
trees
              newChild :: Tree a
newChild = a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
top ([Tree a] -> Tree a) -> [Tree a] -> Tree a
forall a b. (a -> b) -> a -> b
$ (NonEmpty a -> b) -> [a] -> [Tree a] -> [a] -> [Tree a]
forall a b.
(Eq a, Ord b) =>
(NonEmpty a -> b) -> [a] -> [Tree a] -> [a] -> [Tree a]
orderedTreeInsertPath NonEmpty a -> b
pathOrder [a]
rest [Tree a]
grandChildren (a
top a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ancestors)
           in (NonEmpty a -> b) -> [Tree a] -> [Tree a]
sortChildrenOn NonEmpty a -> b
pathOrder ([Tree a]
oneDead [Tree a] -> [Tree a] -> [Tree a]
forall a. Semigroup a => a -> a -> a
<> OneItem [Tree a] -> [Tree a]
forall x. One x => OneItem x -> x
one Tree a
OneItem [Tree a]
newChild)
      where
        treeFindChild :: a -> t (Tree a) -> Maybe (Tree a)
treeFindChild a
x t (Tree a)
xs =
          (Tree a -> Bool) -> t (Tree a) -> Maybe (Tree a)
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Maybe a
List.find (\Tree a
n -> Tree a -> a
forall a. Tree a -> a
Tree.rootLabel Tree a
n a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x) t (Tree a)
xs
        sortChildrenOn :: (NonEmpty a -> b) -> [Tree a] -> [Tree a]
sortChildrenOn NonEmpty a -> b
f =
          (Tree a -> b) -> [Tree a] -> [Tree a]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn ((Tree a -> b) -> [Tree a] -> [Tree a])
-> (Tree a -> b) -> [Tree a] -> [Tree a]
forall a b. (a -> b) -> a -> b
$ (\a
s -> NonEmpty a -> b
f (NonEmpty a -> b) -> NonEmpty a -> b
forall a b. (a -> b) -> a -> b
$ NonEmpty a -> NonEmpty a
forall a. NonEmpty a -> NonEmpty a
NE.reverse (NonEmpty a -> NonEmpty a) -> NonEmpty a -> NonEmpty a
forall a b. (a -> b) -> a -> b
$ a
s a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
ancestors) (a -> b) -> (Tree a -> a) -> Tree a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> a
forall a. Tree a -> a
Tree.rootLabel

treeDeletePath :: Eq a => NonEmpty a -> [Tree a] -> [Tree a]
treeDeletePath :: NonEmpty a -> [Tree a] -> [Tree a]
treeDeletePath =
  (a -> [Tree a] -> [Tree a]) -> NonEmpty a -> [Tree a] -> [Tree a]
forall a.
Eq a =>
(a -> [Tree a] -> [Tree a]) -> NonEmpty a -> [Tree a] -> [Tree a]
treeDeletePathWithLastBehavingAs ((a -> [Tree a] -> [Tree a]) -> NonEmpty a -> [Tree a] -> [Tree a])
-> (a -> [Tree a] -> [Tree a])
-> NonEmpty a
-> [Tree a]
-> [Tree a]
forall a b. (a -> b) -> a -> b
$ \a
lastInPath [Tree a]
ts ->
    (Tree a -> Tree a -> Bool) -> Tree a -> [Tree a] -> [Tree a]
forall a. (a -> a -> Bool) -> a -> [a] -> [a]
List.deleteBy (\Tree a
x Tree a
y -> Tree a -> a
forall a. Tree a -> a
Tree.rootLabel Tree a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Tree a -> a
forall a. Tree a -> a
Tree.rootLabel Tree a
y) (a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
lastInPath []) [Tree a]
ts

treeDeleteLeafPath :: Eq a => NonEmpty a -> [Tree a] -> [Tree a]
treeDeleteLeafPath :: NonEmpty a -> [Tree a] -> [Tree a]
treeDeleteLeafPath =
  (a -> [Tree a] -> [Tree a]) -> NonEmpty a -> [Tree a] -> [Tree a]
forall a.
Eq a =>
(a -> [Tree a] -> [Tree a]) -> NonEmpty a -> [Tree a] -> [Tree a]
treeDeletePathWithLastBehavingAs ((a -> [Tree a] -> [Tree a]) -> NonEmpty a -> [Tree a] -> [Tree a])
-> (a -> [Tree a] -> [Tree a])
-> NonEmpty a
-> [Tree a]
-> [Tree a]
forall a b. (a -> b) -> a -> b
$ \a
lastInPath [Tree a]
ts ->
    case [Tree a]
ts of
      [Tree a
t] -> [Tree a
t | Tree a -> a
forall a. Tree a -> a
Tree.rootLabel Tree a
t a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
lastInPath]
      [Tree a]
_ -> [Tree a]
ts

treeDeletePathWithLastBehavingAs :: forall a. Eq a => (a -> [Tree a] -> [Tree a]) -> NonEmpty a -> [Tree a] -> [Tree a]
treeDeletePathWithLastBehavingAs :: (a -> [Tree a] -> [Tree a]) -> NonEmpty a -> [Tree a] -> [Tree a]
treeDeletePathWithLastBehavingAs a -> [Tree a] -> [Tree a]
f NonEmpty a
slugs =
  [a] -> [Tree a] -> [Tree a]
go (NonEmpty a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList NonEmpty a
slugs)
  where
    go :: [a] -> [Tree a] -> [Tree a]
    go :: [a] -> [Tree a] -> [Tree a]
go [] [Tree a]
t = [Tree a]
t
    go [a
p] [Tree a]
ts =
      a -> [Tree a] -> [Tree a]
f a
p [Tree a]
ts
    go (a
p : [a]
ps) [Tree a]
t =
      [Tree a]
t [Tree a] -> (Tree a -> Tree a) -> [Tree a]
forall (f :: Type -> Type) a b. Functor f => f a -> (a -> b) -> f b
<&> \node :: Tree a
node@(Node a
x [Tree a]
xs) ->
        if a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
p
          then a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
x ([Tree a] -> Tree a) -> [Tree a] -> Tree a
forall a b. (a -> b) -> a -> b
$ [a] -> [Tree a] -> [Tree a]
go [a]
ps [Tree a]
xs
          else Tree a
node

treeDeleteChild :: Eq a => a -> [Tree a] -> [Tree a]
treeDeleteChild :: a -> [Tree a] -> [Tree a]
treeDeleteChild a
x =
  (Tree a -> Tree a -> Bool) -> Tree a -> [Tree a] -> [Tree a]
forall a. (a -> a -> Bool) -> a -> [a] -> [a]
List.deleteBy (\Tree a
p Tree a
q -> Tree a -> a
forall a. Tree a -> a
Tree.rootLabel Tree a
p a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Tree a -> a
forall a. Tree a -> a
Tree.rootLabel Tree a
q) (a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
x [])