Optimise whole module exports

We directly build up the correct AvailInfos rather than generating
lots of singleton instances and combining them with expensive calls to
unionLists.

There are two other small changes.

* Pushed the nubAvails call into the explicit export list
  branch as we construct them correctly and uniquely ourselves.
* fix_faminst only needs to check the first element of the export
  list as we maintain the (yucky) invariant that the parent is the
  first thing in it.

Reviewers: simonpj, austin, bgamari

Reviewed By: simonpj, bgamari

Subscribers: simonpj, thomie, niteria

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

GHC Trac Issues: #12754

2 files changed, 65 insertions, 15 deletions

diff --git a/compiler/basicTypes/RdrName.hs b/compiler/basicTypes/RdrName.hs
index 40c152b664..d60522f41c 100644
--- a/compiler/basicTypes/RdrName.hs
+++ b/compiler/basicTypes/RdrName.hs
@@ -52,6 +52,7 @@ module RdrName (
         -- * GlobalRdrElts
         gresFromAvails, gresFromAvail, localGREsFromAvail, availFromGRE,
         greUsedRdrName, greRdrNames, greSrcSpan, greQualModName,
+        gresToAvailInfo,
 
         -- ** Global 'RdrName' mapping elements: 'GlobalRdrElt', 'Provenance', 'ImportSpec'
         GlobalRdrElt(..), isLocalGRE, isRecFldGRE, greLabel,
@@ -77,9 +78,10 @@ import Unique
 import UniqFM
 import Util
 import StaticFlags( opt_PprStyle_Debug )
+import NameEnv
 
 import Data.Data
-import Data.List( sortBy )
+import Data.List( sortBy, foldl', nub )
 
 {-
 ************************************************************************
@@ -453,7 +455,7 @@ data GlobalRdrElt
         , gre_par  :: Parent
         , gre_lcl :: Bool          -- ^ True <=> the thing was defined locally
         , gre_imp :: [ImportSpec]  -- ^ In scope through these imports
-    } deriving Data
+    } deriving (Data, Eq)
          -- INVARIANT: either gre_lcl = True or gre_imp is non-empty
          -- See Note [GlobalRdrElt provenance]
 
@@ -687,15 +689,60 @@ mkParent _ (Avail _)           = NoParent
 mkParent n (AvailTC m _ _) | n == m    = NoParent
                          | otherwise = ParentIs m
 
+greParentName :: GlobalRdrElt -> Maybe Name
+greParentName gre = case gre_par gre of
+                      NoParent -> Nothing
+                      ParentIs n -> Just n
+                      FldParent n _ -> Just n
+
+-- | Takes a list of distinct GREs and folds them
+-- into AvailInfos. This is more efficient than mapping each individual
+-- GRE to an AvailInfo and the folding using `plusAvail` but needs the
+-- uniqueness assumption.
+gresToAvailInfo :: [GlobalRdrElt] -> [AvailInfo]
+gresToAvailInfo gres
+  = ASSERT( nub gres == gres ) nameEnvElts avail_env
+  where
+    avail_env :: NameEnv AvailInfo -- keyed by the parent
+    avail_env = foldl' add emptyNameEnv gres
+
+    add :: NameEnv AvailInfo -> GlobalRdrElt -> NameEnv AvailInfo
+    add env gre = extendNameEnv_Acc comb availFromGRE env
+                    (fromMaybe (gre_name gre)
+                               (greParentName gre)) gre
+
+      where
+        -- We want to insert the child `k` into a list of children but
+        -- need to maintain the invariant that the parent is first.
+        --
+        -- We also use the invariant that `k` is not already in `ns`.
+        insertChildIntoChildren :: Name -> [Name] -> Name -> [Name]
+        insertChildIntoChildren _ [] k = [k]
+        insertChildIntoChildren p (n:ns) k
+          | p == k = k:n:ns
+          | otherwise = n:k:ns
+
+        comb :: GlobalRdrElt -> AvailInfo -> AvailInfo
+        comb _ (Avail n) = Avail n -- Duplicated name
+        comb gre (AvailTC m ns fls) =
+          let n = gre_name gre
+          in case gre_par gre of
+              NoParent -> AvailTC m (n:ns) fls -- Not sure this ever happens
+              ParentIs {} -> AvailTC m (insertChildIntoChildren m ns n) fls
+              FldParent _ mb_lbl ->  AvailTC m ns (mkFieldLabel n mb_lbl : fls)
+
 availFromGRE :: GlobalRdrElt -> AvailInfo
 availFromGRE (GRE { gre_name = me, gre_par = parent })
   = case parent of
       ParentIs p                  -> AvailTC p [me] []
       NoParent   | isTyConName me -> AvailTC me [me] []
                  | otherwise      -> avail   me
-      FldParent p mb_lbl -> AvailTC p [] [fld]
+      FldParent p mb_lbl -> AvailTC p [] [mkFieldLabel me mb_lbl]
         where
-         fld = case mb_lbl of
+
+mkFieldLabel :: Name -> Maybe FastString -> FieldLabel
+mkFieldLabel me mb_lbl =
+          case mb_lbl of
                  Nothing  -> FieldLabel { flLabel = occNameFS (nameOccName me)
                                         , flIsOverloaded = False
                                         , flSelector = me }
diff --git a/compiler/typecheck/TcRnExports.hs b/compiler/typecheck/TcRnExports.hs
index 14c151b62b..35ff65f92a 100644
--- a/compiler/typecheck/TcRnExports.hs
+++ b/compiler/typecheck/TcRnExports.hs
@@ -136,10 +136,9 @@ tcRnExports explicit_mod exports
                         -- ToDo: the 'noLoc' here is unhelpful if 'main'
                         --       turns out to be out of scope
 
-        ; (rn_exports, avails) <- exports_from_avail real_exports rdr_env imports this_mod
-        ; traceRn "Exported Avails" (ppr avails)
-        ; let final_avails = nubAvails avails    -- Combine families
-              final_ns     = availsToNameSetWithSelectors final_avails
+        ; (rn_exports, final_avails)
+            <- exports_from_avail real_exports rdr_env imports this_mod
+        ; let final_ns     = availsToNameSetWithSelectors final_avails
 
         ; traceRn "rnExports: Exports:" (ppr final_avails)
 
@@ -164,9 +163,9 @@ exports_from_avail Nothing rdr_env _imports _this_mod
    -- The same as (module M) where M is the current module name,
    -- so that's how we handle it, except we also export the data family
    -- when a data instance is exported.
-  = let avails = [ fix_faminst $ availFromGRE gre
-                 | gre <- globalRdrEnvElts rdr_env
-                 , isLocalGRE gre ]
+  = let avails =
+          map fix_faminst . gresToAvailInfo
+            . filter isLocalGRE . globalRdrEnvElts $ rdr_env
     in return (Nothing, avails)
   where
     -- #11164: when we define a data instance
@@ -174,9 +173,12 @@ exports_from_avail Nothing rdr_env _imports _this_mod
     -- Even though we don't check whether this is actually a data family
     -- only data families can locally define subordinate things (`ns` here)
     -- without locally defining (and instead importing) the parent (`n`)
-    fix_faminst (AvailTC n ns flds)
-      | n `notElem` ns
-      = AvailTC n (n:ns) flds
+    fix_faminst (AvailTC n ns flds) =
+      let new_ns =
+            case ns of
+              [] -> [n]
+              (p:_) -> if p == n then ns else n:ns
+      in AvailTC n new_ns flds
 
     fix_faminst avail = avail
 
@@ -184,7 +186,8 @@ exports_from_avail Nothing rdr_env _imports _this_mod
 exports_from_avail (Just (L _ rdr_items)) rdr_env imports this_mod
   = do ExportAccum ie_names _ exports
         <-  checkNoErrs $ foldAndRecoverM do_litem emptyExportAccum rdr_items
-       return (Just ie_names, exports)
+       let final_exports = nubAvails exports -- Combine families
+       return (Just ie_names, final_exports)
   where
     do_litem :: ExportAccum -> LIE RdrName -> RnM ExportAccum
     do_litem acc lie = setSrcSpan (getLoc lie) (exports_from_item acc lie)