Cache TypeRep kinds aggressively

Cache `TypeRep k` in each `TrApp` or `TrTyCon` constructor of
`TypeRep (a :: k)`. This makes `typeRepKind` cheap.

With this change, we won't need any special effort to deserialize
typereps efficiently. The downside, of course, is that we make
`TypeRep`s slightly larger.

Reviewers: austin, hvr, bgamari, simonpj

Reviewed By: bgamari, simonpj

Subscribers: carter, simonpj, rwbarton, thomie

GHC Trac Issues: #14254

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

3 files changed, 36 insertions, 22 deletions

diff --git a/compiler/deSugar/DsBinds.hs b/compiler/deSugar/DsBinds.hs
index e11f580842..3048871d7f 100644
--- a/compiler/deSugar/DsBinds.hs
+++ b/compiler/deSugar/DsBinds.hs
@@ -1239,10 +1239,12 @@ ds_ev_typeable ty (EvTypeableTyCon tc kind_ev)
          -- Note that we use the kind of the type, not the TyCon from which it
          -- is constructed since the latter may be kind polymorphic whereas the
          -- former we know is not (we checked in the solver).
-       ; return $ mkApps (Var mkTrCon) [ Type (typeKind ty)
-                                       , Type ty
-                                       , tc_rep
-                                       , kind_args ]
+       ; let expr = mkApps (Var mkTrCon) [ Type (typeKind ty)
+                                         , Type ty
+                                         , tc_rep
+                                         , kind_args ]
+       -- ; pprRuntimeTrace "Trace mkTrTyCon" (ppr expr) expr
+       ; return expr
        }
 
 ds_ev_typeable ty (EvTypeableTyApp ev1 ev2)
@@ -1253,8 +1255,11 @@ ds_ev_typeable ty (EvTypeableTyApp ev1 ev2)
                     -- mkTrApp :: forall k1 k2 (a :: k1 -> k2) (b :: k1).
                     --            TypeRep a -> TypeRep b -> TypeRep (a b)
        ; let (k1, k2) = splitFunTy (typeKind t1)
-       ; return $ mkApps (mkTyApps (Var mkTrApp) [ k1, k2, t1, t2 ])
-                         [ e1, e2 ] }
+       ; let expr =  mkApps (mkTyApps (Var mkTrApp) [ k1, k2, t1, t2 ])
+                            [ e1, e2 ]
+       -- ; pprRuntimeTrace "Trace mkTrApp" (ppr expr) expr
+       ; return expr
+       }
 
 ds_ev_typeable ty (EvTypeableTrFun ev1 ev2)
   | Just (t1,t2) <- splitFunTy_maybe ty
diff --git a/compiler/prelude/PrelNames.hs b/compiler/prelude/PrelNames.hs
index f418348fcd..47b146559f 100644
--- a/compiler/prelude/PrelNames.hs
+++ b/compiler/prelude/PrelNames.hs
@@ -240,6 +240,7 @@ basicKnownKeyNames
         typeLitSymbolDataConName,
         typeLitNatDataConName,
         typeRepIdName,
+        mkTrTypeName,
         mkTrConName,
         mkTrAppName,
         mkTrFunName,
@@ -1256,6 +1257,7 @@ typeableClassName
   , typeRepTyConName
   , someTypeRepTyConName
   , someTypeRepDataConName
+  , mkTrTypeName
   , mkTrConName
   , mkTrAppName
   , mkTrFunName
@@ -1269,6 +1271,7 @@ typeRepTyConName      = tcQual  tYPEABLE_INTERNAL (fsLit "TypeRep")        typeR
 someTypeRepTyConName   = tcQual tYPEABLE_INTERNAL (fsLit "SomeTypeRep")    someTypeRepTyConKey
 someTypeRepDataConName = dcQual tYPEABLE_INTERNAL (fsLit "SomeTypeRep")    someTypeRepDataConKey
 typeRepIdName         = varQual tYPEABLE_INTERNAL (fsLit "typeRep#")       typeRepIdKey
+mkTrTypeName          = varQual tYPEABLE_INTERNAL (fsLit "mkTrType")       mkTrTypeKey
 mkTrConName           = varQual tYPEABLE_INTERNAL (fsLit "mkTrCon")        mkTrConKey
 mkTrAppName           = varQual tYPEABLE_INTERNAL (fsLit "mkTrApp")        mkTrAppKey
 mkTrFunName           = varQual tYPEABLE_INTERNAL (fsLit "mkTrFun")        mkTrFunKey
@@ -2329,6 +2332,7 @@ proxyHashKey = mkPreludeMiscIdUnique 502
 
 -- Used to make `Typeable` dictionaries
 mkTyConKey
+  , mkTrTypeKey
   , mkTrConKey
   , mkTrAppKey
   , mkTrFunKey
@@ -2337,12 +2341,13 @@ mkTyConKey
   , typeRepIdKey
   :: Unique
 mkTyConKey            = mkPreludeMiscIdUnique 503
-mkTrConKey            = mkPreludeMiscIdUnique 504
-mkTrAppKey            = mkPreludeMiscIdUnique 505
-typeNatTypeRepKey     = mkPreludeMiscIdUnique 506
-typeSymbolTypeRepKey  = mkPreludeMiscIdUnique 507
-typeRepIdKey          = mkPreludeMiscIdUnique 508
-mkTrFunKey            = mkPreludeMiscIdUnique 509
+mkTrTypeKey           = mkPreludeMiscIdUnique 504
+mkTrConKey            = mkPreludeMiscIdUnique 505
+mkTrAppKey            = mkPreludeMiscIdUnique 506
+typeNatTypeRepKey     = mkPreludeMiscIdUnique 507
+typeSymbolTypeRepKey  = mkPreludeMiscIdUnique 508
+typeRepIdKey          = mkPreludeMiscIdUnique 509
+mkTrFunKey            = mkPreludeMiscIdUnique 510
 
 -- Representations for primitive types
 trTYPEKey
@@ -2350,10 +2355,10 @@ trTYPEKey
   , trRuntimeRepKey
   , tr'PtrRepLiftedKey
   :: Unique
-trTYPEKey              = mkPreludeMiscIdUnique 510
-trTYPE'PtrRepLiftedKey = mkPreludeMiscIdUnique 511
-trRuntimeRepKey        = mkPreludeMiscIdUnique 512
-tr'PtrRepLiftedKey     = mkPreludeMiscIdUnique 513
+trTYPEKey              = mkPreludeMiscIdUnique 511
+trTYPE'PtrRepLiftedKey = mkPreludeMiscIdUnique 512
+trRuntimeRepKey        = mkPreludeMiscIdUnique 513
+tr'PtrRepLiftedKey     = mkPreludeMiscIdUnique 514
 
 -- KindReps for common cases
 starKindRepKey, starArrStarKindRepKey, starArrStarArrStarKindRepKey :: Unique
diff --git a/compiler/typecheck/TcTypeable.hs b/compiler/typecheck/TcTypeable.hs
index ed4b5483ef..6fa875b8d3 100644
--- a/compiler/typecheck/TcTypeable.hs
+++ b/compiler/typecheck/TcTypeable.hs
@@ -655,17 +655,20 @@ The TypeRep encoding of `Proxy Type Int` looks like this:
 
     $tcProxy :: GHC.Types.TyCon
     $trInt   :: TypeRep Int
-    $trType  :: TypeRep Type
+    TrType   :: TypeRep Type
 
     $trProxyType :: TypeRep (Proxy Type :: Type -> Type)
     $trProxyType = TrTyCon $tcProxy
-                           [$trType]  -- kind variable instantiation
+                           [TrType]  -- kind variable instantiation
+                           (tyConKind $tcProxy [TrType]) -- The TypeRep of
+                                                         -- Type -> Type
 
     $trProxy :: TypeRep (Proxy Type Int)
-    $trProxy = TrApp $trProxyType $trInt
+    $trProxy = TrApp $trProxyType $trInt TrType
 
     $tkProxy :: GHC.Types.KindRep
-    $tkProxy = KindRepFun (KindRepVar 0) (KindRepTyConApp $trType [])
+    $tkProxy = KindRepFun (KindRepVar 0)
+                          (KindRepTyConApp (KindRepTYPE LiftedRep) [])
 
 Note how $trProxyType cannot use 'TrApp', because TypeRep cannot represent
 polymorphic types.  So instead
@@ -679,9 +682,10 @@ polymorphic types.  So instead
        Proxy :: forall k. k->Type
 
  * A KindRep is just a recipe that we can instantiate with the
-   argument kinds, using Data.Typeable.Internal.instantiateKindRep.
+   argument kinds, using Data.Typeable.Internal.tyConKind and
+   store in the relevant 'TypeRep' constructor.
 
-   Data.Typeable.Internal.typeRepKind uses instantiateKindRep
+   Data.Typeable.Internal.typeRepKind looks up the stored kinds.
 
  * In a KindRep, the kind variables are represented by 0-indexed
    de Bruijn numbers: