Add tree structure generation

src/Dodge/Layout/Tree/Polymorphic.hs

@@ -0,0 +1,111 @@
+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 x xs) = Node (f x) xs
+
+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
+