112 lines
2.8 KiB
Haskell
112 lines
2.8 KiB
Haskell
module Dodge.Layout.Tree.Polymorphic
|
|
where
|
|
import Dodge.RandomHelp
|
|
|
|
import Data.Tree
|
|
import Control.Monad.State
|
|
import System.Random
|
|
|
|
{-
|
|
Creates a linear tree.
|
|
Safe.
|
|
-}
|
|
treePost :: [a] -> a -> Tree a
|
|
treePost [] y = Node y []
|
|
treePost (x:xs) y = Node x [treePost xs y]
|
|
|
|
{-
|
|
Creates a tree with one trunk branch,
|
|
input as a list, that ends in another tree.
|
|
-}
|
|
treeTrunk
|
|
:: [a] -- ^ The trunk
|
|
-> Tree a -- ^ The end of the tree
|
|
-> Tree a
|
|
treeTrunk [] t = t
|
|
treeTrunk (x:xs) t = Node x [treeTrunk xs t]
|
|
|
|
{-
|
|
Applies a function to the root of a tree.
|
|
-}
|
|
applyToRoot :: (a -> a) -> Tree a -> Tree a
|
|
applyToRoot f (Node t ts) = Node (f t) ts
|
|
|
|
treeSize = length . flatten
|
|
|
|
{-
|
|
Applies a function to a specific node determined by a list of indices.
|
|
Unsafe (partial function).
|
|
-}
|
|
applyToNode :: [Int] -> (a -> a) -> Tree a -> Tree a
|
|
applyToNode [] f t = applyToRoot f t
|
|
applyToNode (i:is) f (Node x xs) = Node x (ys ++ [applyToNode is f z] ++ zs)
|
|
where
|
|
(ys, z:zs) = splitAt i xs
|
|
|
|
zipTree :: Tree a -> Tree b -> Tree (a,b)
|
|
zipTree (Node x xs) (Node y ys) = Node (x,y) $ zipWith zipTree xs ys
|
|
|
|
{-
|
|
Makes each node into its child number, i.e. the index it has
|
|
in the list of children of its parent.
|
|
-}
|
|
treeChildNums :: Tree a -> Tree Int
|
|
treeChildNums t = setRoot 0 t
|
|
where
|
|
setRoot :: Int -> Tree a -> Tree Int
|
|
setRoot i (Node x xs) = Node i (zipWith setRoot [0..] xs)
|
|
|
|
{-
|
|
Makes each node into its path, i.e. the list of indices that,
|
|
when followed from the root, lead to the node.
|
|
-}
|
|
treePaths :: Tree a -> Tree [a]
|
|
treePaths (Node x xs) = fmap (x :) $ Node [] (map treePaths xs)
|
|
|
|
{-
|
|
Picks a random path in the tree.
|
|
Uniform probability that the path leads to any specific node.
|
|
-}
|
|
randomPath :: RandomGen g => Tree a -> State g [Int]
|
|
randomPath = takeOne . flatten . treePaths . treeChildNums
|
|
|
|
{-
|
|
Apply a function to a node picked uniformly at random.
|
|
-}
|
|
applyToRandomNode :: RandomGen g => (a -> a) -> Tree a -> State g (Tree a)
|
|
applyToRandomNode f t = do
|
|
p <- randomPath t
|
|
return $ applyToNode p f t
|
|
|
|
{-
|
|
Add a forest to the end of a tree.
|
|
-}
|
|
addToTrunk :: Tree a -> [Tree a] -> Tree a
|
|
addToTrunk (Node x []) f = Node x f
|
|
addToTrunk (Node x (t:ts)) f = Node x (addToTrunk t f : ts)
|
|
|
|
{-
|
|
Find the depth of a tree along the trunk.
|
|
-}
|
|
trunkDepth :: Tree a -> Int
|
|
trunkDepth (Node _ []) = 0
|
|
trunkDepth (Node _ (x:xs)) = trunkDepth x + 1
|
|
|
|
{-
|
|
Split a tree at a given point along its trunk.
|
|
-}
|
|
splitTrunkAt :: Int -> Tree a -> (Tree a, [Tree a])
|
|
splitTrunkAt 0 (Node x xs) = (Node x [],xs)
|
|
splitTrunkAt i (Node y (x:xs)) =
|
|
let (t, ts) = (splitTrunkAt (i-1) x)
|
|
in (Node y (t : xs) , ts)
|
|
|
|
{-
|
|
Split a tree at a random point along its trunk.
|
|
-}
|
|
splitTrunk :: RandomGen g => Tree a -> State g (Tree a, [Tree a])
|
|
splitTrunk t = do
|
|
i <- state $ randomR (0, trunkDepth t)
|
|
return $ splitTrunkAt i t
|
|
|