From d8174c7ccc60cd2859e5da684ef636050502770f Mon Sep 17 00:00:00 2001
From: justin <justin.g.kruger@gmail.com>
Date: Thu, 9 Jun 2022 18:52:21 +0100
Subject: [PATCH] Add functions for composing trees

---
 src/Dodge/Tree/Compose.hs | 18 ++++++++++++++++++
 src/TreeHelp.hs           | 16 +++++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/Dodge/Tree/Compose.hs b/src/Dodge/Tree/Compose.hs
index 1dd162665..5a532770f 100644
--- a/src/Dodge/Tree/Compose.hs
+++ b/src/Dodge/Tree/Compose.hs
@@ -11,6 +11,7 @@ module Dodge.Tree.Compose
     , overwriteLabel
     , applyToCompRoot
     , composeTree
+    , composeTree'
     , composeTreeRand
 --    , compTree
     ) where
@@ -19,6 +20,7 @@ import TreeHelp
 import Dodge.Base
 import LensHelp
 
+import Data.Bifunctor
 import Control.Monad.State
 import Data.Foldable
 import System.Random
@@ -84,6 +86,22 @@ newfRand t = updateRandNode (isJust . upf . snd) ( (root . _2 %~ g) . (root . _1
 composeTree :: Tree (a -> Maybe a, Tree a) -> Tree a
 composeTree (Node (_,t) ts) = attachList t ts
 
+composeTree' :: Monoid b => Tree (a -> Maybe (b, a), Tree a) -> (b, Tree a)
+composeTree' (Node (_,t) ts) = attachList' t ts
+
+attachList' :: Monoid b => Tree a -> [Tree (a -> Maybe (b, a), Tree a)] -> (b,Tree a)
+attachList' t xs = shiftChildren <$> foldrM newf' (fmap ([],) t) xs
+
+newf' :: Monoid b => Tree (a -> Maybe (b, a),Tree a) -> Tree ([Tree a],a) -> (b , Tree ([Tree a],a))
+newf' t = first (s<>) . msafeUpdateSingleNode (isJust . upf . snd) upt
+  where
+    (s,t') = composeTree' t
+    upf = t ^?! root . _1
+    upt = (root . _2 %%~ g) . (root . _1 .:~ t')
+    g x = fromMaybe (mempty,x) (upf x)
+--    mtf = undefined
+--     ( (root . _2 %~ g) . (root . _1 .:~ composeTree t ) )
+
 attachList :: Tree a -> [Tree (a -> Maybe a, Tree a)] -> Tree a
 attachList t xs = shiftChildren $ foldr newf (fmap ([],) t) xs
 
diff --git a/src/TreeHelp.hs b/src/TreeHelp.hs
index c13087b94..0167d8e35 100644
--- a/src/TreeHelp.hs
+++ b/src/TreeHelp.hs
@@ -21,6 +21,7 @@ module TreeHelp
     , updateSingleNodes
     , updateRandNode
     , safeUpdateSingleNode
+    , msafeUpdateSingleNode
     ) where
 import RandomHelp
 
@@ -88,13 +89,26 @@ updateSingleNodes f update t@(Node x ts)
   where
     updateChildren = map (Node x) (subMap (updateSingleNodes f update) ts)
 
+mupdateSingleNodes :: Monad m => (a -> Bool) -> (Tree a -> m (Tree a)) -> Tree a -> [m (Tree a)]
+mupdateSingleNodes f update t@(Node x ts)
+    | f x       = update t : updateChildren
+    | otherwise = updateChildren
+  where
+    updateChildren = map (fmap $ Node x) (msubMap (mupdateSingleNodes 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
 
+msafeUpdateSingleNode :: Monoid b => (a -> Bool) -> (Tree a -> (b, Tree a)) -> Tree a -> (b,Tree a)
+msafeUpdateSingleNode f g t = fromMaybe (mempty,t) $ listToMaybe $ mupdateSingleNodes f g t
+
+msubMap :: Functor m => (a -> [m a]) -> [a] -> [m [a]]
+msubMap f (x:xs) = (f x <&> fmap (: xs))  ++ ( (fmap (x :)) <$> msubMap f xs ) 
+msubMap _ [] = []
+
 subMap :: (a -> [a]) -> [a] -> [[a]]
 subMap f (x:xs) = (f x <&> (: xs)) ++ ( (x :) <$> subMap f xs ) 
 subMap _ [] = []