Allow for random composition of trees

src/TreeHelp.hs

@@ -0,0 +1,170 @@
+{-
+Helpers for the manipulation of rose trees.
+Throughout, the _trunk_ refers to successive first children in the tree.
+For example, in the tree 
+
+>   Node a [ Node b [], Node c [Node d []] ]
+
+the nodes in the trunk are [a,b].
+ -}
+module TreeHelp
+    ( module Data.Tree
+    , module Data.Tree.Lens
+    , applyToSubtree
+    , applyToSubforest
+    , treeFromPost
+    , treeFromTrunk
+    , splitTrunk
+    , applyToRandomNode
+    , addToTrunk
+    , inorderNumberTree
+    , updateSingleNodes
+    , updateRandNode
+    , safeUpdateSingleNode
+    ) where
+import Dodge.RandomHelp
+
+import Data.Maybe
+import Data.Tree.Lens
+import Data.Tree
+import Control.Monad.State
+import System.Random
+import Control.Lens
+
+{- | Creates a linear tree.
+Safe.  -}
+treeFromPost :: [a] -> a -> Tree a
+treeFromPost xs = treeFromTrunk xs . pure
+
+{- | Creates a tree with one trunk branch, 
+input as a list, that ends in another tree.  -}
+treeFromTrunk 
+    :: [a]      -- ^ The trunk
+    -> Tree a   -- ^ The end of the tree
+    -> Tree a
+treeFromTrunk = flip $ foldr f
+  where
+    f x t = Node x [t]
+
+-- find use for?
+---- | Consider defining this using generalised recursion patterns
+--treeSize :: Tree a -> Int
+--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 is = applyToSubtree is . over root
+
+{- | Applies a function to a specific subtree determined by a list of indices.
+Unsafe (partial function).  -}
+applyToSubtree :: [Int] -> (Tree a -> Tree a) -> Tree a -> Tree a
+applyToSubtree [] f t = f t
+applyToSubtree (i:is) f (Node x xs) = Node x (xs & ix i %~ applyToSubtree is f)
+
+{- | Applies a function to a specific subforest determined by a list of indices.
+Unsafe (partial function).  -}
+applyToSubforest :: [Int] -> ([Tree a] -> [Tree a]) -> Tree a -> Tree a
+applyToSubforest is = applyToSubtree is . over branches
+--applyToSubforest [] f (Node p cs) = Node p (f cs)
+--applyToSubforest (i:is) f (Node x xs) = Node x (ys ++ [applyToSubforest is f z] ++ zs)
+--  where
+--    (ys, z:zs) = splitAt i xs
+
+-- do not delete: find use for
+--{- |
+--Applies a function to the first node along a trunk that satisfies a given property.
+---}
+--applyToSubTrunkBy :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
+--applyToSubTrunkBy cond f (Node x (t:ts))
+--    | cond x    = f (Node x (t:ts))
+--    | otherwise = Node x (applyToSubTrunkBy cond f t : ts)
+--applyToSubTrunkBy _ _ t = t
+
+-- gives the list of all updates to a single node
+-- there must be a better way of doing something like this
+updateSingleNodes :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> [Tree a]
+updateSingleNodes f update t@(Node x ts)
+    | f x       = update t : updateChildren
+    | otherwise = updateChildren
+  where
+    updateChildren = map (Node x) (subMap (updateSingleNodes f update) ts)
+
+updateRandNode :: RandomGen g => (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> State g (Tree a)
+updateRandNode t f = takeOne . updateSingleNodes t f
+
+
+safeUpdateSingleNode :: (a -> Bool) -> (Tree a -> Tree a) -> Tree a -> Tree a
+safeUpdateSingleNode f g t = fromMaybe t $ listToMaybe $ updateSingleNodes f g t
+
+subMap :: (a -> [a]) -> [a] -> [[a]]
+subMap f (x:xs) = (f x <&> (: xs)) ++ ( (x :) <$> subMap f xs ) 
+subMap _ [] = []
+
+-- find use for?
+--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 = setRoot 0
+  where 
+    setRoot :: Int -> Tree a -> Tree Int
+    setRoot i (Node _ 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) = (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 the value of a node; 
+the node is 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 (along the trunk).  -}
+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:_)) = trunkDepth x + 1
+
+{- | Split a tree at a given point along its trunk.  -}
+splitTrunkAt 
+    :: Int      -- ^ Split depth
+    -> 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)
+splitTrunkAt _ (Node _ []) = error "Trying to split to short a trunk"
+
+{- | 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
+
+-- untested
+inorderNumberTree :: Tree a -> Tree (a,Int)
+inorderNumberTree = fst . f 0
+  where
+    f i (Node x ts) =
+        let (ts',i') = g (i+1) ts
+        in (Node (x,i) ts', i')
+    g i (t:ts) = 
+        let (t',i') = f i t
+            (ts',i'') = g i' ts
+        in (t': ts', i'')
+    g i [] = ([], i)