NodeDepthTree.hs

#!/usr/bin/env stack
-- stack runhaskell

{-# OPTIONS_GHC -Wall -Wno-orphans #-}

import Data.Bifunctor
import Data.Foldable
import Data.List.NonEmpty (NonEmpty ((:|)))
import Data.Tree
import Test.QuickCheck

type NodeDepthTree a = NonEmpty (Int, a)

treeToNdt :: Tree a -> NodeDepthTree a
treeToNdt =
  foldTree
    $ \a subforest -> (0, a) :| map (first succ) (foldMap toList subforest)

ndtToTree :: NodeDepthTree a -> Tree a
ndtToTree ndt = case takeTree ndt of
  (tree, []) -> tree
  (_, _) -> error "invalid tree"

takeForest :: Int -> [(Int, a)] -> (Forest a, [(Int, a)])
takeForest parentDepth ndt = case ndt of
  node@(depth, _) : nodes
    | depth > parentDepth ->
      let (child, rest) = takeTree $ node :| nodes
       in case rest of
            [] -> ([child], [])
            _ : _ ->
              let (subforest, rest') = takeForest parentDepth rest
               in (child : subforest, rest')
  _ ->
    ([], ndt)

takeTree :: NonEmpty (Int, a) -> (Tree a, [(Int, a)])
takeTree ((depth, node) :| nodes) = case nodes of
  [] -> (Node node [], [])
  _ : _ ->
    let (subforest, rest) = takeForest depth nodes
     in (Node node subforest, rest)

--------------------------------------------------------------------------------

instance Arbitrary a => Arbitrary (Tree a) where
  arbitrary = do
    s <- getSize
    Node <$> arbitrary <*> resize (s `div` 2) arbitrary

main :: IO ()
main = do
  result <-
    quickCheckResult
      $ \t -> ndtToTree (treeToNdt t) === (t :: Tree Char)
  case result of
    Success {} -> pure ()
    _ -> fail $ show result