1{- 2(c) The University of Glasgow 2006 3(c) The GRASP/AQUA Project, Glasgow University, 1992-1998 4-} 5 6{-# LANGUAGE RankNTypes #-} 7{-# LANGUAGE TypeFamilies #-} 8{-# LANGUAGE FlexibleContexts #-} 9{-# LANGUAGE TypeSynonymInstances #-} 10{-# LANGUAGE FlexibleInstances #-} 11{-# LANGUAGE UndecidableInstances #-} 12module TrieMap( 13 -- * Maps over 'Maybe' values 14 MaybeMap, 15 -- * Maps over 'List' values 16 ListMap, 17 -- * Maps over 'Literal's 18 LiteralMap, 19 -- * 'TrieMap' class 20 TrieMap(..), insertTM, deleteTM, 21 22 -- * Things helpful for adding additional Instances. 23 (>.>), (|>), (|>>), XT, 24 foldMaybe, 25 -- * Map for leaf compression 26 GenMap, 27 lkG, xtG, mapG, fdG, 28 xtList, lkList 29 30 ) where 31 32import GhcPrelude 33 34import Literal 35import UniqDFM 36import Unique( Unique ) 37 38import qualified Data.Map as Map 39import qualified Data.IntMap as IntMap 40import Outputable 41import Control.Monad( (>=>) ) 42 43{- 44This module implements TrieMaps, which are finite mappings 45whose key is a structured value like a CoreExpr or Type. 46 47This file implements tries over general data structures. 48Implementation for tries over Core Expressions/Types are 49available in coreSyn/TrieMap. 50 51The regular pattern for handling TrieMaps on data structures was first 52described (to my knowledge) in Connelly and Morris's 1995 paper "A 53generalization of the Trie Data Structure"; there is also an accessible 54description of the idea in Okasaki's book "Purely Functional Data 55Structures", Section 10.3.2 56 57************************************************************************ 58* * 59 The TrieMap class 60* * 61************************************************************************ 62-} 63 64type XT a = Maybe a -> Maybe a -- How to alter a non-existent elt (Nothing) 65 -- or an existing elt (Just) 66 67class TrieMap m where 68 type Key m :: * 69 emptyTM :: m a 70 lookupTM :: forall b. Key m -> m b -> Maybe b 71 alterTM :: forall b. Key m -> XT b -> m b -> m b 72 mapTM :: (a->b) -> m a -> m b 73 74 foldTM :: (a -> b -> b) -> m a -> b -> b 75 -- The unusual argument order here makes 76 -- it easy to compose calls to foldTM; 77 -- see for example fdE below 78 79insertTM :: TrieMap m => Key m -> a -> m a -> m a 80insertTM k v m = alterTM k (\_ -> Just v) m 81 82deleteTM :: TrieMap m => Key m -> m a -> m a 83deleteTM k m = alterTM k (\_ -> Nothing) m 84 85---------------------- 86-- Recall that 87-- Control.Monad.(>=>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c 88 89(>.>) :: (a -> b) -> (b -> c) -> a -> c 90-- Reverse function composition (do f first, then g) 91infixr 1 >.> 92(f >.> g) x = g (f x) 93infixr 1 |>, |>> 94 95(|>) :: a -> (a->b) -> b -- Reverse application 96x |> f = f x 97 98---------------------- 99(|>>) :: TrieMap m2 100 => (XT (m2 a) -> m1 (m2 a) -> m1 (m2 a)) 101 -> (m2 a -> m2 a) 102 -> m1 (m2 a) -> m1 (m2 a) 103(|>>) f g = f (Just . g . deMaybe) 104 105deMaybe :: TrieMap m => Maybe (m a) -> m a 106deMaybe Nothing = emptyTM 107deMaybe (Just m) = m 108 109{- 110************************************************************************ 111* * 112 IntMaps 113* * 114************************************************************************ 115-} 116 117instance TrieMap IntMap.IntMap where 118 type Key IntMap.IntMap = Int 119 emptyTM = IntMap.empty 120 lookupTM k m = IntMap.lookup k m 121 alterTM = xtInt 122 foldTM k m z = IntMap.foldr k z m 123 mapTM f m = IntMap.map f m 124 125xtInt :: Int -> XT a -> IntMap.IntMap a -> IntMap.IntMap a 126xtInt k f m = IntMap.alter f k m 127 128instance Ord k => TrieMap (Map.Map k) where 129 type Key (Map.Map k) = k 130 emptyTM = Map.empty 131 lookupTM = Map.lookup 132 alterTM k f m = Map.alter f k m 133 foldTM k m z = Map.foldr k z m 134 mapTM f m = Map.map f m 135 136 137{- 138Note [foldTM determinism] 139~~~~~~~~~~~~~~~~~~~~~~~~~ 140We want foldTM to be deterministic, which is why we have an instance of 141TrieMap for UniqDFM, but not for UniqFM. Here's an example of some things that 142go wrong if foldTM is nondeterministic. Consider: 143 144 f a b = return (a <> b) 145 146Depending on the order that the typechecker generates constraints you 147get either: 148 149 f :: (Monad m, Monoid a) => a -> a -> m a 150 151or: 152 153 f :: (Monoid a, Monad m) => a -> a -> m a 154 155The generated code will be different after desugaring as the dictionaries 156will be bound in different orders, leading to potential ABI incompatibility. 157 158One way to solve this would be to notice that the typeclasses could be 159sorted alphabetically. 160 161Unfortunately that doesn't quite work with this example: 162 163 f a b = let x = a <> a; y = b <> b in x 164 165where you infer: 166 167 f :: (Monoid m, Monoid m1) => m1 -> m -> m1 168 169or: 170 171 f :: (Monoid m1, Monoid m) => m1 -> m -> m1 172 173Here you could decide to take the order of the type variables in the type 174according to depth first traversal and use it to order the constraints. 175 176The real trouble starts when the user enables incoherent instances and 177the compiler has to make an arbitrary choice. Consider: 178 179 class T a b where 180 go :: a -> b -> String 181 182 instance (Show b) => T Int b where 183 go a b = show a ++ show b 184 185 instance (Show a) => T a Bool where 186 go a b = show a ++ show b 187 188 f = go 10 True 189 190GHC is free to choose either dictionary to implement f, but for the sake of 191determinism we'd like it to be consistent when compiling the same sources 192with the same flags. 193 194inert_dicts :: DictMap is implemented with a TrieMap. In getUnsolvedInerts it 195gets converted to a bag of (Wanted) Cts using a fold. Then in 196solve_simple_wanteds it's merged with other WantedConstraints. We want the 197conversion to a bag to be deterministic. For that purpose we use UniqDFM 198instead of UniqFM to implement the TrieMap. 199 200See Note [Deterministic UniqFM] in UniqDFM for more details on how it's made 201deterministic. 202-} 203 204instance TrieMap UniqDFM where 205 type Key UniqDFM = Unique 206 emptyTM = emptyUDFM 207 lookupTM k m = lookupUDFM m k 208 alterTM k f m = alterUDFM f m k 209 foldTM k m z = foldUDFM k z m 210 mapTM f m = mapUDFM f m 211 212{- 213************************************************************************ 214* * 215 Maybes 216* * 217************************************************************************ 218 219If m is a map from k -> val 220then (MaybeMap m) is a map from (Maybe k) -> val 221-} 222 223data MaybeMap m a = MM { mm_nothing :: Maybe a, mm_just :: m a } 224 225instance TrieMap m => TrieMap (MaybeMap m) where 226 type Key (MaybeMap m) = Maybe (Key m) 227 emptyTM = MM { mm_nothing = Nothing, mm_just = emptyTM } 228 lookupTM = lkMaybe lookupTM 229 alterTM = xtMaybe alterTM 230 foldTM = fdMaybe 231 mapTM = mapMb 232 233mapMb :: TrieMap m => (a->b) -> MaybeMap m a -> MaybeMap m b 234mapMb f (MM { mm_nothing = mn, mm_just = mj }) 235 = MM { mm_nothing = fmap f mn, mm_just = mapTM f mj } 236 237lkMaybe :: (forall b. k -> m b -> Maybe b) 238 -> Maybe k -> MaybeMap m a -> Maybe a 239lkMaybe _ Nothing = mm_nothing 240lkMaybe lk (Just x) = mm_just >.> lk x 241 242xtMaybe :: (forall b. k -> XT b -> m b -> m b) 243 -> Maybe k -> XT a -> MaybeMap m a -> MaybeMap m a 244xtMaybe _ Nothing f m = m { mm_nothing = f (mm_nothing m) } 245xtMaybe tr (Just x) f m = m { mm_just = mm_just m |> tr x f } 246 247fdMaybe :: TrieMap m => (a -> b -> b) -> MaybeMap m a -> b -> b 248fdMaybe k m = foldMaybe k (mm_nothing m) 249 . foldTM k (mm_just m) 250 251{- 252************************************************************************ 253* * 254 Lists 255* * 256************************************************************************ 257-} 258 259data ListMap m a 260 = LM { lm_nil :: Maybe a 261 , lm_cons :: m (ListMap m a) } 262 263instance TrieMap m => TrieMap (ListMap m) where 264 type Key (ListMap m) = [Key m] 265 emptyTM = LM { lm_nil = Nothing, lm_cons = emptyTM } 266 lookupTM = lkList lookupTM 267 alterTM = xtList alterTM 268 foldTM = fdList 269 mapTM = mapList 270 271instance (TrieMap m, Outputable a) => Outputable (ListMap m a) where 272 ppr m = text "List elts" <+> ppr (foldTM (:) m []) 273 274mapList :: TrieMap m => (a->b) -> ListMap m a -> ListMap m b 275mapList f (LM { lm_nil = mnil, lm_cons = mcons }) 276 = LM { lm_nil = fmap f mnil, lm_cons = mapTM (mapTM f) mcons } 277 278lkList :: TrieMap m => (forall b. k -> m b -> Maybe b) 279 -> [k] -> ListMap m a -> Maybe a 280lkList _ [] = lm_nil 281lkList lk (x:xs) = lm_cons >.> lk x >=> lkList lk xs 282 283xtList :: TrieMap m => (forall b. k -> XT b -> m b -> m b) 284 -> [k] -> XT a -> ListMap m a -> ListMap m a 285xtList _ [] f m = m { lm_nil = f (lm_nil m) } 286xtList tr (x:xs) f m = m { lm_cons = lm_cons m |> tr x |>> xtList tr xs f } 287 288fdList :: forall m a b. TrieMap m 289 => (a -> b -> b) -> ListMap m a -> b -> b 290fdList k m = foldMaybe k (lm_nil m) 291 . foldTM (fdList k) (lm_cons m) 292 293foldMaybe :: (a -> b -> b) -> Maybe a -> b -> b 294foldMaybe _ Nothing b = b 295foldMaybe k (Just a) b = k a b 296 297{- 298************************************************************************ 299* * 300 Basic maps 301* * 302************************************************************************ 303-} 304 305type LiteralMap a = Map.Map Literal a 306 307{- 308************************************************************************ 309* * 310 GenMap 311* * 312************************************************************************ 313 314Note [Compressed TrieMap] 315~~~~~~~~~~~~~~~~~~~~~~~~~ 316 317The GenMap constructor augments TrieMaps with leaf compression. This helps 318solve the performance problem detailed in #9960: suppose we have a handful 319H of entries in a TrieMap, each with a very large key, size K. If you fold over 320such a TrieMap you'd expect time O(H). That would certainly be true of an 321association list! But with TrieMap we actually have to navigate down a long 322singleton structure to get to the elements, so it takes time O(K*H). This 323can really hurt on many type-level computation benchmarks: 324see for example T9872d. 325 326The point of a TrieMap is that you need to navigate to the point where only one 327key remains, and then things should be fast. So the point of a SingletonMap 328is that, once we are down to a single (key,value) pair, we stop and 329just use SingletonMap. 330 331'EmptyMap' provides an even more basic (but essential) optimization: if there is 332nothing in the map, don't bother building out the (possibly infinite) recursive 333TrieMap structure! 334 335Compressed triemaps are heavily used by CoreMap. So we have to mark some things 336as INLINEABLE to permit specialization. 337-} 338 339data GenMap m a 340 = EmptyMap 341 | SingletonMap (Key m) a 342 | MultiMap (m a) 343 344instance (Outputable a, Outputable (m a)) => Outputable (GenMap m a) where 345 ppr EmptyMap = text "Empty map" 346 ppr (SingletonMap _ v) = text "Singleton map" <+> ppr v 347 ppr (MultiMap m) = ppr m 348 349-- TODO undecidable instance 350instance (Eq (Key m), TrieMap m) => TrieMap (GenMap m) where 351 type Key (GenMap m) = Key m 352 emptyTM = EmptyMap 353 lookupTM = lkG 354 alterTM = xtG 355 foldTM = fdG 356 mapTM = mapG 357 358--We want to be able to specialize these functions when defining eg 359--tries over (GenMap CoreExpr) which requires INLINEABLE 360 361{-# INLINEABLE lkG #-} 362lkG :: (Eq (Key m), TrieMap m) => Key m -> GenMap m a -> Maybe a 363lkG _ EmptyMap = Nothing 364lkG k (SingletonMap k' v') | k == k' = Just v' 365 | otherwise = Nothing 366lkG k (MultiMap m) = lookupTM k m 367 368{-# INLINEABLE xtG #-} 369xtG :: (Eq (Key m), TrieMap m) => Key m -> XT a -> GenMap m a -> GenMap m a 370xtG k f EmptyMap 371 = case f Nothing of 372 Just v -> SingletonMap k v 373 Nothing -> EmptyMap 374xtG k f m@(SingletonMap k' v') 375 | k' == k 376 -- The new key matches the (single) key already in the tree. Hence, 377 -- apply @f@ to @Just v'@ and build a singleton or empty map depending 378 -- on the 'Just'/'Nothing' response respectively. 379 = case f (Just v') of 380 Just v'' -> SingletonMap k' v'' 381 Nothing -> EmptyMap 382 | otherwise 383 -- We've hit a singleton tree for a different key than the one we are 384 -- searching for. Hence apply @f@ to @Nothing@. If result is @Nothing@ then 385 -- we can just return the old map. If not, we need a map with *two* 386 -- entries. The easiest way to do that is to insert two items into an empty 387 -- map of type @m a@. 388 = case f Nothing of 389 Nothing -> m 390 Just v -> emptyTM |> alterTM k' (const (Just v')) 391 >.> alterTM k (const (Just v)) 392 >.> MultiMap 393xtG k f (MultiMap m) = MultiMap (alterTM k f m) 394 395{-# INLINEABLE mapG #-} 396mapG :: TrieMap m => (a -> b) -> GenMap m a -> GenMap m b 397mapG _ EmptyMap = EmptyMap 398mapG f (SingletonMap k v) = SingletonMap k (f v) 399mapG f (MultiMap m) = MultiMap (mapTM f m) 400 401{-# INLINEABLE fdG #-} 402fdG :: TrieMap m => (a -> b -> b) -> GenMap m a -> b -> b 403fdG _ EmptyMap = \z -> z 404fdG k (SingletonMap _ v) = \z -> k v z 405fdG k (MultiMap m) = foldTM k m 406