Batch merge

5 files changed, 156 insertions, 60 deletions

diff --git a/compiler/hieFile/HieAst.hs b/compiler/hieFile/HieAst.hs
index 401b861e30..b6b5f0ccb7 100644
--- a/compiler/hieFile/HieAst.hs
+++ b/compiler/hieFile/HieAst.hs
@@ -28,9 +28,11 @@ import HscTypes
 import Module                     ( ModuleName, ml_hs_file )
 import MonadUtils                 ( concatMapM, liftIO )
 import Name                       ( Name, nameSrcSpan, setNameLoc )
+import NameEnv                    ( NameEnv, emptyNameEnv, extendNameEnv, lookupNameEnv )
 import SrcLoc
-import TcHsSyn                    ( hsPatType )
-import Type                       ( Type )
+import TcHsSyn                    ( hsLitType, hsPatType )
+import Type                       ( mkFunTys, Type )
+import TysWiredIn                 ( mkListTy, mkSumTy )
 import Var                        ( Id, Var, setVarName, varName, varType )
 
 import HieTypes
@@ -60,11 +62,11 @@ We don't care about the distinction between mono and poly bindings,
 so we replace all occurrences of the mono name with the poly name.
 -}
 newtype HieState = HieState
-  { name_remapping :: M.Map Name Id
+  { name_remapping :: NameEnv Id
   }
 
 initState :: HieState
-initState = HieState M.empty
+initState = HieState emptyNameEnv
 
 class ModifyState a where -- See Note [Name Remapping]
   addSubstitution :: a -> a -> HieState -> HieState
@@ -74,7 +76,7 @@ instance ModifyState Name where
 
 instance ModifyState Id where
   addSubstitution mono poly hs =
-    hs{name_remapping = M.insert (varName mono) poly (name_remapping hs)}
+    hs{name_remapping = extendNameEnv (name_remapping hs) (varName mono) poly}
 
 modifyState :: ModifyState (IdP p) => [ABExport p] -> HieState -> HieState
 modifyState = foldr go id
@@ -377,7 +379,9 @@ instance ToHie (Context (Located Var)) where
       C context (L (RealSrcSpan span) name')
         -> do
         m <- asks name_remapping
-        let name = M.findWithDefault name' (varName name') m
+        let name = case lookupNameEnv m (varName name') of
+              Just var -> var
+              Nothing-> name'
         pure
           [Node
             (NodeInfo S.empty [] $
@@ -392,7 +396,7 @@ instance ToHie (Context (Located Name)) where
   toHie c = case c of
       C context (L (RealSrcSpan span) name') -> do
         m <- asks name_remapping
-        let name = case M.lookup name' m of
+        let name = case lookupNameEnv m name' of
               Just var -> varName var
               Nothing -> name'
         pure
@@ -432,13 +436,67 @@ instance HasType (LPat GhcTc) where
 instance HasType (LHsExpr GhcRn) where
   getTypeNode (L spn e) = makeNode e spn
 
+-- | This instance tries to construct 'HieAST' nodes which include the type of
+-- the expression. It is not yet possible to do this efficiently for all
+-- expression forms, so we skip filling in the type for those inputs.
+--
+-- 'HsApp', for example, doesn't have any type information available directly on
+-- the node. Our next recourse would be to desugar it into a 'CoreExpr' then
+-- query the type of that. Yet both the desugaring call and the type query both
+-- involve recursive calls to the function and argument! This is particularly
+-- problematic when you realize that the HIE traversal will eventually visit
+-- those nodes too and ask for their types again.
+--
+-- Since the above is quite costly, we just skip cases where computing the
+-- expression's type is going to be expensive.
+--
+-- See #16233
 instance HasType (LHsExpr GhcTc) where
-  getTypeNode e@(L spn e') = lift $ do
-    hs_env <- Hsc $ \e w -> return (e,w)
-    (_,mbe) <- liftIO $ deSugarExpr hs_env e
-    case mbe of
-      Just te -> makeTypeNode e' spn (exprType te)
-      Nothing -> makeNode e' spn
+  getTypeNode e@(L spn e') = lift $
+    -- Some expression forms have their type immediately available
+    let tyOpt = case e' of
+          HsLit _ l -> Just (hsLitType l)
+          HsOverLit _ o -> Just (overLitType o)
+
+          HsLam     _ (MG { mg_ext = groupTy }) -> Just (matchGroupType groupTy)
+          HsLamCase _ (MG { mg_ext = groupTy }) -> Just (matchGroupType groupTy)
+          HsCase _  _ (MG { mg_ext = groupTy }) -> Just (mg_res_ty groupTy)
+
+          ExplicitList  ty _ _   -> Just (mkListTy ty)
+          ExplicitSum   ty _ _ _ -> Just (mkSumTy ty)
+          HsDo          ty _ _   -> Just ty
+          HsMultiIf     ty _     -> Just ty
+
+          _ -> Nothing
+
+    in
+    case tyOpt of
+      _ | skipDesugaring e' -> fallback
+        | otherwise -> do
+            hs_env <- Hsc $ \e w -> return (e,w)
+            (_,mbe) <- liftIO $ deSugarExpr hs_env e
+            maybe fallback (makeTypeNode e' spn . exprType) mbe
+    where
+      fallback = makeNode e' spn
+
+      matchGroupType :: MatchGroupTc -> Type
+      matchGroupType (MatchGroupTc args res) = mkFunTys args res
+
+      -- | Skip desugaring of these expressions for performance reasons.
+      --
+      -- See impact on Haddock output (esp. missing type annotations or links)
+      -- before marking more things here as 'False'. See impact on Haddock
+      -- performance before marking more things as 'True'.
+      skipDesugaring :: HsExpr a -> Bool
+      skipDesugaring e = case e of
+        HsVar{}        -> False
+        HsUnboundVar{} -> False
+        HsConLikeOut{} -> False
+        HsRecFld{}     -> False
+        HsOverLabel{}  -> False
+        HsIPVar{}      -> False
+        HsWrap{}       -> False
+        _              -> True
 
 instance ( ToHie (Context (Located (IdP a)))
          , ToHie (MatchGroup a (LHsExpr a))
diff --git a/compiler/nativeGen/X86/CodeGen.hs b/compiler/nativeGen/X86/CodeGen.hs
index 37080b990e..9591c42ede 100644
--- a/compiler/nativeGen/X86/CodeGen.hs
+++ b/compiler/nativeGen/X86/CodeGen.hs
@@ -2045,25 +2045,37 @@ genCCall dflags is32Bit (PrimTarget (MO_Clz width)) dest_regs@[dst] args@[src] b
 
   | otherwise = do
     code_src <- getAnyReg src
-    src_r <- getNewRegNat format
-    tmp_r <- getNewRegNat format
     let dst_r = getRegisterReg platform False (CmmLocal dst)
-
-    -- The following insn sequence makes sure 'clz 0' has a defined value.
-    -- starting with Haswell, one could use the LZCNT insn instead.
-    return $ code_src src_r `appOL` toOL
-             ([ MOVZxL  II8    (OpReg src_r) (OpReg src_r) | width == W8 ] ++
-              [ BSR     format (OpReg src_r) tmp_r
-              , MOV     II32   (OpImm (ImmInt (2*bw-1))) (OpReg dst_r)
-              , CMOV NE format (OpReg tmp_r) dst_r
-              , XOR     format (OpImm (ImmInt (bw-1))) (OpReg dst_r)
-              ]) -- NB: We don't need to zero-extend the result for the
-                 -- W8/W16 cases because the 'MOV' insn already
-                 -- took care of implicitly clearing the upper bits
+    if isBmi2Enabled dflags
+        then do
+            src_r <- getNewRegNat (intFormat width)
+            return $ appOL (code_src src_r) $ case width of
+                W8 -> toOL
+                    [ MOVZxL II8  (OpReg src_r)       (OpReg src_r) -- zero-extend to 32 bit
+                    , LZCNT  II32 (OpReg src_r)       dst_r         -- lzcnt with extra 24 zeros
+                    , SUB    II32 (OpImm (ImmInt 24)) (OpReg dst_r) -- compensate for extra zeros
+                    ]
+                W16 -> toOL
+                    [ LZCNT  II16 (OpReg src_r) dst_r
+                    , MOVZxL II16 (OpReg dst_r) (OpReg dst_r) -- zero-extend from 16 bit
+                    ]
+                _ -> unitOL (LZCNT (intFormat width) (OpReg src_r) dst_r)
+        else do
+            let format = if width == W8 then II16 else intFormat width
+            src_r <- getNewRegNat format
+            tmp_r <- getNewRegNat format
+            return $ code_src src_r `appOL` toOL
+                     ([ MOVZxL  II8    (OpReg src_r) (OpReg src_r) | width == W8 ] ++
+                      [ BSR     format (OpReg src_r) tmp_r
+                      , MOV     II32   (OpImm (ImmInt (2*bw-1))) (OpReg dst_r)
+                      , CMOV NE format (OpReg tmp_r) dst_r
+                      , XOR     format (OpImm (ImmInt (bw-1))) (OpReg dst_r)
+                      ]) -- NB: We don't need to zero-extend the result for the
+                         -- W8/W16 cases because the 'MOV' insn already
+                         -- took care of implicitly clearing the upper bits
   where
     bw = widthInBits width
     platform = targetPlatform dflags
-    format = if width == W8 then II16 else intFormat width
     lbl = mkCmmCodeLabel primUnitId (fsLit (clzLabel width))
 
 genCCall dflags is32Bit (PrimTarget (MO_Ctz width)) [dst] [src] bid
@@ -2073,6 +2085,7 @@ genCCall dflags is32Bit (PrimTarget (MO_Ctz width)) [dst] [src] bid
           dst_r   = getRegisterReg platform False (CmmLocal dst)
       lbl1 <- getBlockIdNat
       lbl2 <- getBlockIdNat
+      let format = if width == W8 then II16 else intFormat width
       tmp_r <- getNewRegNat format
 
       -- New CFG Edges:
@@ -2109,24 +2122,38 @@ genCCall dflags is32Bit (PrimTarget (MO_Ctz width)) [dst] [src] bid
 
   | otherwise = do
     code_src <- getAnyReg src
-    src_r <- getNewRegNat format
-    tmp_r <- getNewRegNat format
     let dst_r = getRegisterReg platform False (CmmLocal dst)
 
-    -- The following insn sequence makes sure 'ctz 0' has a defined value.
-    -- starting with Haswell, one could use the TZCNT insn instead.
-    return $ code_src src_r `appOL` toOL
-             ([ MOVZxL  II8    (OpReg src_r) (OpReg src_r) | width == W8 ] ++
-              [ BSF     format (OpReg src_r) tmp_r
-              , MOV     II32   (OpImm (ImmInt bw)) (OpReg dst_r)
-              , CMOV NE format (OpReg tmp_r) dst_r
-              ]) -- NB: We don't need to zero-extend the result for the
-                 -- W8/W16 cases because the 'MOV' insn already
-                 -- took care of implicitly clearing the upper bits
+    if isBmi2Enabled dflags
+    then do
+        src_r <- getNewRegNat (intFormat width)
+        return $ appOL (code_src src_r) $ case width of
+            W8 -> toOL
+                [ OR    II32 (OpImm (ImmInt 0xFFFFFF00)) (OpReg src_r)
+                , TZCNT II32 (OpReg src_r)        dst_r
+                ]
+            W16 -> toOL
+                [ TZCNT  II16 (OpReg src_r) dst_r
+                , MOVZxL II16 (OpReg dst_r) (OpReg dst_r)
+                ]
+            _ -> unitOL $ TZCNT (intFormat width) (OpReg src_r) dst_r
+    else do
+        -- The following insn sequence makes sure 'ctz 0' has a defined value.
+        -- starting with Haswell, one could use the TZCNT insn instead.
+        let format = if width == W8 then II16 else intFormat width
+        src_r <- getNewRegNat format
+        tmp_r <- getNewRegNat format
+        return $ code_src src_r `appOL` toOL
+                 ([ MOVZxL  II8    (OpReg src_r) (OpReg src_r) | width == W8 ] ++
+                  [ BSF     format (OpReg src_r) tmp_r
+                  , MOV     II32   (OpImm (ImmInt bw)) (OpReg dst_r)
+                  , CMOV NE format (OpReg tmp_r) dst_r
+                  ]) -- NB: We don't need to zero-extend the result for the
+                     -- W8/W16 cases because the 'MOV' insn already
+                     -- took care of implicitly clearing the upper bits
   where
     bw = widthInBits width
     platform = targetPlatform dflags
-    format = if width == W8 then II16 else intFormat width
 
 genCCall dflags is32Bit (PrimTarget (MO_UF_Conv width)) dest_regs args bid = do
     targetExpr <- cmmMakeDynamicReference dflags
diff --git a/compiler/nativeGen/X86/Instr.hs b/compiler/nativeGen/X86/Instr.hs
index c47e1fae83..5e790e481e 100644
--- a/compiler/nativeGen/X86/Instr.hs
+++ b/compiler/nativeGen/X86/Instr.hs
@@ -342,6 +342,8 @@ data Instr
 
     -- bit counting instructions
         | POPCNT      Format Operand Reg -- [SSE4.2] count number of bits set to 1
+        | LZCNT       Format Operand Reg -- [BMI2] count number of leading zeros
+        | TZCNT       Format Operand Reg -- [BMI2] count number of trailing zeros
         | BSF         Format Operand Reg -- bit scan forward
         | BSR         Format Operand Reg -- bit scan reverse
 
@@ -471,6 +473,8 @@ x86_regUsageOfInstr platform instr
     DELTA   _           -> noUsage
 
     POPCNT _ src dst -> mkRU (use_R src []) [dst]
+    LZCNT  _ src dst -> mkRU (use_R src []) [dst]
+    TZCNT  _ src dst -> mkRU (use_R src []) [dst]
     BSF    _ src dst -> mkRU (use_R src []) [dst]
     BSR    _ src dst -> mkRU (use_R src []) [dst]
 
@@ -653,6 +657,8 @@ x86_patchRegsOfInstr instr env
     CLTD _              -> instr
 
     POPCNT fmt src dst -> POPCNT fmt (patchOp src) (env dst)
+    LZCNT  fmt src dst -> LZCNT  fmt (patchOp src) (env dst)
+    TZCNT  fmt src dst -> TZCNT  fmt (patchOp src) (env dst)
     PDEP   fmt src mask dst -> PDEP   fmt (patchOp src) (patchOp mask) (env dst)
     PEXT   fmt src mask dst -> PEXT   fmt (patchOp src) (patchOp mask) (env dst)
     BSF    fmt src dst -> BSF    fmt (patchOp src) (env dst)
diff --git a/compiler/nativeGen/X86/Ppr.hs b/compiler/nativeGen/X86/Ppr.hs
index bf449d044e..075bb26337 100644
--- a/compiler/nativeGen/X86/Ppr.hs
+++ b/compiler/nativeGen/X86/Ppr.hs
@@ -693,6 +693,8 @@ pprInstr (XOR FF64 src dst) = pprOpOp (sLit "xorpd") FF64 src dst
 pprInstr (XOR format src dst) = pprFormatOpOp (sLit "xor")  format src dst
 
 pprInstr (POPCNT format src dst) = pprOpOp (sLit "popcnt") format src (OpReg dst)
+pprInstr (LZCNT format src dst)  = pprOpOp (sLit "lzcnt")  format src (OpReg dst)
+pprInstr (TZCNT format src dst)  = pprOpOp (sLit "tzcnt")  format src (OpReg dst)
 pprInstr (BSF format src dst)    = pprOpOp (sLit "bsf")    format src (OpReg dst)
 pprInstr (BSR format src dst)    = pprOpOp (sLit "bsr")    format src (OpReg dst)
 
diff --git a/compiler/parser/RdrHsSyn.hs b/compiler/parser/RdrHsSyn.hs
index c1777759da..45fc5a0972 100644
--- a/compiler/parser/RdrHsSyn.hs
+++ b/compiler/parser/RdrHsSyn.hs
@@ -151,10 +151,11 @@ mkClassDecl loc (dL->L _ (mcxt, tycl_hdr)) fds where_cls
   = do { (binds, sigs, ats, at_insts, _, docs) <- cvBindsAndSigs where_cls
        ; let cxt = fromMaybe (noLoc []) mcxt
        ; (cls, tparams, fixity, ann) <- checkTyClHdr True tycl_hdr
-       ; mapM_ (\a -> a loc) ann -- Add any API Annotations to the top SrcSpan
-       ; tyvars <- checkTyVarsP (text "class") whereDots cls tparams
-       ; (at_defs, anns) <- fmap unzip $ mapM (eitherToP . mkATDefault) at_insts
-       ; sequence_ anns
+       ; addAnnsAt loc ann -- Add any API Annotations to the top SrcSpan
+       ; (tyvars,annst) <- checkTyVarsP (text "class") whereDots cls tparams
+       ; addAnnsAt loc annst -- Add any API Annotations to the top SrcSpan
+       ; (at_defs, annsi) <- mapAndUnzipM (eitherToP . mkATDefault) at_insts
+       ; sequence_ annsi
        ; return (cL loc (ClassDecl { tcdCExt = noExt, tcdCtxt = cxt
                                    , tcdLName = cls, tcdTyVars = tyvars
                                    , tcdFixity = fixity
@@ -186,7 +187,7 @@ mkATDefault (dL->L loc (TyFamInstDecl { tfid_eqn = HsIB { hsib_body = e }}))
                                     , feqn_pats   = tvs
                                     , feqn_fixity = fixity
                                     , feqn_rhs    = rhs })
-           ; pure (f, anns) }
+           ; pure (f, addAnnsAt loc anns) }
 mkATDefault (dL->L _ (TyFamInstDecl (HsIB _ (XFamEqn _)))) = panic "mkATDefault"
 mkATDefault (dL->L _ (TyFamInstDecl (XHsImplicitBndrs _))) = panic "mkATDefault"
 mkATDefault _ = panic "mkATDefault: Impossible Match"
@@ -203,8 +204,9 @@ mkTyData :: SrcSpan
 mkTyData loc new_or_data cType (dL->L _ (mcxt, tycl_hdr))
          ksig data_cons maybe_deriv
   = do { (tc, tparams, fixity, ann) <- checkTyClHdr False tycl_hdr
-       ; mapM_ (\a -> a loc) ann -- Add any API Annotations to the top SrcSpan
-       ; tyvars <- checkTyVarsP (ppr new_or_data) equalsDots tc tparams
+       ; addAnnsAt loc ann -- Add any API Annotations to the top SrcSpan
+       ; (tyvars, anns) <- checkTyVarsP (ppr new_or_data) equalsDots tc tparams
+       ; addAnnsAt loc anns -- Add any API Annotations to the top SrcSpan
        ; defn <- mkDataDefn new_or_data cType mcxt ksig data_cons maybe_deriv
        ; return (cL loc (DataDecl { tcdDExt = noExt,
                                     tcdLName = tc, tcdTyVars = tyvars,
@@ -235,8 +237,9 @@ mkTySynonym :: SrcSpan
             -> P (LTyClDecl GhcPs)
 mkTySynonym loc lhs rhs
   = do { (tc, tparams, fixity, ann) <- checkTyClHdr False lhs
-       ; mapM_ (\a -> a loc) ann -- Add any API Annotations to the top SrcSpan
-       ; tyvars <- checkTyVarsP (text "type") equalsDots tc tparams
+       ; addAnnsAt loc ann -- Add any API Annotations to the top SrcSpan
+       ; (tyvars, anns) <- checkTyVarsP (text "type") equalsDots tc tparams
+       ; addAnnsAt loc anns -- Add any API Annotations to the top SrcSpan
        ; return (cL loc (SynDecl { tcdSExt = noExt
                                  , tcdLName = tc, tcdTyVars = tyvars
                                  , tcdFixity = fixity
@@ -293,8 +296,9 @@ mkFamDecl :: SrcSpan
           -> P (LTyClDecl GhcPs)
 mkFamDecl loc info lhs ksig injAnn
   = do { (tc, tparams, fixity, ann) <- checkTyClHdr False lhs
-       ; mapM_ (\a -> a loc) ann -- Add any API Annotations to the top SrcSpan
-       ; tyvars <- checkTyVarsP (ppr info) equals_or_where tc tparams
+       ; addAnnsAt loc ann -- Add any API Annotations to the top SrcSpan
+       ; (tyvars, anns) <- checkTyVarsP (ppr info) equals_or_where tc tparams
+       ; addAnnsAt loc anns -- Add any API Annotations to the top SrcSpan
        ; return (cL loc (FamDecl noExt (FamilyDecl
                                            { fdExt       = noExt
                                            , fdInfo      = info, fdLName = tc
@@ -804,13 +808,11 @@ really doesn't matter!
 -}
 
 checkTyVarsP :: SDoc -> SDoc -> Located RdrName -> [LHsTypeArg GhcPs]
-             -> P (LHsQTyVars GhcPs)
+             -> P (LHsQTyVars GhcPs, [AddAnn])
 -- Same as checkTyVars, but in the P monad
 checkTyVarsP pp_what equals_or_where tc tparms
   = do { let checkedTvs = checkTyVars pp_what equals_or_where tc tparms
-       ; (tvs, anns) <- eitherToP checkedTvs
-       ; anns
-       ; pure tvs }
+       ; eitherToP checkedTvs }
 
 eitherToP :: Either (SrcSpan, SDoc) a -> P a
 -- Adapts the Either monad to the P monad
@@ -820,14 +822,14 @@ eitherToP (Right thing)     = return thing
 checkTyVars :: SDoc -> SDoc -> Located RdrName -> [LHsTypeArg GhcPs]
             -> Either (SrcSpan, SDoc)
                       ( LHsQTyVars GhcPs  -- the synthesized type variables
-                      , P () )            -- action which adds annotations
+                      , [AddAnn] )        -- action which adds annotations
 -- ^ Check whether the given list of type parameters are all type variables
 -- (possibly with a kind signature).
 -- We use the Either monad because it's also called (via 'mkATDefault') from
 -- "Convert".
 checkTyVars pp_what equals_or_where tc tparms
   = do { (tvs, anns) <- fmap unzip $ mapM check tparms
-       ; return (mkHsQTvs tvs, sequence_ anns) }
+       ; return (mkHsQTvs tvs, concat anns) }
   where
     check (HsTypeArg ki@(L loc _)) = Left (loc,
                                       vcat [ text "Unexpected type application" <+>
@@ -839,14 +841,15 @@ checkTyVars pp_what equals_or_where tc tparms
                            <+> text "declaration for" <+> quotes (ppr tc)])
         -- Keep around an action for adjusting the annotations of extra parens
     chkParens :: [AddAnn] -> LHsType GhcPs
-              -> Either (SrcSpan, SDoc) (LHsTyVarBndr GhcPs, P ())
+              -> Either (SrcSpan, SDoc) (LHsTyVarBndr GhcPs, [AddAnn])
     chkParens acc (dL->L l (HsParTy _ ty)) = chkParens (mkParensApiAnn l
                                                         ++ acc) ty
     chkParens acc ty = case chk ty of
       Left err -> Left err
-      Right tv@(dL->L l _) -> Right (tv, addAnnsAt l (reverse acc))
+      Right tv -> Right (tv, reverse acc)
 
         -- Check that the name space is correct!
+    chk :: LHsType GhcPs -> Either (SrcSpan, SDoc) (LHsTyVarBndr GhcPs)
     chk (dL->L l (HsKindSig _ (dL->L lv (HsTyVar _ _ (dL->L _ tv))) k))
         | isRdrTyVar tv    = return (cL l (KindedTyVar noExt (cL lv tv) k))
     chk (dL->L l (HsTyVar _ _ (dL->L ltv tv)))