From d12b2178de447a0a12b26730232a821bcf763725 Mon Sep 17 00:00:00 2001
From: Neil Brown <neil@twistedsquare.com>
Date: Thu, 6 Nov 2008 18:32:35 +0000
Subject: [PATCH] Added an augmented version of Adam's gmapMFor that uses the
 same generics techniques, but also gives a route to the node to the
 transformation function

---
 common/GenericUtils.hs | 30 ++++++++++++++++++++++++++++++
 common/TreeUtils.hs    |  1 +
 2 files changed, 31 insertions(+)

diff --git a/common/GenericUtils.hs b/common/GenericUtils.hs
index 9a4257d..c95be11 100644
--- a/common/GenericUtils.hs
+++ b/common/GenericUtils.hs
@@ -27,8 +27,11 @@ module GenericUtils (
   , TypeSet, makeTypeSet
   , containsTypes
   , gmapMFor
+  , gmapMForRoute
   ) where
 
+import Control.Monad.Identity()
+import Control.Monad.State
 import Data.Generics
 import Data.IntMap (IntMap)
 import qualified Data.IntMap as IntMap
@@ -39,6 +42,7 @@ import Data.Typeable
 import System.IO.Unsafe
 
 import qualified AST as A
+import TreeUtils
 import Utils
 
 -- | A type identifier.
@@ -153,3 +157,29 @@ gmapMFor typeset f = gmapM (each f)
             Just Through -> gmapM (each f) x
             Just Miss -> return x
             Nothing -> return x
+
+gmapMForRoute :: forall m t. (Monad m, Data t) =>
+  TypeSet ->
+  (forall s. Data s => (s, (forall n. Monad n => (s -> n s) -> (t -> n t))) -> m s)
+  -> (t -> m t)
+gmapMForRoute typeset f = gmapMWithRoute (each f)
+  where
+    each :: Data u => (forall s. Data s => (s, (forall n. Monad n => (s -> n s) -> (t -> n t))) -> m s)
+             -> ((u, (forall n. Monad n => (u -> n u) -> (t -> n t))) -> m u)
+    each f (x, route)
+        = case IntMap.lookup (typeKey x) typeset of
+            Just Hit -> f (x, route)
+            Just Through -> gmapMWithRoute (\(y, route') -> each f (y, route . route')) x
+            Just Miss -> return x
+            Nothing -> return x
+
+gmapMWithRoute :: forall a m. (Monad m, Data a) => (forall b. Data b => (b, (forall n. Monad n => (b -> n b) ->
+  (a -> n a))) -> m b) -> a -> m a
+gmapMWithRoute f = gmapFuncs [GM {unGM = f' n} | n <- [0..]]
+  where
+    f' :: Int -> (forall b. Data b => b -> m b)
+    f' n x = f (x, makeRoute n)
+
+-- Given a number, makes a route function for that child:
+makeRoute :: (Data s, Data t, Monad m) => Int -> (s -> m s) -> (t -> m t)
+makeRoute target f = gmapFuncs [mkM' (if n == target then f else return) | n <- [0..]]
diff --git a/common/TreeUtils.hs b/common/TreeUtils.hs
index d3aa646..9cda044 100644
--- a/common/TreeUtils.hs
+++ b/common/TreeUtils.hs
@@ -26,6 +26,7 @@ module TreeUtils (
   checkTreeForConstr,
   con0, con1, con2, con3, con4, con5, con6, con7
   , decomp1, decomp2, decomp3, decomp4, decomp5
+  , mkM', gmapFuncs
   ) where
 
 import Control.Monad.State