210 lines
7.3 KiB
Haskell
210 lines
7.3 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 (
|
|
TypeKey, typeKey
|
|
, TypeSet, makeTypeSet
|
|
, containsTypes
|
|
, gmapMFor
|
|
, gmapMForRoute
|
|
, routeModify, routeGet, routeSet, Route
|
|
) where
|
|
|
|
import Control.Monad.Identity
|
|
import Control.Monad.State
|
|
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
|
|
import TreeUtils
|
|
import Utils
|
|
|
|
-- | A type identifier.
|
|
type TypeKey = Int
|
|
|
|
-- | Given a witness for a type, return its 'TypeKey'.
|
|
typeKey :: Typeable a => a -> TypeKey
|
|
typeKey x = unsafePerformIO $ typeRepKey $ typeOf x
|
|
|
|
-- | 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 a map from type keys to witnesses for
|
|
-- all the types it contains recursively.
|
|
containedTypes :: (Data a, Typeable a) => a -> IntMap DataBox
|
|
containedTypes start = containedTypes' (DataBox start) IntMap.empty
|
|
where
|
|
containedTypes' :: DataBox -> IntMap DataBox -> IntMap DataBox
|
|
containedTypes' box@(DataBox thisType) seen
|
|
= if thisKey `IntMap.member` seen
|
|
then seen
|
|
else foldl (\s t -> containedTypes' 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 $ containedTypes t)
|
|
| DataBox t <- IntMap.elems allTypes]
|
|
where
|
|
allTypes = containedTypes (undefined :: A.AST)
|
|
|
|
-- | Does a value contain any of the listed types?
|
|
-- (A value always contains its own type.)
|
|
containsTypes :: Data t => t -> [TypeKey] -> Bool
|
|
containsTypes x targets
|
|
= or $ map containsType targets
|
|
where
|
|
start :: TypeKey
|
|
start = typeKey x
|
|
|
|
containsType :: TypeKey -> Bool
|
|
containsType target
|
|
| start == target = True
|
|
| otherwise = case IntMap.lookup start contains of
|
|
Just set -> target `IntSet.member` set
|
|
Nothing -> True -- can't tell, so it might be
|
|
|
|
-- | A decision about what to do when we find a particular type during a
|
|
-- generic operation.
|
|
data TypeDecision =
|
|
-- | This is one of the types we're looking for.
|
|
Hit
|
|
-- | This isn't one of the types we're looking for, but it might contain one
|
|
-- of them.
|
|
| Through
|
|
-- | This isn't one of the types we're looking for, and there's no need to
|
|
-- look inside it.
|
|
| Miss
|
|
|
|
-- | A set of type information for use by 'gmapMFor'.
|
|
type TypeSet = IntMap TypeDecision
|
|
|
|
-- | Make a 'TypeSet' from a list of 'TypeKey's.
|
|
makeTypeSet :: [TypeKey] -> TypeSet
|
|
makeTypeSet targets
|
|
= IntMap.fromList [(tk, decide tk)
|
|
| tk <- IntMap.keys contains]
|
|
where
|
|
decide :: TypeKey -> TypeDecision
|
|
decide tk
|
|
| tk `elem` targets = Hit
|
|
| tk `IntSet.member` allThrough = Through
|
|
| otherwise = Miss
|
|
|
|
allThrough :: IntSet
|
|
allThrough
|
|
= IntSet.fromList $ filter containsThis $ IntMap.keys contains
|
|
where
|
|
containsThis tk
|
|
= case IntMap.lookup tk contains of
|
|
Just set -> or $ map (`IntSet.member` set) targets
|
|
Nothing -> False
|
|
|
|
-- | Type-smart generic mapM.
|
|
-- This is like 'gmapM', but it only applies the function to arguments that
|
|
-- could contain any of the target types.
|
|
gmapMFor :: (Monad m, Data t) =>
|
|
TypeSet -- ^ Target types
|
|
-> (forall s. Data s => s -> m s) -- ^ Function to apply
|
|
-> (t -> m t) -- ^ Generic operation
|
|
gmapMFor typeset f = gmapM (each f)
|
|
where
|
|
each :: (Monad m, Data t) =>
|
|
(forall s. Data s => s -> m s) -> (t -> m t)
|
|
each f x
|
|
= case IntMap.lookup (typeKey x) typeset of
|
|
Just Hit -> f x
|
|
Just Through -> gmapM (each f) x
|
|
Just Miss -> return x
|
|
Nothing -> return x
|
|
|
|
|
|
data Route inner outer = Route [Int] (forall m. Monad m => (inner -> m inner) -> (outer -> m outer))
|
|
|
|
instance Eq (Route inner outer) where
|
|
(==) (Route xns _) (Route yns _) = xns == yns
|
|
|
|
instance Ord (Route inner outer) where
|
|
compare (Route xns _) (Route yns _) = compare xns yns
|
|
|
|
routeModify :: Monad m => Route inner outer -> (inner -> m inner) -> (outer -> m
|
|
outer)
|
|
routeModify (Route _ wrap) = wrap
|
|
|
|
routeGet :: Route inner outer -> outer -> inner
|
|
routeGet route = flip execState undefined . routeModify route (\x -> put x >> return x)
|
|
|
|
routeSet :: Route inner outer -> inner -> outer -> outer
|
|
routeSet route x = runIdentity . routeModify route (const $ return x)
|
|
|
|
(@->) :: Route mid outer -> Route inner mid -> Route inner outer
|
|
(@->) (Route outInds outF) (Route inInds inF) = Route (outInds ++ inInds) (outF
|
|
. inF)
|
|
|
|
gmapMForRoute :: forall m t. (Monad m, Data t) =>
|
|
TypeSet ->
|
|
(forall s. Data s => (s, Route s t) -> m s)
|
|
-> (t -> m t)
|
|
gmapMForRoute typeset f = gmapMWithRoute (each f)
|
|
where
|
|
each :: Data u => (forall s. Data s => (s, Route s t) -> m s)
|
|
-> ((u, Route u 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, Route
|
|
b 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) => Int -> Route s t
|
|
makeRoute target = Route [target] (\f -> gmapFuncs [mkM' (if n == target then f else return) | n <- [0..]])
|