[commit: ghc] type-nats-simple: Add a function to lookup all things with the same top tycon in a TypeMap (0d7649a)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/ghc

On branch  : type-nats-simple
Link       : http://ghc.haskell.org/trac/ghc/changeset/0d7649a689af0f1fa715ceb2501fc46c89ca6ee9/ghc

>---------------------------------------------------------------

commit 0d7649a689af0f1fa715ceb2501fc46c89ca6ee9
Author: Iavor S. Diatchki <iavor.diatchki at gmail.com>
Date:   Sat Sep 7 17:06:43 2013 -0700

    Add a function to lookup all things with the same top tycon in a TypeMap


>---------------------------------------------------------------

0d7649a689af0f1fa715ceb2501fc46c89ca6ee9
 compiler/coreSyn/TrieMap.lhs |   13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/compiler/coreSyn/TrieMap.lhs b/compiler/coreSyn/TrieMap.lhs
index b7b3a56..255ab89 100644
--- a/compiler/coreSyn/TrieMap.lhs
+++ b/compiler/coreSyn/TrieMap.lhs
@@ -18,7 +18,8 @@ module TrieMap(
    CoercionMap, 
    MaybeMap, 
    ListMap,
-   TrieMap(..)
+   TrieMap(..),
+   lookupTypeMapTyCon
  ) where
 
 import CoreSyn
@@ -27,6 +28,7 @@ import Literal
 import Name
 import Type
 import TypeRep
+import TyCon(TyCon)
 import Var
 import UniqFM
 import Unique( Unique )
@@ -648,6 +650,15 @@ emptyTypeMap = EmptyTM
 lookupTypeMap :: TypeMap a -> Type -> Maybe a
 lookupTypeMap cm t = lkT emptyCME t cm
 
+-- Returns the type map entries that have keys starting with the given tycon.
+-- This only considers saturated applications (i.e. TyConApp ones).
+lookupTypeMapTyCon :: TypeMap a -> TyCon -> [a]
+lookupTypeMapTyCon EmptyTM _ = []
+lookupTypeMapTyCon TM { tm_tc_app = cs } tc =
+  case lookupUFM cs tc of
+    Nothing -> []
+    Just xs -> foldTM (:) xs []
+
 extendTypeMap :: TypeMap a -> Type -> a -> TypeMap a
 extendTypeMap m t v = xtT emptyCME t (\_ -> Just v) m