Extra.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE RankNTypes #-}

module Liquid.GHC.API.Extra (
    module StableModule
  , ApiComment(..)
  , addNoInlinePragmasToBinds
  , apiComments
  , apiCommentsParsedSource
  , dataConSig
  , directImports
  , fsToUnitId
  , isPatErrorAlt
  , minus_RDR
  , qualifiedNameFS
  , renderWithStyle
  , showPprQualified
  , showPprDebug
  , showSDocQualified
  , splitDollarApp
  , strictNothing
  , thisPackage
  , tyConRealArity
  , untick
  , withTimingWallClock
  ) where

import Control.Monad
import Control.Monad.IO.Class
import GHC.Conc (getAllocationCounter)
import Debug.Trace
import GHC.Clock (getMonotonicTimeNSec)
import           Liquid.GHC.API.StableModule      as StableModule
import GHC hiding (modInfoLookupName)
import Data.Data (Data, gmapQr, gmapT)
import Data.Generics (extQ, extT)
import Data.Foldable                  (asum)
import Data.List                      (sortOn)
import qualified Data.Map as Map
import GHC.Builtin.Names ( dollarIdKey, minusName )
import GHC.Core                       as Ghc
import GHC.Core.Coercion              as Ghc
import GHC.Core.DataCon               as Ghc
import GHC.Core.Make                  (pAT_ERROR_ID)
import GHC.Core.Type                  as Ghc hiding (typeKind , isPredTy, extendCvSubst, linear)
import GHC.Data.FastString            as Ghc
import GHC.Data.Maybe
import qualified GHC.Data.Strict
import GHC.Driver.Session             as Ghc
import GHC.Tc.Types
import GHC.Types.Id
import GHC.Types.Basic
import GHC.Types.Name                 (isSystemName, nameModule_maybe, occNameFS)
import GHC.Types.Name.Reader          (nameRdrName)
import GHC.Types.SrcLoc               as Ghc
import GHC.Types.Unique               (getUnique, hasKey)

import GHC.Utils.Error                as Ghc
import GHC.Utils.Logger               as Ghc
import GHC.Utils.Outputable           as Ghc hiding ((<>))
import qualified GHC.Utils.Outputable as Ghc

import GHC.Unit.Module

-- 'fsToUnitId' is gone in GHC 9, but we can bring code it in terms of 'fsToUnit' and 'toUnitId'.
fsToUnitId :: FastString -> UnitId
fsToUnitId = toUnitId . fsToUnit

thisPackage :: DynFlags -> UnitId
thisPackage = homeUnitId_

-- See NOTE [tyConRealArity].
tyConRealArity :: TyCon -> Int
tyConRealArity tc = go 0 (tyConKind tc)
  where
    go :: Int -> Kind -> Int
    go !acc k =
      case asum [fmap (\(_, _, _, c) -> c) (splitFunTy_maybe k), fmap snd (splitForAllTyCoVar_maybe k)] of
        Nothing -> acc
        Just ks -> go (acc + 1) ks

renderWithStyle :: DynFlags -> SDoc -> PprStyle -> String
renderWithStyle dynflags sdoc style = Ghc.renderWithContext (Ghc.initSDocContext dynflags style) sdoc

-- This function is gone in GHC 9.
dataConSig :: DataCon -> ([TyCoVar], ThetaType, [Type], Type)
dataConSig dc
  = (dataConUnivAndExTyCoVars dc, dataConTheta dc, map irrelevantMult $ dataConOrigArgTys dc, dataConOrigResTy dc)

-- | Extracts the direct imports of a module.
directImports :: TcGblEnv -> [Module]
directImports = Map.keys . imp_mods . tcg_imports

-- | Abstraction of 'EpaComment'.
data ApiComment
  = ApiLineComment String
  | ApiBlockComment String
  deriving (Eq, Show)

-- | Extract top-level comments from a module.
apiComments :: HsParsedModule -> [Ghc.Located ApiComment]
apiComments = apiCommentsParsedSource . hpm_module

apiCommentsParsedSource :: Located (HsModule GhcPs) -> [Ghc.Located ApiComment]
apiCommentsParsedSource ps =
    let hs = unLoc ps
        go :: forall a. Data a => a -> [LEpaComment]
        go = gmapQr (++) [] go `extQ` (id @[LEpaComment])
     in Data.List.sortOn (spanToLineColumn . getLoc) $
          mapMaybe (tokComment . toRealSrc) $ go hs
  where
    tokComment (L sp (EpaComment (EpaLineComment s) _)) = Just (L sp (ApiLineComment s))
    tokComment (L sp (EpaComment (EpaBlockComment s) _)) = Just (L sp (ApiBlockComment s))
    tokComment _ = Nothing

    -- TODO: take into account anchor_op, which only matters if the source was
    -- pre-processed by an exact-print-aware tool.
    toRealSrc (L a e) = L (RealSrcSpan (epaLocationRealSrcSpan a) strictNothing) e

    spanToLineColumn =
      fmap (\s -> (srcSpanStartLine s, srcSpanStartCol s)) . srcSpanToRealSrcSpan

-- | Adds NOINLINE pragmas to all bindings in the module.
--
-- This prevents the simple optimizer from inlining such bindings which might
-- have specs that would otherwise be left dangling.
--
-- https://gitlab.haskell.org/ghc/ghc/-/issues/24386
--
addNoInlinePragmasToBinds :: TcGblEnv -> TcGblEnv
addNoInlinePragmasToBinds tcg = tcg{ tcg_binds = go (tcg_binds tcg) }
  where
    go :: forall a. Data a => a -> a
    go = gmapT $ go `extT` markHsBind

    -- Mark all user-originating `Id` binders as `NOINLINE`.
    markHsBind :: HsBind GhcTc -> HsBind GhcTc
    markHsBind = \case
        bind@VarBind{ var_id = var, var_rhs = rhs } -> bind{ var_id = markId var, var_rhs = go rhs }
        bind@FunBind{ fun_id = var, fun_matches = matches } -> bind{ fun_id = markId <$> var, fun_matches = go matches }
        bind@PatBind{ pat_lhs = lhs, pat_rhs = rhs } -> bind{ pat_lhs = markPat <$> lhs, pat_rhs = go rhs }
        PatSynBind{} -> error "markNoInline: unexpected PatSynBind, should have been eliminated by the typechecker"
        XHsBindsLR absBinds -> XHsBindsLR (markAbsBinds absBinds)

    markPat :: Pat GhcTc -> Pat GhcTc
    markPat = \case
        VarPat ext var -> VarPat ext (markId <$> var)
        pat -> gmapT (id `extT` markPat) pat

    markId :: Id -> Id
    markId var = var `setInlinePragma` neverInlinePragma

    -- The AbsBinds come from the GHC typechecker to handle polymorphism,
    -- overloading, and recursion, so those don't correspond directly to
    -- user-written `Id`s except for those in @abs_exports@. For instance,
    -- @tests/pos/Map0.hs@ would fail if Ids in @abs_exports@ are not marked.
    --
    -- See
    -- https://github.com/ucsd-progsys/liquidhaskell/issues/2257 for more
    -- context.
    markAbsBinds :: AbsBinds -> AbsBinds
    markAbsBinds absBinds0@AbsBinds{ abs_binds = binds, abs_exports = exports }  =
        absBinds0
          { abs_binds = fmap skipFirstHsBind <$> binds
          , abs_exports = map markABE exports
          }
      where
        skipFirstHsBind :: HsBind GhcTc -> HsBind GhcTc
        skipFirstHsBind = \case
          XHsBindsLR absBinds -> XHsBindsLR (markAbsBinds absBinds)
          b -> gmapT go b

        markABE :: ABExport -> ABExport
        markABE abe@ABE{ abe_poly = poly
                       , abe_mono = mono } = abe
          { abe_poly = markId poly
          , abe_mono = markId mono }

-- | Tells if a case alternative calls to patError
isPatErrorAlt :: CoreAlt -> Bool
isPatErrorAlt (Alt _ _ exprCoreBndr) = hasPatErrorCall exprCoreBndr
  where
   hasPatErrorCall :: CoreExpr -> Bool
   -- auto generated undefined case: (\_ -> (patError @levity @type "error message")) void
   -- Type arguments are erased before calling isUndefined
   hasPatErrorCall (App e _)
     | Var x <- unTick e = x == pAT_ERROR_ID
     | otherwise = hasPatErrorCall e
   -- another auto generated undefined case:
   -- let lqanf_... = patError "error message") in case lqanf_... of {}
   hasPatErrorCall (Let (NonRec x e) ec)
     | Case e0 _ _ [] <- unTick ec
     , Var v <- unTick e0
     , x == v = hasPatErrorCall e
   hasPatErrorCall (Case e _ _ _) = hasPatErrorCall e
   hasPatErrorCall (Let _ e) = hasPatErrorCall e
   hasPatErrorCall (Tick _ e) = hasPatErrorCall e
   -- otherwise
   hasPatErrorCall _ = False

   unTick (Tick _ e) = unTick e
   unTick e = e


qualifiedNameFS :: Name -> FastString
qualifiedNameFS n = concatFS [modFS, occFS, uniqFS]
  where
  modFS = case nameModule_maybe n of
            Nothing -> fsLit ""
            Just m  -> concatFS [moduleNameFS (moduleName m), fsLit "."]

  occFS = occNameFS (getOccName n)
  uniqFS
    | isSystemName n
    = concatFS [fsLit "_",  fsLit (showPprQualified (getUnique n))]
    | otherwise
    = fsLit ""

-- Variants of Outputable functions which now require DynFlags!
showPprQualified :: Outputable a => a -> String
showPprQualified = showSDocQualified . ppr

showSDocQualified :: Ghc.SDoc -> String
showSDocQualified = Ghc.renderWithContext ctx
  where
    style = Ghc.mkUserStyle myQualify Ghc.AllTheWay
    ctx = Ghc.defaultSDocContext { sdocStyle = style }

myQualify :: Ghc.NamePprCtx
myQualify = Ghc.neverQualify { Ghc.queryQualifyName = Ghc.alwaysQualifyNames }
-- { Ghc.queryQualifyName = \_ _ -> Ghc.NameNotInScope1 }

showPprDebug :: Outputable a => a -> String
showPprDebug = showSDocDebug . ppr

showSDocDebug :: Ghc.SDoc -> String
showSDocDebug = Ghc.renderWithContext ctx
  where
    style = Ghc.mkUserStyle myQualify Ghc.AllTheWay
    ctx = Ghc.defaultSDocContext {
        sdocStyle = style
      , sdocPprDebug = True
      }

strictNothing :: GHC.Data.Strict.Maybe a
strictNothing = GHC.Data.Strict.Nothing

splitDollarApp :: CoreExpr -> Maybe (CoreExpr, CoreExpr)
splitDollarApp e
     -- matches `$ t1 t2 t3 t4 f a`
     | App e1 a  <- untick e
     , App e2 f  <- untick e1
     , App e3 t4 <- untick e2
     , App e4 t3 <- untick e3
     , App e5 t2 <- untick e4
     , App d t1  <- untick e5
     , Var v     <- untick d
     , v `hasKey` dollarIdKey
     , Type _    <- untick t1
     , Type _    <- untick t2
     , Type _    <- untick t3
     , Type _    <- untick t4
     = Just (f, a)
     | otherwise
     = Nothing

untick :: CoreExpr -> CoreExpr
untick (Tick _ e) = untick e
untick e = e

minus_RDR :: RdrName
minus_RDR = nameRdrName minusName

-- | A version of 'withTiming' that uses wall clock time instead of CPU time.
--
-- This version is copied and modified from GHC's 'GHC.Utils.Error.withTiming'
withTimingWallClock :: MonadIO m
           => Logger
           -> SDoc         -- ^ The name of the phase
           -> (a -> ())    -- ^ A function to force the result
                           -- (often either @const ()@ or 'rnf')
           -> m a          -- ^ The body of the phase to be timed
           -> m a
withTimingWallClock logger what force_result action =
    if logVerbAtLeast logger 2 || logHasDumpFlag logger Opt_D_dump_timings
    then do when printTimingsNotDumpToFile $ liftIO $
              logInfo logger $ withPprStyle defaultUserStyle $
                text "***" <+> what Ghc.<> colon
            let ctx = log_default_user_context (logFlags logger)
            alloc0 <- liftIO getAllocationCounter
            !start <- liftIO getMonotonicTimeNSec
            eventBegins ctx what
            recordAllocs alloc0
            !r <- action
            () <- pure $ force_result r
            eventEnds ctx what
            !end <- liftIO getMonotonicTimeNSec
            alloc1 <- liftIO getAllocationCounter
            recordAllocs alloc1
            -- recall that allocation counter counts down
            let alloc = alloc0 - alloc1
                time = (end - start) `div` 1000000

            when (logVerbAtLeast logger 2 && printTimingsNotDumpToFile)
                $ liftIO $ logInfo logger $ withPprStyle defaultUserStyle
                    (text "!!!" <+> what Ghc.<> colon <+> text "finished in"
                     <+> word64 time
                     <+> text "milliseconds"
                     Ghc.<> comma
                     <+> text "allocated"
                     <+> doublePrec 3 (realToFrac alloc / 1024 / 1024)
                     <+> text "megabytes")

            liftIO $ putDumpFileMaybe logger Opt_D_dump_timings "" FormatText
                    $ text $ showSDocOneLine ctx
                    $ hsep [ what Ghc.<> colon
                           , text "alloc=" Ghc.<> ppr alloc
                           , text "time=" Ghc.<> word64 time
                           ]
            pure r
     else action

    where -- Avoid both printing to console and dumping to a file (#20316).
          printTimingsNotDumpToFile =
            not (log_dump_to_file (logFlags logger))

          recordAllocs alloc =
            liftIO $ traceMarkerIO $ "GHC:allocs:" ++ show alloc

          eventBegins ctx w = do
            let doc = eventBeginsDoc ctx w
            liftIO $ traceMarkerIO doc
            liftIO $ traceEventIO doc

          eventEnds ctx w = do
            let doc = eventEndsDoc ctx w
            liftIO $ traceMarkerIO doc
            liftIO $ traceEventIO doc

          eventBeginsDoc ctx w = showSDocOneLine ctx $ text "GHC:started:" <+> w
          eventEndsDoc   ctx w = showSDocOneLine ctx $ text "GHC:finished:" <+> w