Files
loop/src/Dodge/Layout/Tree.hs
T
2021-04-14 12:00:12 +02:00

50 lines
1.9 KiB
Haskell

{-|
Combining and composing trees.
-}
module Dodge.Layout.Tree
where
import Data.Tree
import Control.Monad.State
import System.Random
-- | 'Left' elements get new children, 'Right' elements inherit the children from the
-- mapped over node
expandTreeBy :: (a -> Tree (Either b b)) -> Tree a -> Tree b
expandTreeBy f (Node x []) = fmap removeEither (f x)
where removeEither (Left y) = y
removeEither (Right y) = y
expandTreeBy f (Node x xs) = appendAndRemove $ f x
where appendAndRemove (Node (Left y) ys) = Node y (map appendAndRemove ys)
appendAndRemove (Node (Right y) _ ) = Node y (map (expandTreeBy f) xs)
expandTreeRand :: RandomGen g =>
(a -> State g (Tree (Either b b))) -> Tree a -> State g (Tree b)
expandTreeRand f (Node x []) = fmap (fmap removeEither) (f x)
where removeEither (Left y) = y
removeEither (Right y) = y
expandTreeRand f (Node x xs) = do
root <- f x
branches <- sequence $ map (expandTreeRand f) xs
return (appendAndRemove branches root)
where appendAndRemove :: [Tree a] -> Tree (Either a a) -> Tree a
appendAndRemove bran (Node (Left y) ys) = Node y (map (appendAndRemove bran) ys)
appendAndRemove bran (Node (Right y) _) = Node y bran
collapseTree :: Tree (Tree (Either a a)) -> Tree (Either a a)
collapseTree = g . expandTreeBy f
where f (Node (Right x) xs) = Node (Right (Right x)) $ map f xs
f (Node (Left x) xs) = Node (Left (Left x)) $ map f xs
g (Node (Right x) []) = Node (Right x) []
g (Node (Right x) xs) = Node (Left x) $ map g xs
g (Node y ys) = Node y $ map g ys
treePost :: [a] -> a -> Tree a
treePost [] y = Node y []
treePost (x:xs) y = Node x [treePost xs y]
treeTrunk :: [a] -> Tree a -> Tree a
treeTrunk [] t = t
treeTrunk (x:xs) t = Node x [treeTrunk xs t]
applyToRoot :: (a -> a) -> Tree a -> Tree a
applyToRoot f (Node x xs) = Node (f x) xs