More compact buildTree impl using foldr

src/Data/Graph/DirectedAcyclic/View/Tree.hs

@@ -106,12 +106,20 @@ mkGraph nodeMap edges =
 keySet :: HashMap k v -> HashSet k
 keySet = S.fromMap . M.map (const ())
 
+-- | Traverse a graph DFS-style and build a tree recording the traversal.
+--
+-- The code looks like a simple fold, because the edge labels are the ones
+-- responsible for limiting the recursion into a tree structure.
+--
+-- The graph should have at most one full out-edge per node, and\/or have no
+-- cycles, otherwise this function isn't guaranteed to stop.
 buildTree
     :: (Eq n, Hashable n)
     => [(n, Maybe b)]
     -> Graph n a (Maybe b)
     -> [DagViewTree a (a, b)]
-buildTree nodes graph = go nodes
+buildTree nodes graph = -- go nodes
+    {-
     where
     go []               = []
     go ((n, full) : ps) =
@@ -126,6 +134,18 @@ buildTree nodes graph = go nodes
                     Just info ->
                         let ts = go ps
                         in  LinkNode (fst c, info) : ts
+    -}
+
+    let f (n, full) ts =
+            case M.lookup n graph of
+                Nothing -> ts
+                Just c  ->
+                    let t = case full of
+                                Nothing   -> FullNode (fst c) (go $ snd c)
+                                Just info -> LinkNode (fst c, info)
+                    in  t : ts
+        go = foldr f []
+    in  go nodes
 
 dagViewTree
     :: (Eq n, Ord n, Hashable n)