tock-mirror/common/GenericUtils.hs

{-
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(..), (@->), routeIdentity, routeId, routeList
  , route11, route22, route23, route33, route34, route44, route45, route55
  ) 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

-- | A Route is a way of navigating to a particular node in a tree structure.
--
-- Let's say that you have some binary tree structure:
--
-- > data BinTree a = Leaf a | Branch (BinTree a) (BinTree a)
--
-- Suppose you then have a big binary tree of integers, potentially with duplicate values,
-- and you want to be able to modify a particular integer.  You can't modify in-place,
-- because this is a functional language.  So you instead want to be able to apply
-- a modify function to the whole tree that really just modifies the particular
-- integer.
--
-- To do this you can use:
-- 
-- > myRoute :: Route Int (BinTree Int)
--
-- You apply it as follows (for example, to increment the integer):
--
-- > runIdentity $ routeModify myRoute (return . (+1)) myTree
--
-- The modifier is monadic because that's usually how we want to use it, but we
-- can use the identity monad as above for pure functions.  This will only work
-- if the route is valid on the given tree.
--
-- Another useful aspect is composition.  If your tree was in a tree of trees:
--
-- > routeToInnerTree :: Route (BinTree Int) (BinTree (BinTree Int))
--
-- You could compose this with the earlier route:
--
-- > routeToInnerTree @-> myRoute :: Route Int (BinTree (BinTree Int))
--
-- These routes are a little like zippers, but (in my opinion) easier to use, and
-- tack on to existing code with complex data structures (without needing any code
-- generation).  You can either compose routes yourself (as the flow-graph building
-- does) or by using 'gmapMForRoute'.
--
-- Routes support Eq, Show and Ord.  All these instances represent a route as a
-- list of integers: a route-map.  [0,2,1] means first child (zero-based), then
-- third child, then first child of the given data-type.  Routes are ordered using
-- the standard list ordering (lexicographic) over this representation.
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

instance Show (Route inner outer) where
  show (Route ns _) = "Route " ++ show ns

-- | Gets the integer-list version of a route.  See the documentation of 'Route'.
routeId :: Route inner outer -> [Int]
routeId (Route ns _) = ns

-- | Given an index (zero is the first item), forms a route to that index item
-- in the list.  So for example:
--
-- > runIdentity $ routeModify (routeList 3) (return . (*10)) [0,1,2,3,4,5] == [0,1,2,30,4,5]
-- 
routeList :: Int -> Route a [a]
routeList 0 = Route [0] (\f (x:xs) -> f x >>* (: xs))
routeList n = Route [1] (\f (x:xs) -> f xs >>* (x:)) @-> routeList (n-1)

-- | Applies a monadic modification function using the given route.
routeModify :: Monad m => Route inner outer -> (inner -> m inner) -> (outer -> m
  outer)
routeModify (Route _ wrap) = wrap

-- | Given a route, gets the value in the large data structure that is pointed
-- to by that route.
routeGet :: Route inner outer -> outer -> inner
routeGet route = flip execState undefined . routeModify route (\x -> put x >> return x)

-- | Given a route, sets the value in the large data structure that is pointed
-- to by that route.
routeSet :: Route inner outer -> inner -> outer -> outer
routeSet route x = runIdentity . routeModify route (const $ return x)

-- | Composes two routes together.  The outer-to-mid route goes on the left hand
-- side, and the mid-to-inner goes on the right hand side to form an outer-to-inner
-- route.
(@->) :: Route mid outer -> Route inner mid -> Route inner outer
(@->) (Route outInds outF) (Route inInds inF) = Route (outInds ++ inInds) (outF
  . inF)

-- | The identity route.  This has various obvious properties:
--
-- > routeGet routeIdentity == id
-- > routeSet routeIdentity == const
-- > routeModify routeIdentity == id
-- > routeIdentity @-> route == route
-- > route @-> routeIdentity == route
routeIdentity :: Route a a
routeIdentity = Route [] id

-- | Acts just like gmapMFor, except that it also tells you the route to the node
-- that your generic function is being applied to.
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

-- A helper for gmapMForRoute that maps over the direct children and supplies routes
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..]])

decomp11 :: (Monad m, Data a, Typeable a0) => (a0 -> a) -> (a0 -> m a0) -> (a -> m a)
decomp11 con f1 = decomp1 con f1

decomp22 :: (Monad m, Data a, Typeable a0, Typeable a1) => (a0 -> a1 -> a) -> (a1 -> m a1) -> (a -> m a)
decomp22 con f1 = decomp2 con return f1

decomp23 :: (Monad m, Data a, Typeable a0, Typeable a1, Typeable a2) => (a0 -> a1 -> a2 -> a) -> (a1 -> m a1) -> (a -> m a)
decomp23 con f1 = decomp3 con return f1 return

decomp33 :: (Monad m, Data a, Typeable a0, Typeable a1, Typeable a2) => (a0 -> a1 -> a2 -> a) -> (a2 -> m a2) -> (a -> m a)
decomp33 con f2 = decomp3 con return return f2

decomp34 :: (Monad m, Data a, Typeable a0, Typeable a1, Typeable a2, Typeable a3) =>
  (a0 -> a1 -> a2 -> a3 -> a) -> (a2 -> m a2) -> (a -> m a)
decomp34 con f2 = decomp4 con return return f2 return

decomp44 :: (Monad m, Data a, Typeable a0, Typeable a1, Typeable a2, Typeable a3) =>
  (a0 -> a1 -> a2 -> a3 -> a) -> (a3 -> m a3) -> (a -> m a)
decomp44 con f3 = decomp4 con return return return f3

decomp45 :: (Monad m, Data a, Typeable a0, Typeable a1, Typeable a2, Typeable a3, Typeable a4) =>
  (a0 -> a1 -> a2 -> a3 -> a4 -> a) -> (a3 -> m a3) -> (a -> m a)
decomp45 con f3 = decomp5 con return return return f3 return

decomp55 :: (Monad m, Data a, Typeable a0, Typeable a1, Typeable a2, Typeable a3, Typeable a4) =>
  (a0 -> a1 -> a2 -> a3 -> a4 -> a) -> (a4 -> m a4) -> (a -> m a)
decomp55 con f4 = decomp5 con return return return return f4

route11 :: (Data a, Typeable a0) => Route a b -> (a0 -> a) -> Route a0 b
route11 route con = route @-> Route [0] (decomp11 con)

route22 :: (Data a, Typeable a0, Typeable a1) => Route a b -> (a0 -> a1 -> a) -> Route a1 b
route22 route con = route @-> Route [1] (decomp22 con)

route23 :: (Data a, Typeable a0, Typeable a1, Typeable a2) => Route a b -> (a0 -> a1 -> a2 -> a) -> Route a1 b
route23 route con = route @-> Route [1] (decomp23 con)

route33 :: (Data a, Typeable a0, Typeable a1, Typeable a2) => Route a b -> (a0 -> a1 -> a2 -> a) -> Route a2 b
route33 route con = route @-> Route [2] (decomp33 con)

route34 :: (Data a, Typeable a0, Typeable a1, Typeable a2, Typeable a3) =>
  Route a b -> (a0 -> a1 -> a2 -> a3 -> a) -> Route a2 b
route34 route con = route @-> Route [2] (decomp34 con)

route44 :: (Data a, Typeable a0, Typeable a1, Typeable a2, Typeable a3) =>
  Route a b -> (a0 -> a1 -> a2 -> a3 -> a) -> Route a3 b
route44 route con = route @-> Route [3] (decomp44 con)

route45 :: (Data a, Typeable a0, Typeable a1, Typeable a2, Typeable a3, Typeable a4) =>
  Route a b -> (a0 -> a1 -> a2 -> a3 -> a4 -> a) -> Route a3 b
route45 route con = route @-> Route [3] (decomp45 con)

route55 :: (Data a, Typeable a0, Typeable a1, Typeable a2, Typeable a3, Typeable a4) =>
  Route a b -> (a0 -> a1 -> a2 -> a3 -> a4 -> a) -> Route a4 b
route55 route con = route @-> Route [4] (decomp55 con)

-- TODO we should be able to provide versions of these that do not need to know
-- the constructor or the arity