Lift constructor tag allocation out of a loop

Before this change, for each constructor that we want
to allocate a tag for we would traverse a list of all
the constructors in a datatype to determine which tag
a constructor should get.

This is obviously quadratic and for datatypes with 10k
constructors it actually makes a big difference.

This change implements the plan outlined by @simonpj in
https://mail.haskell.org/pipermail/ghc-devs/2017-October/014974.html
which is basically about using a map and constructing it outside the
loop.

One place where things got a bit awkward was TysWiredIn.hs,
it would have been possible to just assign the tags by hand, but
that seemed error-prone to me, so I decided to go through a map
there as well.

Test Plan:
./validate
On a file with 10k constructors
Before:
   8,130,522,344 bytes allocated in the heap
  Total   time    3.682s  (  3.920s elapsed)
After:
   4,133,478,744 bytes allocated in the heap
  Total   time    2.509s  (  2.750s elapsed)

Reviewers: simonpj, bgamari

Reviewed By: simonpj

Subscribers: goldfire, rwbarton, thomie, simonmar, carter, simonpj

GHC Trac Issues: #14657

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

2 files changed, 14 insertions, 4 deletions

diff --git a/compiler/iface/BuildTyCl.hs b/compiler/iface/BuildTyCl.hs
index 113ec12b63..43e9408430 100644
--- a/compiler/iface/BuildTyCl.hs
+++ b/compiler/iface/BuildTyCl.hs
@@ -27,6 +27,7 @@ import Var
 import VarSet
 import BasicTypes
 import Name
+import NameEnv
 import MkId
 import Class
 import TyCon
@@ -107,13 +108,16 @@ buildDataCon :: FamInstEnvs
                                        -- or the GADT equalities
            -> [Type] -> Type           -- Argument and result types
            -> TyCon                    -- Rep tycon
+           -> NameEnv ConTag           -- Maps the Name of each DataCon to its
+                                       -- ConTag
            -> TcRnIf m n DataCon
 -- A wrapper for DataCon.mkDataCon that
 --   a) makes the worker Id
 --   b) makes the wrapper Id if necessary, including
 --      allocating its unique (hence monadic)
-buildDataCon fam_envs src_name declared_infix prom_info src_bangs impl_bangs field_lbls
-             univ_tvs ex_tvs user_tvbs eq_spec ctxt arg_tys res_ty rep_tycon
+buildDataCon fam_envs src_name declared_infix prom_info src_bangs impl_bangs
+             field_lbls univ_tvs ex_tvs user_tvbs eq_spec ctxt arg_tys res_ty
+             rep_tycon tag_map
   = do  { wrap_name <- newImplicitBinder src_name mkDataConWrapperOcc
         ; work_name <- newImplicitBinder src_name mkDataConWorkerOcc
         -- This last one takes the name of the data constructor in the source
@@ -124,10 +128,12 @@ buildDataCon fam_envs src_name declared_infix prom_info src_bangs impl_bangs fie
         ; us <- newUniqueSupply
         ; dflags <- getDynFlags
         ; let stupid_ctxt = mkDataConStupidTheta rep_tycon arg_tys univ_tvs
+              tag = lookupNameEnv_NF tag_map src_name
+              -- See Note [Constructor tag allocation], fixes #14657
               data_con = mkDataCon src_name declared_infix prom_info
                                    src_bangs field_lbls
                                    univ_tvs ex_tvs user_tvbs eq_spec ctxt
-                                   arg_tys res_ty NoRRI rep_tycon
+                                   arg_tys res_ty NoRRI rep_tycon tag
                                    stupid_ctxt dc_wrk dc_rep
               dc_wrk = mkDataConWorkId work_name data_con
               dc_rep = initUs_ us (mkDataConRep dflags fam_envs wrap_name
@@ -307,6 +313,7 @@ buildClass tycon_name binders roles fds
                                    arg_tys
                                    (mkTyConApp rec_tycon (mkTyVarTys univ_tvs))
                                    rec_tycon
+                                   (mkTyConTagMap rec_tycon)
 
         ; rhs <- case () of
                   _ | use_newtype
diff --git a/compiler/iface/TcIface.hs b/compiler/iface/TcIface.hs
index 6fad8da87c..70438f6337 100644
--- a/compiler/iface/TcIface.hs
+++ b/compiler/iface/TcIface.hs
@@ -897,6 +897,9 @@ tcIfaceDataCons tycon_name tycon tc_tybinders if_cons
     univ_tvs :: [TyVar]
     univ_tvs = binderVars (tyConTyVarBinders tc_tybinders)
 
+    tag_map :: NameEnv ConTag
+    tag_map = mkTyConTagMap tycon
+
     tc_con_decl (IfCon { ifConInfix = is_infix,
                          ifConExTvs = ex_bndrs,
                          ifConUserTvBinders = user_bndrs,
@@ -960,7 +963,7 @@ tcIfaceDataCons tycon_name tycon tc_tybinders if_cons
                        lbl_names
                        univ_tvs ex_tvs user_tv_bndrs
                        eq_spec theta
-                       arg_tys orig_res_ty tycon
+                       arg_tys orig_res_ty tycon tag_map
         ; traceIf (text "Done interface-file tc_con_decl" <+> ppr dc_name)
         ; return con }
     mk_doc con_name = text "Constructor" <+> ppr con_name