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,
|
|
IntMap.keysSet $ 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
|
|
|