From 5740d17fb13d789f9e5b2725f19ecf88d45003da Mon Sep 17 00:00:00 2001
From: Jakub Hampl <kopomir@gmail.com>
Date: Mon, 26 Jun 2023 23:02:47 +0100
Subject: [PATCH 1/3] Adds a number of useful functions

- depth
- leaves
- links
- context maps (bottom up and top down, with and without accumulator)
- findBfs
- sortWith
---
 src/Tree.elm | 405 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 399 insertions(+), 6 deletions(-)

diff --git a/src/Tree.elm b/src/Tree.elm
index adf28f1..b7fd6f5 100644
--- a/src/Tree.elm
+++ b/src/Tree.elm
@@ -1,9 +1,12 @@
 module Tree exposing
     ( Tree, singleton, tree, label, children
     , mapLabel, replaceLabel, mapChildren, replaceChildren, prependChild, appendChild
-    , foldl, foldr, count, flatten
-    , map, indexedMap, mapAccumulate, map2, indexedMap2, mapAccumulate2, andMap
-    , unfold, restructure
+    , count, depth
+    , foldl, foldr
+    , flatten, leaves, links
+    , map, indexedMap, mapAccumulate, map2, indexedMap2, mapAccumulate2, andMap, mapWithContextBottomUp, mapAccumulateWithContextBottomUp, mapWithContextTopDown, mapAccumulateWithContextTopDown
+    , findBfs
+    , sortWith, unfold, restructure
     )
 
 {-| A multiway tree or rosetree is a labeled tree where each node can have zero,
@@ -40,19 +43,34 @@ children of their own, and so on.
 @docs mapLabel, replaceLabel, mapChildren, replaceChildren, prependChild, appendChild
 
 
+# Describing a tree
+
+@docs count, depth
+
+
 # Folds
 
-@docs foldl, foldr, count, flatten
+@docs foldl, foldr
+
+
+# Converting to lists
+
+@docs flatten, leaves, links
 
 
 # Mapping and traversing
 
-@docs map, indexedMap, mapAccumulate, map2, indexedMap2, mapAccumulate2, andMap
+@docs map, indexedMap, mapAccumulate, map2, indexedMap2, mapAccumulate2, andMap, mapWithContextBottomUp, mapAccumulateWithContextBottomUp, mapWithContextTopDown, mapAccumulateWithContextTopDown
+
+
+# Search
+
+@docs findBfs
 
 
 # Fancy stuff
 
-@docs unfold, restructure
+@docs sortWith, unfold, restructure
 
 -}
 
@@ -293,6 +311,71 @@ flatten t =
     foldr (::) [] t
 
 
+{-| Returns the nodes that have no children.
+-}
+leaves : Tree a -> List a
+leaves t =
+    leavesHelp [] [ t ] []
+
+
+leavesHelp : List a -> List (Tree a) -> List (List (Tree a)) -> List a
+leavesHelp soFar trees nextSets =
+    case trees of
+        [] ->
+            case nextSets of
+                set :: sets ->
+                    leavesHelp soFar set sets
+
+                [] ->
+                    soFar
+
+        (Tree d []) :: rest ->
+            leavesHelp (d :: soFar) rest nextSets
+
+        (Tree _ xs) :: rest ->
+            leavesHelp soFar rest (xs :: nextSets)
+
+
+{-| Returns pairs representing parent-child relationships in the tree.
+
+The left item is the label of the parent, the right item is the label of
+the child. Useful for visualising trees.
+
+    tree 1
+        [ singleton 2
+        , tree 3
+            [ singleton 4
+            , singleton 5
+            ]
+        , singleton 6
+        ]
+        |> links
+    --> [ ( 1, 2 ), ( 1, 3 ), ( 1, 6 ), ( 3, 4 ), ( 3, 5 ) ]
+
+-}
+links : Tree a -> List ( a, a )
+links (Tree l cr) =
+    linksHelp l [] cr []
+
+
+linksHelp : a -> List ( a, a ) -> List (Tree a) -> List ( a, List (Tree a) ) -> List ( a, a )
+linksHelp parent soFar trees nextSets =
+    case trees of
+        [] ->
+            case nextSets of
+                ( newParent, set ) :: sets ->
+                    linksHelp newParent soFar set sets
+
+                [] ->
+                    List.reverse soFar
+
+        (Tree d []) :: rest ->
+            linksHelp parent (( parent, d ) :: soFar) rest nextSets
+
+        (Tree d xs) :: rest ->
+            linksHelp parent (( parent, d ) :: soFar) rest (( d, xs ) :: nextSets)
+
+
 {-| Create a tree from a seed.
 
 Running the function on the seed should return a label and a list of seeds to
@@ -651,6 +734,316 @@ type alias Map2Acc a b c =
     }
 
 
+{-| Counts the number of levels in a tree (where the root is 0).
+
+    depth (tree 2 [ tree 1 [tree 0 []]])
+    --> 2
+
+-}
+depth : Tree a -> Int
+depth t =
+    Tuple.first (mapAccumulateWithContextTopDown (\s ctx -> ( max s (List.length ctx.ancestors), () )) 0 t)
+
+
+{-| Maps the tree, but gives the mapping function access to the context of the tree.
+
+This is quite useful when using a tree to model hierarchical relationships, as the
+position of the nodes of the tree is often as significant to operations on the tree as
+are the values contained.
+
+This is a pre-order traversal of the tree, meaning that the tree is transformed from the
+root to the leaves, giving you access to untransformed children and transformed parent nodes.
+
+-}
+mapWithContextTopDown : ({ ancestors : List b, node : a, children : List (Tree a) } -> b) -> Tree a -> Tree b
+mapWithContextTopDown f t =
+    mapAccumulateWithContextTopDown (\_ e -> ( (), f e )) () t
+        |> Tuple.second
+
+
+{-| Maps the tree, but gives the mapping function access to the context of the tree.
+
+This is a post-order traversal of the tree, meaning that the tree is transformed from the
+leaves to the root, giving you access to transformed children and untransformed parent nodes.
+
+-}
+mapWithContextBottomUp : ({ ancestors : List a, node : a, children : List (Tree b) } -> b) -> Tree a -> Tree b
+mapWithContextBottomUp f t =
+    mapAccumulateWithContextBottomUp (\_ e -> ( (), f e )) () t
+        |> Tuple.second
+
+
+{-| Like `mapWithContextTopDown`, but with an accumulator argument.
+-}
+mapAccumulateWithContextTopDown : (s -> { ancestors : List b, node : a, children : List (Tree a) } -> ( s, b )) -> s -> Tree a -> ( s, Tree b )
+mapAccumulateWithContextTopDown f s (Tree d cs) =
+    let
+        ( s_, d_ ) =
+            f s { ancestors = [], node = d, children = cs }
+    in
+    mapAccumulateWithContextTopDownHelp f
+        s_
+        { todo = cs
+        , done = []
+        , label = d_
+        , parents = [ d_ ]
+        }
+        []
+
+
+mapAccumulateWithContextTopDownHelp :
+    (s -> { ancestors : List b, node : a, children : List (Tree a) } -> ( s, b ))
+    -> s
+    -> MapAccCtxTopDown a b
+    -> List (MapAccCtxTopDown a b)
+    -> ( s, Tree b )
+mapAccumulateWithContextTopDownHelp f state acc stack =
+    case acc.todo of
+        [] ->
+            let
+                node =
+                    Tree acc.label (List.reverse acc.done)
+            in
+            case stack of
+                [] ->
+                    ( state, node )
+
+                top :: rest ->
+                    mapAccumulateWithContextTopDownHelp f state { top | done = node :: top.done } rest
+
+        (Tree d []) :: rest ->
+            let
+                ( state_, label_ ) =
+                    f state { children = [], node = d, ancestors = acc.parents }
+            in
+            mapAccumulateWithContextTopDownHelp f
+                state_
+                { acc
+                    | todo = rest
+                    , done = Tree label_ [] :: acc.done
+                }
+                stack
+
+        (Tree d cs) :: rest ->
+            let
+                ( state_, label_ ) =
+                    f state { children = cs, node = d, ancestors = acc.parents }
+            in
+            mapAccumulateWithContextTopDownHelp f
+                state_
+                { todo = cs
+                , done = []
+                , label = label_
+                , parents = label_ :: acc.parents
+                }
+                ({ acc | todo = rest } :: stack)
+
+
+type alias MapAccCtxTopDown a b =
+    { todo : List (Tree a)
+    , done : List (Tree b)
+    , parents : List b
+    , label : b
+    }
+
+
+{-| Like `mapWithContextBottomUp`, but with an accumulator argument.
+-}
+mapAccumulateWithContextBottomUp : (s -> { ancestors : List a, node : a, children : List (Tree b) } -> ( s, b )) -> s -> Tree a -> ( s, Tree b )
+mapAccumulateWithContextBottomUp f s (Tree d cs) =
+    mapAccumulateWithContextBottomUpHelp f
+        s
+        { todo = cs
+        , done = []
+        , label = d
+        , parents = [ d ]
+        }
+        []
+
+
+mapAccumulateWithContextBottomUpHelp :
+    (s -> { ancestors : List a, node : a, children : List (Tree b) } -> ( s, b ))
+    -> s
+    -> MapAccCtxBottomUp a b
+    -> List (MapAccCtxBottomUp a b)
+    -> ( s, Tree b )
+mapAccumulateWithContextBottomUpHelp f state acc stack =
+    case acc.todo of
+        [] ->
+            let
+                children_ =
+                    List.reverse acc.done
+
+                ( state_, label_ ) =
+                    f state { children = children_, node = acc.label, ancestors = acc.parents }
+
+                node =
+                    Tree label_ children_
+            in
+            case stack of
+                [] ->
+                    ( state_, node )
+
+                top :: rest ->
+                    mapAccumulateWithContextBottomUpHelp f state { top | done = node :: top.done } rest
+
+        (Tree d []) :: rest ->
+            let
+                ( state_, label_ ) =
+                    f state { children = [], node = d, ancestors = acc.parents }
+            in
+            mapAccumulateWithContextBottomUpHelp f
+                state_
+                { acc
+                    | todo = rest
+                    , done = Tree label_ [] :: acc.done
+                }
+                stack
+
+        (Tree d cs) :: rest ->
+            mapAccumulateWithContextBottomUpHelp f
+                state
+                { todo = cs
+                , done = []
+                , label = d
+                , parents = d :: acc.parents
+                }
+                ({ acc | todo = rest } :: stack)
+
+
+type alias MapAccCtxBottomUp a b =
+    { todo : List (Tree a)
+    , done : List (Tree b)
+    , parents : List a
+    , label : a
+    }
+
+
+{-| Sorts all children of each node based on the comparator function (the function recieves a list of ancestors).
+
+    tree 1
+        [ tree 3
+            [ singleton 5
+            , singleton 4
+            ]
+        , singleton 2
+        , singleton 6
+        ]
+        |> sortWith (\_ a b -> compare (label a) (label b))
+    --> tree 1
+    -->   [ singleton 2
+    -->    , tree 3
+    -->        [ singleton 4
+    -->        , singleton 5
+    -->        ]
+    -->    , singleton 6
+    -->    ]
+
+-}
+sortWith : (List a -> Tree a -> Tree a -> Order) -> Tree a -> Tree a
+sortWith compareFn (Tree d cs) =
+    sortWithHelp compareFn
+        { todo = cs
+        , done = []
+        , label = d
+        , parents = [ d ]
+        }
+        []
+
+
+sortWithHelp : (List a -> Tree a -> Tree a -> Order) -> MapAccCtxBottomUp a a -> List (MapAccCtxBottomUp a a) -> Tree a
+sortWithHelp compareFn acc stack =
+    case acc.todo of
+        [] ->
+            let
+                children_ =
+                    List.sortWith (compareFn acc.parents) acc.done
+
+                node =
+                    Tree acc.label children_
+            in
+            case stack of
+                [] ->
+                    node
+
+                top :: rest ->
+                    sortWithHelp compareFn { top | done = node :: top.done } rest
+
+        (Tree d []) :: rest ->
+            sortWithHelp compareFn
+                { acc
+                    | todo = rest
+                    , done = Tree d [] :: acc.done
+                }
+                stack
+
+        (Tree d cs) :: rest ->
+            sortWithHelp compareFn
+                { todo = cs
+                , done = []
+                , label = d
+                , parents = d :: acc.parents
+                }
+                ({ acc | todo = rest } :: stack)
+
+
+{-| Finds a subtree whose label matches the predicate.
+
+Searches the tree in a breadth-first manner.
+
+    tree 1
+        [ tree 3
+            [ singleton 5
+            , singleton 4
+            ]
+        , singleton 2
+        , singleton 6
+        ]
+        |> findBfs (\a -> label a == 3)
+    --> Just (tree 3 [ singleton 5, singleton 4 ])
+
+-}
+findBfs : (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
+findBfs predicate t =
+    findBfsHelp predicate (Fifo [ t ] [])
+
+
+findBfsHelp : (Tree a -> Bool) -> Queue (Tree a) -> Maybe (Tree a)
+findBfsHelp predicate queue =
+    case removeQueue queue of
+        ( Just t, rest ) ->
+            if predicate t then
+                Just t
+
+            else
+                findBfsHelp predicate (insertQueue t rest)
+
+        ( Nothing, _ ) ->
+            Nothing
+
+
+type Queue a
+    = Fifo (List a) (List a)
+
+
+insertQueue : Tree a -> Queue (Tree a) -> Queue (Tree a)
+insertQueue (Tree _ cs) (Fifo front back) =
+    Fifo front (List.reverse cs ++ back)
+
+
+removeQueue : Queue a -> ( Maybe a, Queue a )
+removeQueue fifo =
+    case fifo of
+        Fifo [] [] ->
+            ( Nothing, fifo )
+
+        Fifo [] back ->
+            removeQueue <| Fifo (List.reverse back) []
+
+        Fifo (next :: rest) back ->
+            ( Just next, Fifo rest back )
+
+
 {-| Restructure a `Tree` into another type of structure.
 
 Imagine you have a `Tree String` and you can to turn it into nested `<ul>`s.

From 67a9245f44e996dace288db1ecc8a2e5582db7d5 Mon Sep 17 00:00:00 2001
From: Jakub Hampl <kopomir@gmail.com>
Date: Tue, 18 Jul 2023 10:41:34 +0300
Subject: [PATCH 2/3] Adds generic traversals

---
 src/Tree.elm       | 791 +++++++++++++++++++--------------------------
 tests/TreeTest.elm |  68 +++-
 2 files changed, 399 insertions(+), 460 deletions(-)

diff --git a/src/Tree.elm b/src/Tree.elm
index b7fd6f5..46afad1 100644
--- a/src/Tree.elm
+++ b/src/Tree.elm
@@ -4,9 +4,10 @@ module Tree exposing
     , count, depth
     , foldl, foldr
     , flatten, leaves, links
-    , map, indexedMap, mapAccumulate, map2, indexedMap2, mapAccumulate2, andMap, mapWithContextBottomUp, mapAccumulateWithContextBottomUp, mapWithContextTopDown, mapAccumulateWithContextTopDown
+    , map, indexedMap, mapAccumulate, map2, indexedMap2, mapAccumulate2, andMap
     , findBfs
     , sortWith, unfold, restructure
+    , Step(..), breadthFirstFold, depthFirstFold, depthFirstTraversal
     )
 
 {-| A multiway tree or rosetree is a labeled tree where each node can have zero,
@@ -60,7 +61,7 @@ children of their own, and so on.
 
 # Mapping and traversing
 
-@docs map, indexedMap, mapAccumulate, map2, indexedMap2, mapAccumulate2, andMap, mapWithContextBottomUp, mapAccumulateWithContextBottomUp, mapWithContextTopDown, mapAccumulateWithContextTopDown
+@docs map, indexedMap, mapAccumulate, map2, indexedMap2, mapAccumulate2, andMap
 
 
 # Search
@@ -72,6 +73,30 @@ children of their own, and so on.
 
 @docs sortWith, unfold, restructure
 
+
+# Advanced: Generic traversals
+
+These functions have highly complex type signatures, but they abstract very
+generic ways of working with trees and in fact nearly all the other functions
+in this library are built using them. In general it is better to prefer the simpler
+interfaces, but there are situations that may not be covered by other functions
+in this library, where these more powerful functions can come in handy.
+
+Note that all the callbacks passed receive four arguments:
+
+  - the `state` variable that is accumulated throughout the whole computation
+  - a list of `ancestors`, that is all the labels that lie above the current node
+  - the current `label` of the node being processed
+  - its `children`
+
+I like to call these `\s a l c`, since "salc" is nice and pronouncable and quite
+easy to remember.
+
+@docs Step, breadthFirstFold, depthFirstFold, depthFirstTraversal
+
+You may want to read the source of this module for inspiration on how to use these
+functions.
+
 -}
 
 
@@ -264,7 +289,7 @@ count t =
 -}
 foldl : (a -> b -> b) -> b -> Tree a -> b
 foldl f acc t =
-    foldlHelp f acc [ t ] []
+    depthFirstFold (\s _ l _ -> Continue (f l s)) acc t
 
 
 {-| Fold over all the labels in a tree, right to left, depth first.
@@ -286,22 +311,20 @@ foldr f acc t =
     List.foldl f acc <| foldl (::) [] t
 
 
-foldlHelp : (a -> b -> b) -> b -> List (Tree a) -> List (List (Tree a)) -> b
-foldlHelp f acc trees nextSets =
-    case trees of
-        [] ->
-            case nextSets of
-                set :: sets ->
-                    foldlHelp f acc set sets
 
-                [] ->
-                    acc
-
-        (Tree d []) :: rest ->
-            foldlHelp f (f d acc) rest nextSets
-
-        (Tree d xs) :: rest ->
-            foldlHelp f (f d acc) xs (rest :: nextSets)
+-- foldlHelp : (a -> b -> b) -> b -> List (Tree a) -> List (List (Tree a)) -> b
+-- foldlHelp f acc trees nextSets =
+--     case trees of
+--         [] ->
+--             case nextSets of
+--                 set :: sets ->
+--                     foldlHelp f acc set sets
+--                 [] ->
+--                     acc
+--         (Tree d []) :: rest ->
+--             foldlHelp f (f d acc) rest nextSets
+--         (Tree d xs) :: rest ->
+--             foldlHelp f (f d acc) xs (rest :: nextSets)
 
 
 {-| Flattens the tree into a list. This is equivalent to `foldr (::) []`
@@ -312,28 +335,34 @@ flatten t =
 
 
 {-| Returns the nodes that have no children.
+
+    tree 1
+        [ singleton 2
+        , tree 3
+            [ singleton 4
+            , tree 5
+                [ singleton 6]
+            ]
+        , singleton 7
+        ]
+        |> leaves
+    --> [ 2, 7, 4, 6 ]
+
 -}
 leaves : Tree a -> List a
 leaves t =
-    leavesHelp [] [ t ] []
-
-
-leavesHelp : List a -> List (Tree a) -> List (List (Tree a)) -> List a
-leavesHelp soFar trees nextSets =
-    case trees of
-        [] ->
-            case nextSets of
-                set :: sets ->
-                    leavesHelp soFar set sets
-
+    breadthFirstFold
+        (\s _ l c ->
+            case c of
                 [] ->
-                    soFar
+                    Continue (l :: s)
 
-        (Tree d []) :: rest ->
-            leavesHelp (d :: soFar) rest nextSets
-
-        (Tree _ xs) :: rest ->
-            leavesHelp soFar rest (xs :: nextSets)
+                _ ->
+                    Continue s
+        )
+        []
+        t
+        |> List.reverse
 
 
 {-| Returns pairs representing parent-child relationships in the tree.
@@ -345,35 +374,45 @@ the child. Useful for visualising trees.
         [ singleton 2
         , tree 3
             [ singleton 4
-            , singleton 5
+            , tree 5
+                [ singleton 6]
             ]
-        , singleton 6
+        , singleton 7
         ]
         |> links
-    --> [ ( 1, 2 ), ( 1, 3 ), ( 1, 6 ), ( 3, 4 ), ( 3, 5 ) ]
+    --> [ ( 1, 2 ), ( 1, 3 ), ( 1, 7 ), ( 3, 4 ), ( 3, 5 ), ( 5, 6 ) ]
 
 -}
 links : Tree a -> List ( a, a )
-links (Tree l cr) =
-    linksHelp l [] cr []
-
-
-linksHelp : a -> List ( a, a ) -> List (Tree a) -> List ( a, List (Tree a) ) -> List ( a, a )
-linksHelp parent soFar trees nextSets =
-    case trees of
-        [] ->
-            case nextSets of
-                ( newParent, set ) :: sets ->
-                    linksHelp newParent soFar set sets
+links t =
+    breadthFirstFold
+        (\s a l _ ->
+            case a of
+                parent :: _ ->
+                    Continue (( parent, l ) :: s)
+
+                _ ->
+                    Continue s
+        )
+        []
+        t
+        |> List.reverse
 
-                [] ->
-                    List.reverse soFar
 
-        (Tree d []) :: rest ->
-            linksHelp parent (( parent, d ) :: soFar) rest nextSets
 
-        (Tree d xs) :: rest ->
-            linksHelp parent (( parent, d ) :: soFar) rest (( d, xs ) :: nextSets)
+-- linksHelp : a -> List ( a, a ) -> List (Tree a) -> List ( a, List (Tree a) ) -> List ( a, a )
+-- linksHelp parent soFar trees nextSets =
+--     case trees of
+--         [] ->
+--             case nextSets of
+--                 ( newParent, set ) :: sets ->
+--                     linksHelp newParent soFar set sets
+--                 [] ->
+--                     List.reverse soFar
+--         (Tree d []) :: rest ->
+--             linksHelp parent (( parent, d ) :: soFar) rest nextSets
+--         (Tree d xs) :: rest ->
+--             linksHelp parent (( parent, d ) :: soFar) rest (( d, xs ) :: nextSets)
 
 
 {-| Create a tree from a seed.
@@ -404,58 +443,18 @@ has the number of children mentioned in the label, recursively.
 -}
 unfold : (b -> ( a, List b )) -> b -> Tree a
 unfold f seed =
-    let
-        ( v, next ) =
-            f seed
-    in
-    unfoldHelp f { todo = next, label = v, done = [] } []
-
-
-unfoldHelp :
-    (b -> ( a, List b ))
-    -> UnfoldAcc a b
-    -> List (UnfoldAcc a b)
-    -> Tree a
-unfoldHelp f acc stack =
-    case acc.todo of
-        [] ->
+    depthFirstTraversal
+        (\s _ l _ ->
             let
-                node =
-                    Tree acc.label (List.reverse acc.done)
+                ( l_, c ) =
+                    f l
             in
-            case stack of
-                [] ->
-                    node
-
-                top :: rest ->
-                    unfoldHelp f
-                        { top | done = node :: top.done }
-                        rest
-
-        x :: xs ->
-            case f x of
-                ( label_, [] ) ->
-                    unfoldHelp f
-                        { acc
-                            | todo = xs
-                            , done = singleton label_ :: acc.done
-                        }
-                        stack
-
-                ( label_, todo ) ->
-                    unfoldHelp f
-                        { todo = todo
-                        , label = label_
-                        , done = []
-                        }
-                        ({ acc | todo = xs } :: stack)
-
-
-type alias UnfoldAcc a b =
-    { todo : List b
-    , done : List (Tree a)
-    , label : a
-    }
+            ( s, l_, List.map singleton c )
+        )
+        defaultBottomUp
+        ()
+        (singleton seed)
+        |> Tuple.second
 
 
 {-| Run a function on every label in the tree.
@@ -520,72 +519,18 @@ indexedMap f t =
 
 -}
 mapAccumulate : (s -> a -> ( s, b )) -> s -> Tree a -> ( s, Tree b )
-mapAccumulate f s (Tree d cs) =
-    let
-        ( s_, d_ ) =
-            f s d
-    in
-    mapAccumulateHelp f
-        s_
-        { todo = cs
-        , done = []
-        , label = d_
-        }
-        []
-
-
-mapAccumulateHelp :
-    (s -> a -> ( s, b ))
-    -> s
-    -> MapAcc a b
-    -> List (MapAcc a b)
-    -> ( s, Tree b )
-mapAccumulateHelp f state acc stack =
-    case acc.todo of
-        [] ->
-            let
-                node =
-                    Tree acc.label (List.reverse acc.done)
-            in
-            case stack of
-                [] ->
-                    ( state, node )
-
-                top :: rest ->
-                    mapAccumulateHelp f state { top | done = node :: top.done } rest
-
-        (Tree d []) :: rest ->
+mapAccumulate f state t =
+    depthFirstTraversal
+        (\s _ l c ->
             let
-                ( state_, label_ ) =
-                    f state d
+                ( s_, l_ ) =
+                    f s l
             in
-            mapAccumulateHelp f
-                state_
-                { acc
-                    | todo = rest
-                    , done = Tree label_ [] :: acc.done
-                }
-                stack
-
-        (Tree d cs) :: rest ->
-            let
-                ( state_, label_ ) =
-                    f state d
-            in
-            mapAccumulateHelp f
-                state_
-                { todo = cs
-                , done = []
-                , label = label_
-                }
-                ({ acc | todo = rest } :: stack)
-
-
-type alias MapAcc a b =
-    { todo : List (Tree a)
-    , done : List (Tree b)
-    , label : b
-    }
+            ( s_, l_, c )
+        )
+        defaultBottomUp
+        state
+        t
 
 
 {-| Map over 2 trees. Much like `List.map2`, the result will be truncated to the shorter result.
@@ -742,181 +687,17 @@ type alias Map2Acc a b c =
 -}
 depth : Tree a -> Int
 depth t =
-    Tuple.first (mapAccumulateWithContextTopDown (\s ctx -> ( max s (List.length ctx.ancestors), () )) 0 t)
-
-
-{-| Maps the tree, but gives the mapping function access to the context of the tree.
-
-This is quite useful when using a tree to model hierarchical relationships, as the
-position of the nodes of the tree is often as significant to operations on the tree as
-are the values contained.
-
-This is a pre-order traversal of the tree, meaning that the tree is transformed from the
-root to the leaves, giving you access to untransformed children and transformed parent nodes.
-
--}
-mapWithContextTopDown : ({ ancestors : List b, node : a, children : List (Tree a) } -> b) -> Tree a -> Tree b
-mapWithContextTopDown f t =
-    mapAccumulateWithContextTopDown (\_ e -> ( (), f e )) () t
-        |> Tuple.second
-
-
-{-| Maps the tree, but gives the mapping function access to the context of the tree.
-
-This is a post-order traversal of the tree, meaning that the tree is transformed from the
-leaves to the root, giving you access to transformed children and untransformed parent nodes.
-
--}
-mapWithContextBottomUp : ({ ancestors : List a, node : a, children : List (Tree b) } -> b) -> Tree a -> Tree b
-mapWithContextBottomUp f t =
-    mapAccumulateWithContextBottomUp (\_ e -> ( (), f e )) () t
-        |> Tuple.second
-
-
-{-| Like `mapWithContextTopDown`, but with an accumulator argument.
--}
-mapAccumulateWithContextTopDown : (s -> { ancestors : List b, node : a, children : List (Tree a) } -> ( s, b )) -> s -> Tree a -> ( s, Tree b )
-mapAccumulateWithContextTopDown f s (Tree d cs) =
-    let
-        ( s_, d_ ) =
-            f s { ancestors = [], node = d, children = cs }
-    in
-    mapAccumulateWithContextTopDownHelp f
-        s_
-        { todo = cs
-        , done = []
-        , label = d_
-        , parents = [ d_ ]
-        }
-        []
-
-
-mapAccumulateWithContextTopDownHelp :
-    (s -> { ancestors : List b, node : a, children : List (Tree a) } -> ( s, b ))
-    -> s
-    -> MapAccCtxTopDown a b
-    -> List (MapAccCtxTopDown a b)
-    -> ( s, Tree b )
-mapAccumulateWithContextTopDownHelp f state acc stack =
-    case acc.todo of
-        [] ->
-            let
-                node =
-                    Tree acc.label (List.reverse acc.done)
-            in
-            case stack of
+    depthFirstFold
+        (\s a _ c ->
+            case c of
                 [] ->
-                    ( state, node )
-
-                top :: rest ->
-                    mapAccumulateWithContextTopDownHelp f state { top | done = node :: top.done } rest
-
-        (Tree d []) :: rest ->
-            let
-                ( state_, label_ ) =
-                    f state { children = [], node = d, ancestors = acc.parents }
-            in
-            mapAccumulateWithContextTopDownHelp f
-                state_
-                { acc
-                    | todo = rest
-                    , done = Tree label_ [] :: acc.done
-                }
-                stack
-
-        (Tree d cs) :: rest ->
-            let
-                ( state_, label_ ) =
-                    f state { children = cs, node = d, ancestors = acc.parents }
-            in
-            mapAccumulateWithContextTopDownHelp f
-                state_
-                { todo = cs
-                , done = []
-                , label = label_
-                , parents = label_ :: acc.parents
-                }
-                ({ acc | todo = rest } :: stack)
-
+                    Continue (max s (List.length a))
 
-type alias MapAccCtxTopDown a b =
-    { todo : List (Tree a)
-    , done : List (Tree b)
-    , parents : List b
-    , label : b
-    }
-
-
-{-| Like `mapWithContextBottomUp`, but with an accumulator argument.
--}
-mapAccumulateWithContextBottomUp : (s -> { ancestors : List a, node : a, children : List (Tree b) } -> ( s, b )) -> s -> Tree a -> ( s, Tree b )
-mapAccumulateWithContextBottomUp f s (Tree d cs) =
-    mapAccumulateWithContextBottomUpHelp f
-        s
-        { todo = cs
-        , done = []
-        , label = d
-        , parents = [ d ]
-        }
-        []
-
-
-mapAccumulateWithContextBottomUpHelp :
-    (s -> { ancestors : List a, node : a, children : List (Tree b) } -> ( s, b ))
-    -> s
-    -> MapAccCtxBottomUp a b
-    -> List (MapAccCtxBottomUp a b)
-    -> ( s, Tree b )
-mapAccumulateWithContextBottomUpHelp f state acc stack =
-    case acc.todo of
-        [] ->
-            let
-                children_ =
-                    List.reverse acc.done
-
-                ( state_, label_ ) =
-                    f state { children = children_, node = acc.label, ancestors = acc.parents }
-
-                node =
-                    Tree label_ children_
-            in
-            case stack of
-                [] ->
-                    ( state_, node )
-
-                top :: rest ->
-                    mapAccumulateWithContextBottomUpHelp f state { top | done = node :: top.done } rest
-
-        (Tree d []) :: rest ->
-            let
-                ( state_, label_ ) =
-                    f state { children = [], node = d, ancestors = acc.parents }
-            in
-            mapAccumulateWithContextBottomUpHelp f
-                state_
-                { acc
-                    | todo = rest
-                    , done = Tree label_ [] :: acc.done
-                }
-                stack
-
-        (Tree d cs) :: rest ->
-            mapAccumulateWithContextBottomUpHelp f
-                state
-                { todo = cs
-                , done = []
-                , label = d
-                , parents = d :: acc.parents
-                }
-                ({ acc | todo = rest } :: stack)
-
-
-type alias MapAccCtxBottomUp a b =
-    { todo : List (Tree a)
-    , done : List (Tree b)
-    , parents : List a
-    , label : a
-    }
+                _ ->
+                    Continue s
+        )
+        0
+        t
 
 
 {-| Sorts all children of each node based on the comparator function (the function recieves a list of ancestors).
@@ -941,50 +722,9 @@ type alias MapAccCtxBottomUp a b =
 
 -}
 sortWith : (List a -> Tree a -> Tree a -> Order) -> Tree a -> Tree a
-sortWith compareFn (Tree d cs) =
-    sortWithHelp compareFn
-        { todo = cs
-        , done = []
-        , label = d
-        , parents = [ d ]
-        }
-        []
-
-
-sortWithHelp : (List a -> Tree a -> Tree a -> Order) -> MapAccCtxBottomUp a a -> List (MapAccCtxBottomUp a a) -> Tree a
-sortWithHelp compareFn acc stack =
-    case acc.todo of
-        [] ->
-            let
-                children_ =
-                    List.sortWith (compareFn acc.parents) acc.done
-
-                node =
-                    Tree acc.label children_
-            in
-            case stack of
-                [] ->
-                    node
-
-                top :: rest ->
-                    sortWithHelp compareFn { top | done = node :: top.done } rest
-
-        (Tree d []) :: rest ->
-            sortWithHelp compareFn
-                { acc
-                    | todo = rest
-                    , done = Tree d [] :: acc.done
-                }
-                stack
-
-        (Tree d cs) :: rest ->
-            sortWithHelp compareFn
-                { todo = cs
-                , done = []
-                , label = d
-                , parents = d :: acc.parents
-                }
-                ({ acc | todo = rest } :: stack)
+sortWith compareFn t =
+    depthFirstTraversal defaultTopDown (\s a l c -> ( s, tree l (List.sortWith (compareFn a) c) )) () t
+        |> Tuple.second
 
 
 {-| Finds a subtree whose label matches the predicate.
@@ -1005,43 +745,16 @@ Searches the tree in a breadth-first manner.
 -}
 findBfs : (Tree a -> Bool) -> Tree a -> Maybe (Tree a)
 findBfs predicate t =
-    findBfsHelp predicate (Fifo [ t ] [])
-
-
-findBfsHelp : (Tree a -> Bool) -> Queue (Tree a) -> Maybe (Tree a)
-findBfsHelp predicate queue =
-    case removeQueue queue of
-        ( Just t, rest ) ->
-            if predicate t then
-                Just t
+    breadthFirstFold
+        (\s _ l c ->
+            if predicate (tree l c) then
+                Stop (Just (tree l c))
 
             else
-                findBfsHelp predicate (insertQueue t rest)
-
-        ( Nothing, _ ) ->
-            Nothing
-
-
-type Queue a
-    = Fifo (List a) (List a)
-
-
-insertQueue : Tree a -> Queue (Tree a) -> Queue (Tree a)
-insertQueue (Tree _ cs) (Fifo front back) =
-    Fifo front (List.reverse cs ++ back)
-
-
-removeQueue : Queue a -> ( Maybe a, Queue a )
-removeQueue fifo =
-    case fifo of
-        Fifo [] [] ->
-            ( Nothing, fifo )
-
-        Fifo [] back ->
-            removeQueue <| Fifo (List.reverse back) []
-
-        Fifo (next :: rest) back ->
-            ( Just next, Fifo rest back )
+                Continue s
+        )
+        Nothing
+        t
 
 
 {-| Restructure a `Tree` into another type of structure.
@@ -1117,63 +830,223 @@ Or perhaps you have your own tree datastructure and you want to convert to it:
 
 -}
 restructure : (a -> b) -> (b -> List c -> c) -> Tree a -> c
-restructure convertLabel convertTree (Tree l c) =
-    restructureHelp convertLabel
+restructure convertLabel convertTree t =
+    depthFirstTraversal (\s _ l c -> ( s, convertLabel l, c )) (\s _ l c -> ( s, convertTree l c )) () t
+        |> Tuple.second
+
+
+{-| The first callback to this function is used when going down the tree, effectively guiding and transforming it.
+The second callback is used to reconstruct the tree when going back up.
+
+This is similar like `restructure`, except this function:
+
+  - has an accumulator argument
+  - the first function can also change the children, not just the label (this allows us to implement `unfold` in terms of this function for instance)
+
+-}
+depthFirstTraversal :
+    (s -> List b -> a -> List (Tree a) -> ( s, b, List (Tree a) ))
+    -> (s -> List b -> b -> List c -> ( s, c ))
+    -> s
+    -> Tree a
+    -> ( s, c )
+depthFirstTraversal convertLabel convertTree s (Tree l c) =
+    let
+        ( state_, label_, children_ ) =
+            convertLabel s [] l c
+    in
+    depthFirstTraversalHelp convertLabel
         convertTree
-        { todo = c
-        , label = convertLabel l
+        state_
+        { todo = children_
+        , label = label_
         , done = []
         }
         []
 
 
-restructureHelp :
-    (a -> b)
-    -> (b -> List c -> c)
-    -> ReAcc a b c
-    -> List (ReAcc a b c)
-    -> c
-restructureHelp fLabel fTree acc stack =
+type alias Acc a b c =
+    { todo : List (Tree a)
+    , done : List c
+    , label : b
+    }
+
+
+defaultTopDown : s -> List a -> a -> List (Tree a) -> ( s, a, List (Tree a) )
+defaultTopDown s _ l c =
+    ( s, l, c )
+
+
+defaultBottomUp : s -> List b -> b -> List (Tree b) -> ( s, Tree b )
+defaultBottomUp s _ l c =
+    ( s, tree l c )
+
+
+depthFirstTraversalHelp :
+    (s -> List b -> a -> List (Tree a) -> ( s, b, List (Tree a) ))
+    -> (s -> List b -> b -> List c -> ( s, c ))
+    -> s
+    -> Acc a b c
+    -> List (Acc a b c)
+    -> ( s, c )
+depthFirstTraversalHelp fLabel fTree state acc stack =
     case acc.todo of
         [] ->
             let
-                node =
-                    fTree acc.label (List.reverse acc.done)
+                ( state_, node ) =
+                    fTree state (List.map .label stack) acc.label (List.reverse acc.done)
             in
             case stack of
                 [] ->
-                    node
+                    ( state_, node )
 
                 top :: rest ->
-                    restructureHelp
+                    depthFirstTraversalHelp
                         fLabel
                         fTree
+                        state_
                         { top | done = node :: top.done }
                         rest
 
-        (Tree l []) :: rest ->
-            restructureHelp
-                fLabel
-                fTree
-                { acc
-                    | todo = rest
-                    , done = fTree (fLabel l) [] :: acc.done
-                }
-                stack
+        (Tree l chs) :: rest ->
+            let
+                ancestors =
+                    acc.label :: List.map .label stack
 
-        (Tree l cs) :: rest ->
-            restructureHelp
-                fLabel
-                fTree
-                { todo = cs
-                , done = []
-                , label = fLabel l
-                }
-                ({ acc | todo = rest } :: stack)
+                ( state0, label_, children_ ) =
+                    fLabel state ancestors l chs
+            in
+            case children_ of
+                [] ->
+                    let
+                        ( state_, newTree ) =
+                            fTree state0 ancestors label_ []
+                    in
+                    depthFirstTraversalHelp
+                        fLabel
+                        fTree
+                        state_
+                        { acc
+                            | todo = rest
+                            , done = newTree :: acc.done
+                        }
+                        stack
+
+                cs ->
+                    depthFirstTraversalHelp
+                        fLabel
+                        fTree
+                        state0
+                        { todo = cs
+                        , done = []
+                        , label = label_
+                        }
+                        ({ acc | todo = rest } :: stack)
 
 
-type alias ReAcc a b c =
-    { todo : List (Tree a)
-    , done : List c
-    , label : b
-    }
+{-| Controls if the fold should continue traversing the tree, or should abort immediately.
+-}
+type Step a
+    = Continue a
+    | Stop a
+
+
+{-| Traverses the tree by "levels".
+
+     tree 1
+        [ tree 3
+            [ singleton 5
+            , singleton 4
+            ]
+        , singleton 2
+        , singleton 6
+        ]
+    |> breadthFirstFold (\s a l c -> Continue (l :: s)) []
+    --> [4, 5, 6, 2, 3, 1]
+
+(The list is reversed here due to how `::` works)
+
+-}
+breadthFirstFold :
+    (s -> List a -> a -> List (Tree a) -> Step s)
+    -> s
+    -> Tree a
+    -> s
+breadthFirstFold f acc t =
+    breadthFirstFoldHelp f acc [] [ t ] []
+
+
+breadthFirstFoldHelp : (s -> List a -> a -> List (Tree a) -> Step s) -> s -> List a -> List (Tree a) -> List ( List a, List (Tree a) ) -> s
+breadthFirstFoldHelp f acc parents trees nextSets =
+    case trees of
+        [] ->
+            case nextSets of
+                ( p, set ) :: sets ->
+                    breadthFirstFoldHelp f acc p set sets
+
+                [] ->
+                    acc
+
+        (Tree d ch) :: rest ->
+            case f acc parents d ch of
+                Continue a ->
+                    case ch of
+                        [] ->
+                            breadthFirstFoldHelp f a parents rest nextSets
+
+                        xs ->
+                            breadthFirstFoldHelp f a parents rest (( d :: parents, xs ) :: nextSets)
+
+                Stop a ->
+                    a
+
+
+{-| Traverses the tree by by going down as far as possible before trying
+out any siblings.
+
+     tree 1
+        [ tree 3
+            [ singleton 5
+            , singleton 4
+            ]
+        , singleton 2
+        , singleton 6
+        ]
+    |> depthFirstFold (\s a l c -> Continue (l :: s)) []
+    --> [6, 2, 4, 5, 3, 1]
+
+(The list is reversed here due to how `::` works)
+
+-}
+depthFirstFold :
+    (s -> List a -> a -> List (Tree a) -> Step s)
+    -> s
+    -> Tree a
+    -> s
+depthFirstFold f acc t =
+    depthFirstFoldHelp f acc [] [ t ] []
+
+
+depthFirstFoldHelp : (s -> List a -> a -> List (Tree a) -> Step s) -> s -> List a -> List (Tree a) -> List ( List a, List (Tree a) ) -> s
+depthFirstFoldHelp f acc parents trees nextSets =
+    case trees of
+        [] ->
+            case nextSets of
+                ( p, set ) :: sets ->
+                    depthFirstFoldHelp f acc p set sets
+
+                [] ->
+                    acc
+
+        (Tree d ch) :: rest ->
+            case f acc parents d ch of
+                Continue a ->
+                    case ch of
+                        [] ->
+                            depthFirstFoldHelp f a parents rest nextSets
+
+                        xs ->
+                            depthFirstFoldHelp f a (d :: parents) xs (( parents, rest ) :: nextSets)
+
+                Stop a ->
+                    a
diff --git a/tests/TreeTest.elm b/tests/TreeTest.elm
index 0829525..b72e504 100644
--- a/tests/TreeTest.elm
+++ b/tests/TreeTest.elm
@@ -1,4 +1,4 @@
-module TreeTest exposing (flattenTest, foldlTest, foldrTest, indexedMap, indexedMap2, lengthTest, map2, mapTest, unfold)
+module TreeTest exposing (depthFirstTraversalTest, flattenTest, foldlTest, foldrTest, indexedMap, indexedMap2, lengthTest, map2, mapTest, unfold)
 
 import Expect
 import Fuzz as F
@@ -297,3 +297,69 @@ indexedMap2 =
             in
             Tree.indexedMap2 (\n i s -> ( n, i, s )) left right
                 |> Expect.equal expected
+
+
+depthFirstTraversalTest : Test
+depthFirstTraversalTest =
+    describe "depthFirstTraversal"
+        [ test "accumulates state correctly" <|
+            \() ->
+                let
+                    t =
+                        Tree.tree 1
+                            [ Tree.singleton 2
+                            , Tree.singleton 3
+                            , Tree.tree 4
+                                [ Tree.singleton 5
+                                , Tree.singleton 6
+                                ]
+                            , Tree.singleton 7
+                            ]
+                in
+                t
+                    |> Tree.depthFirstTraversal (\( d, u ) _ l c -> ( ( d + 1, u ), ( l, d, 0 ), c )) (\( d, u ) _ ( l, p, _ ) c -> ( ( d, u + 1 ), Tree.tree ( l, p, u ) c )) ( 0, 0 )
+                    |> Expect.equal
+                        ( ( 7, 7 )
+                        , Tree.tree ( 1, 0, 6 )
+                            [ Tree.singleton ( 2, 1, 0 )
+                            , Tree.singleton ( 3, 2, 1 )
+                            , Tree.tree ( 4, 3, 4 )
+                                [ Tree.singleton ( 5, 4, 2 )
+                                , Tree.singleton ( 6, 5, 3 )
+                                ]
+                            , Tree.singleton ( 7, 6, 5 )
+                            ]
+                        )
+        , test "provides ancestors nicely" <|
+            \() ->
+                let
+                    tripleFirst : ( a, b, c ) -> a
+                    tripleFirst ( a, _, _ ) =
+                        a
+
+                    t =
+                        Tree.tree 1
+                            [ Tree.singleton 2
+                            , Tree.singleton 3
+                            , Tree.tree 4
+                                [ Tree.singleton 5
+                                , Tree.singleton 6
+                                ]
+                            , Tree.singleton 7
+                            ]
+                in
+                t
+                    |> Tree.depthFirstTraversal (\s a l c -> ( s, ( l, List.map tripleFirst a, [] ), c )) (\s a ( l, p, _ ) c -> ( s, Tree.tree ( l, p, List.map tripleFirst a ) c )) ()
+                    |> Expect.equal
+                        ( ()
+                        , Tree.tree ( 1, [], [] )
+                            [ Tree.singleton ( 2, [ 1 ], [ 1 ] )
+                            , Tree.singleton ( 3, [ 1 ], [ 1 ] )
+                            , Tree.tree ( 4, [ 1 ], [ 1 ] )
+                                [ Tree.singleton ( 5, [ 4, 1 ], [ 4, 1 ] )
+                                , Tree.singleton ( 6, [ 4, 1 ], [ 4, 1 ] )
+                                ]
+                            , Tree.singleton ( 7, [ 1 ], [ 1 ] )
+                            ]
+                        )
+        ]

From f5a80b2f194fe531ef2de50b3fd593f5bb71884f Mon Sep 17 00:00:00 2001
From: Jakub Hampl <kopomir@gmail.com>
Date: Tue, 18 Jul 2023 10:50:22 +0300
Subject: [PATCH 3/3] Removes some forgotten comments

---
 src/Tree.elm | 32 --------------------------------
 1 file changed, 32 deletions(-)

diff --git a/src/Tree.elm b/src/Tree.elm
index 46afad1..041ff09 100644
--- a/src/Tree.elm
+++ b/src/Tree.elm
@@ -311,22 +311,6 @@ foldr f acc t =
     List.foldl f acc <| foldl (::) [] t
 
 
-
--- foldlHelp : (a -> b -> b) -> b -> List (Tree a) -> List (List (Tree a)) -> b
--- foldlHelp f acc trees nextSets =
---     case trees of
---         [] ->
---             case nextSets of
---                 set :: sets ->
---                     foldlHelp f acc set sets
---                 [] ->
---                     acc
---         (Tree d []) :: rest ->
---             foldlHelp f (f d acc) rest nextSets
---         (Tree d xs) :: rest ->
---             foldlHelp f (f d acc) xs (rest :: nextSets)
-
-
 {-| Flattens the tree into a list. This is equivalent to `foldr (::) []`
 -}
 flatten : Tree a -> List a
@@ -399,22 +383,6 @@ links t =
         |> List.reverse
 
 
-
--- linksHelp : a -> List ( a, a ) -> List (Tree a) -> List ( a, List (Tree a) ) -> List ( a, a )
--- linksHelp parent soFar trees nextSets =
---     case trees of
---         [] ->
---             case nextSets of
---                 ( newParent, set ) :: sets ->
---                     linksHelp newParent soFar set sets
---                 [] ->
---                     List.reverse soFar
---         (Tree d []) :: rest ->
---             linksHelp parent (( parent, d ) :: soFar) rest nextSets
---         (Tree d xs) :: rest ->
---             linksHelp parent (( parent, d ) :: soFar) rest (( d, xs ) :: nextSets)
-
-
 {-| Create a tree from a seed.
 
 Running the function on the seed should return a label and a list of seeds to