Compress TypeMap TrieMap leaves with singleton constructor.

Suppose we have a handful H of entries in a TrieMap, each with a very large
key, size K. If you fold over such a TrieMap you'd expect time O(H). That would
certainly be true of an association list! But with TrieMap we actually have to
navigate down a long singleton structure to get to the elements, so it takes
time O(K*H).  The point of a TrieMap is that you need to navigate to the point
where only one key remains, and then things should be fast.

This is a starting point: we can improve the patch by generalizing the
singleton constructor so it applies to CoercionMap and CoreMap; I'll do this
in a later commit.

Summary: Signed-off-by: Edward Z. Yang <ezyang@cs.stanford.edu>

Test Plan: validate

Reviewers: simonpj, austin

Subscribers: carter, thomie

Differential Revision: https://phabricator.haskell.org/D606

GHC Trac Issues: #9960

2 files changed, 60 insertions, 2 deletions

diff --git a/compiler/coreSyn/TrieMap.hs b/compiler/coreSyn/TrieMap.hs
index 9197386ccb..a8ac2b12dc 100644
--- a/compiler/coreSyn/TrieMap.hs
+++ b/compiler/coreSyn/TrieMap.hs
@@ -622,6 +622,7 @@ mapR f = RM . mapTM f . unRM
 
 data TypeMap a
   = EmptyTM
+  | SingletonTM (CmEnv, Type) a
   | TM { tm_var   :: VarMap a
        , tm_app    :: TypeMap (TypeMap a)
        , tm_fun    :: TypeMap (TypeMap a)
@@ -630,6 +631,41 @@ data TypeMap a
        , tm_tylit  :: TyLitMap a
        }
 
+eqTypesModuloDeBruijn :: (CmEnv, [Type]) -> (CmEnv, [Type]) -> Bool
+eqTypesModuloDeBruijn (_, []) (_, []) = True
+eqTypesModuloDeBruijn (env, ty:tys) (env', ty':tys') =
+    eqTypeModuloDeBruijn (env, ty) (env', ty') &&
+    eqTypesModuloDeBruijn (env, tys) (env', tys')
+eqTypesModuloDeBruijn _ _ = False
+
+-- NB: need to coreView!
+eqTypeModuloDeBruijn :: (CmEnv, Type) -> (CmEnv, Type) -> Bool
+eqTypeModuloDeBruijn env_t@(env, t) env_t'@(env', t')
+    -- ToDo: I guess we can make this a little more efficient
+    | Just new_t  <- coreView t  = eqTypeModuloDeBruijn (env, new_t) env_t'
+    | Just new_t' <- coreView t' = eqTypeModuloDeBruijn env_t (env', new_t')
+eqTypeModuloDeBruijn (env, t) (env', t') =
+    case (t, t') of
+        (TyVarTy v, TyVarTy v')
+            -> case (lookupCME env v, lookupCME env' v') of
+                (Just bv, Just bv') -> bv == bv'
+                (Nothing, Nothing) -> v == v'
+                _ -> False
+        (AppTy t1 t2, AppTy t1' t2')
+            -> eqTypeModuloDeBruijn (env, t1) (env', t1') &&
+               eqTypeModuloDeBruijn (env, t2) (env', t2')
+        (FunTy t1 t2, FunTy t1' t2')
+            -> eqTypeModuloDeBruijn (env, t1) (env', t1') &&
+               eqTypeModuloDeBruijn (env, t2) (env', t2')
+        (TyConApp tc tys, TyConApp tc' tys')
+            -> tc == tc' && eqTypesModuloDeBruijn (env, tys) (env', tys')
+        (LitTy l, LitTy l')
+            -> l == l'
+        (ForAllTy tv ty, ForAllTy tv' ty')
+            -> eqTypeModuloDeBruijn (env, tyVarKind tv) (env', tyVarKind tv') &&
+               eqTypeModuloDeBruijn (extendCME env tv, ty)
+                                    (extendCME env' tv', ty')
+        _ -> False
 
 instance Outputable a => Outputable (TypeMap a) where
   ppr m = text "TypeMap elts" <+> ppr (foldTypeMap (:) [] m)
@@ -647,6 +683,10 @@ lookupTypeMap cm t = lkT emptyCME t cm
 -- This only considers saturated applications (i.e. TyConApp ones).
 lookupTypeMapTyCon :: TypeMap a -> TyCon -> [a]
 lookupTypeMapTyCon EmptyTM _ = []
+lookupTypeMapTyCon (SingletonTM (_, TyConApp tc' _) v) tc
+    | tc' == tc = [v]
+    | otherwise = []
+lookupTypeMapTyCon SingletonTM{} _ = []
 lookupTypeMapTyCon TM { tm_tc_app = cs } tc =
   case lookupUFM cs tc of
     Nothing -> []
@@ -673,6 +713,7 @@ instance TrieMap TypeMap where
 
 mapT :: (a->b) -> TypeMap a -> TypeMap b
 mapT _ EmptyTM = EmptyTM
+mapT f (SingletonTM env_ty v) = SingletonTM env_ty (f v)
 mapT f (TM { tm_var  = tvar, tm_app = tapp, tm_fun = tfun
            , tm_tc_app = ttcapp, tm_forall = tforall, tm_tylit = tlit })
   = TM { tm_var    = mapTM f tvar
@@ -686,6 +727,10 @@ mapT f (TM { tm_var  = tvar, tm_app = tapp, tm_fun = tfun
 lkT :: CmEnv -> Type -> TypeMap a -> Maybe a
 lkT env ty m
   | EmptyTM <- m = Nothing
+  | SingletonTM env_ty v <- m =
+        if eqTypeModuloDeBruijn env_ty (env, ty)
+            then Just v
+            else Nothing
   | otherwise    = go ty m
   where
     go ty | Just ty' <- coreView ty = go ty'
@@ -700,7 +745,18 @@ lkT env ty m
 -----------------
 xtT :: CmEnv -> Type -> XT a -> TypeMap a -> TypeMap a
 xtT env ty f m
-  | EmptyTM <- m            = xtT env ty  f wrapEmptyTypeMap
+  | EmptyTM <- m            = case f Nothing of
+                                Just v -> SingletonTM (env, ty) v
+                                Nothing -> EmptyTM
+  | SingletonTM env_ty@(env', ty') v' <- m
+    = if eqTypeModuloDeBruijn env_ty (env, ty)
+        then case f (Just v') of
+                Just v'' -> SingletonTM env_ty v''
+                Nothing -> EmptyTM
+        else case f Nothing of
+                Nothing -> SingletonTM env_ty v'
+                Just v  -> wrapEmptyTypeMap |> xtT env' ty' (const (Just v'))
+                                           >.> xtT env  ty  (const (Just v))
   | Just ty' <- coreView ty = xtT env ty' f m
 
 xtT env (TyVarTy v)       f  m = m { tm_var    = tm_var m |> xtVar env v f }
@@ -714,6 +770,7 @@ xtT _   (LitTy l)         f  m = m { tm_tylit  = tm_tylit m |> xtTyLit l f }
 
 fdT :: (a -> b -> b) -> TypeMap a -> b -> b
 fdT _ EmptyTM = \z -> z
+fdT k (SingletonTM _ v) = \z -> k v z
 fdT k m = foldTM k (tm_var m)
         . foldTM (foldTM k) (tm_app m)
         . foldTM (foldTM k) (tm_fun m)
diff --git a/testsuite/tests/perf/compiler/all.T b/testsuite/tests/perf/compiler/all.T
index 2d964978b6..10136bb3be 100644
--- a/testsuite/tests/perf/compiler/all.T
+++ b/testsuite/tests/perf/compiler/all.T
@@ -607,9 +607,10 @@ test('T9872c',
 test('T9872d',
      [ only_ways(['normal']),
        compiler_stats_num_field('bytes allocated',
-          [(wordsize(64), 739189056, 5),
+          [(wordsize(64), 687562440, 5),
           # 2014-12-18    796071864   Initally created
           # 2014-12-18    739189056   Reduce type families even more eagerly
+          # 2015-01-07    687562440   TrieMap leaf compression
            (wordsize(32), 353644844, 5)
           ]),
       ],