eb29e65bad
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.
122 lines
4.7 KiB
Haskell
122 lines
4.7 KiB
Haskell
{-
|
|
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
|
|
|