Smarter tree traversals: "Scrap Your Uniplate".

This provides gmapMFor and gmapMFor2, which are like gmapM, but know what
they're looking for, and can therefore avoid going down branches of the tree
that won't contain any interesting types.

The basic approach is quite similar to Uniplate's PlateData: there's a function
(containsType) that'll tell you whether one type is contained somewhere within
another. However, unlike Uniplate, we build a static IntMap IntSet of the types
we need to know about, which allows rather quicker lookups. (I did try using
PlateData first.)

The result is that applyDepthM is now much quicker than it was before.
applyDepthM2 is a bit less impressive, which I assume is because it can't
really prune the tree much if it's looking for two types.

Future enhancements:
- convert more passes to use applyDepthM*;
- make gmapMFor* aware of constructors rather than just types, which should
  allow a bit more pruning.

Makefile.am

@@ -120,6 +120,7 @@ tock_SOURCES_hs += checks/UsageCheckUtils.hs
 tock_SOURCES_hs += common/Errors.hs
 tock_SOURCES_hs += common/EvalConstants.hs
 tock_SOURCES_hs += common/EvalLiterals.hs
+tock_SOURCES_hs += common/GenericUtils.hs
 tock_SOURCES_hs += common/Intrinsics.hs
 tock_SOURCES_hs += common/Pattern.hs
 tock_SOURCES_hs += common/PrettyShow.hs

common/CommonTest.hs

@@ -25,6 +25,7 @@ import Data.Generics
 import Test.HUnit hiding (State)
 
 import qualified AST as A
+import GenericUtils
 import Metadata
 import TreeUtils
 import Types
@@ -121,12 +122,37 @@ testDecomp = TestList
    doTest :: (Eq a, Show a) => Int -> Maybe a -> Maybe a -> Test
    doTest n exp act = TestCase $ assertEqual ("testDecomp " ++ show n) exp act
 
+data NotPartOfAST = NotPartOfAST Int
+  deriving (Show, Typeable, Data)
+
+-- | Test 'typeContains'.
+testTypeContains :: Test
+testTypeContains = TestList
+    [
+    -- AST elements that it should know about
+      test   0 True  (undefined :: A.AST) (undefined :: Meta)
+    , test   1 False (undefined :: Meta) (undefined :: A.AST)
+    , test   2 True  (undefined :: A.Process) (undefined :: A.Name)
+    , test   3 False (undefined :: A.AbbrevMode) (undefined :: A.Formal)
+    , test   4 False (undefined :: String) (undefined :: Meta)
+    , test   5 True  (undefined :: String) (undefined :: Char)
+
+    -- Things outside the AST
+    , test 100 True  (NotPartOfAST 42) (undefined :: String)
+    , test 101 False (undefined :: String) (NotPartOfAST 42)
+    ]
+  where
+    test :: (Show a, Data a, Typeable a, Show b, Data b, Typeable b) =>
+            Int -> Bool -> a -> b -> Test
+    test n exp start find
+        = TestCase $ assertEqual ("testTypeContains " ++ show n)
+                                 exp (typeContains start find)
 
 --Returns the list of tests:
 tests :: Test
 tests = TestLabel "CommonTest" $ TestList
- [
+    [ testIsSafeConversion
-   testIsSafeConversion
     , testCheckTreeForConstr
     , testDecomp
+    , testTypeContains
     ]

common/GenericUtils.hs

@@ -0,0 +1,121 @@
+{-
+Tock: a compiler for parallel languages
+Copyright (C) 2008  University of Kent
+
+This program is free software; you can redistribute it and/or modify it
+under the terms of the GNU General Public License as published by the
+Free Software Foundation, either version 2 of the License, or (at your
+option) any later version.
+
+This program is distributed in the hope that it will be useful, but
+WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+General Public License for more details.
+
+You should have received a copy of the GNU General Public License along
+with this program.  If not, see <http://www.gnu.org/licenses/>.
+-}
+
+-- | Utilities for generic operations.
+--
+-- This code was inspired by Neil Mitchell's Uniplate library.
+-- 'typeContains' is faster than PlateData's equivalent at the cost of some
+-- flexibility: it'll only work for types that it knows about (which can be
+-- added to in the definition of 'contains').
+module GenericUtils (typeContains, gmapMFor, gmapMFor2) where
+
+import Data.Generics
+import Data.IntMap (IntMap)
+import qualified Data.IntMap as IntMap
+import Data.IntSet (IntSet)
+import qualified Data.IntSet as IntSet
+import Data.List
+import Data.Typeable
+import System.IO.Unsafe
+
+import qualified AST as A
+
+data DataBox = forall a. (Typeable a, Data a) => DataBox a
+
+-- | Given a witness for a type, return witnesses for all the types that its
+-- constructors take.
+constrArgTypes :: (Data a, Typeable a) => a -> [DataBox]
+constrArgTypes x = if isAlgType dtype then concatMap f constrs else []
+  where
+    f constr = gmapQ DataBox (asTypeOf (fromConstr constr) x)
+    constrs = dataTypeConstrs dtype
+    dtype = dataTypeOf x
+
+-- | Given a witness for a type, return its type key.
+typeKey :: Typeable a => a -> Int
+typeKey x = unsafePerformIO $ typeRepKey $ typeOf x
+
+-- | Given a witness for a type, return a map from type keys to witnesses for
+-- all the types it contains recursively.
+containsTypes :: (Data a, Typeable a) => a -> IntMap DataBox
+containsTypes start = containsTypes' (DataBox start) IntMap.empty
+  where
+    containsTypes' :: DataBox -> IntMap DataBox -> IntMap DataBox
+    containsTypes' box@(DataBox thisType) seen
+        = if thisKey `IntMap.member` seen
+            then seen
+            else foldl (\s t -> containsTypes' t s)
+                       (IntMap.insert thisKey box seen)
+                       (constrArgTypes thisType)
+      where
+        thisKey = typeKey thisType
+
+-- | A map from type keys to the other type keys reachable from them.
+type ContainsMap = IntMap IntSet
+
+-- | A map of reachable types.
+-- At the moment this only knows about types reachable from the AST.
+contains :: ContainsMap
+contains = IntMap.fromList [(typeKey t,
+                             IntSet.fromList $ IntMap.keys $ containsTypes t)
+                            | DataBox t <- IntMap.elems allTypes]
+  where
+    allTypes = containsTypes (undefined :: A.AST)
+
+-- | Does one type contain another?
+-- (A type always contains itself.)
+typeContains :: (Data a, Typeable a, Data b, Typeable b) => a -> b -> Bool
+typeContains start find
+    = if startKey == findKey
+        then True
+        else case IntMap.lookup startKey contains of
+               Just set -> findKey `IntSet.member` set
+               Nothing -> True  -- can't tell, so it might be
+  where
+    startKey = typeKey start
+    findKey = typeKey find
+
+-- | Type-smart generic mapM.
+-- This is like 'gmapM', but it only applies the function to arguments that
+-- could contain the target type.
+gmapMFor :: (Monad m, Data t, Data a) =>
+            a                                 -- ^ Witness for target type
+            -> (forall s. Data s => s -> m s) -- ^ Function to apply
+            -> (t -> m t)                     -- ^ Generic operation
+gmapMFor find top = gmapM (each find top)
+  where
+    each :: (Monad m, Data t, Data a) =>
+            a -> (forall s. Data s => s -> m s) -> (t -> m t)
+    each find top x
+        = if cont then top x else return x
+      where cont = x `typeContains` find
+
+-- | Two-type version of 'gmapMFor'.
+gmapMFor2 :: (Monad m, Data t, Data a1, Data a2) =>
+             a1                                -- ^ Witness for target type 1
+             -> a2                             -- ^ Witness for target type 2
+             -> (forall s. Data s => s -> m s) -- ^ Function to apply
+             -> (t -> m t)                     -- ^ Generic operation
+gmapMFor2 find1 find2 top = gmapM (each find1 find2 top)
+  where
+    each :: (Monad m, Data t, Data a1, Data a2) =>
+            a1 -> a2 -> (forall s. Data s => s -> m s) -> (t -> m t)
+    each find1 find2 top x
+        = if cont then top x else return x
+      where cont = x `typeContains` find1 || x `typeContains` find2
+

pass/Pass.hs

@@ -28,6 +28,7 @@ import System.IO
 import qualified AST as A
 import CompState
 import Errors
+import GenericUtils
 import Metadata
 import PrettyShow
 import TreeUtils
@@ -194,27 +195,26 @@ makeGeneric top
         `extM` (return :: String -> m String)
         `extM` (return :: Meta -> m Meta)
 
--- | Apply a monadic operation everywhere that it matches in the AST, going
+-- | Apply a monadic operation everywhere that it matches, going depth-first.
--- depth-first.
+applyDepthM :: forall a t. (Data a, Data t) => (a -> PassM a) -> t -> PassM t
-applyDepthM :: (Data a, Data t) => (a -> PassM a) -> t -> PassM t
 applyDepthM f = doGeneric `extM` (doSpecific f)
   where
-    doGeneric :: Data t => t -> PassM t
+    doGeneric :: Data t1 => t1 -> PassM t1
-    doGeneric = makeGeneric (applyDepthM f)
+    doGeneric = gmapMFor (undefined :: a) (applyDepthM f)
 
-    doSpecific :: Data t => (t -> PassM t) -> t -> PassM t
+    doSpecific :: Data t2 => (t2 -> PassM t2) -> t2 -> PassM t2
     doSpecific f x = (doGeneric x >>= f)
 
 -- | Apply two monadic operations everywhere they match in the AST, going
 -- depth-first.
-applyDepthM2 :: (Data a, Data b, Data t) =>
+applyDepthM2 :: forall a b t. (Data a, Data b, Data t) =>
                   (a -> PassM a) -> (b -> PassM b) -> t -> PassM t
 applyDepthM2 f1 f2 = doGeneric `extM` (doSpecific f1) `extM` (doSpecific f2)
   where
-    doGeneric :: Data t => t -> PassM t
+    doGeneric :: Data t1 => t1 -> PassM t1
-    doGeneric = makeGeneric (applyDepthM2 f1 f2)
+    doGeneric = gmapMFor2 (undefined :: a) (undefined :: b) (applyDepthM2 f1 f2)
 
-    doSpecific :: Data t => (t -> PassM t) -> t -> PassM t
+    doSpecific :: Data t2 => (t2 -> PassM t2) -> t2 -> PassM t2
     doSpecific f x = (doGeneric x >>= f)
 
 excludeConstr :: (Data a, CSMR m) => [Constr] -> a -> m a
@@ -225,3 +225,4 @@ excludeConstr cons x
 
 mk1M :: (Monad m, Data a, Typeable1 t) => (forall d . Data d => t d -> m (t d)) -> a -> m a
 mk1M = ext1M return
+