GHC/Types/Cpr.hs

    {-# LANGUAGE GeneralisedNewtypeDeriving #-}
    -- | Types for the Constructed Product Result lattice. "GHC.Core.Opt.CprAnal" and "GHC.Core.Opt.WorkWrap.Utils"
    -- are its primary customers via 'GHC.Types.Id.idCprInfo'.
    module GHC.Types.Cpr (
        CprResult, topCpr, botCpr, conCpr, asConCpr,
        CprType (..), topCprType, botCprType, conCprType,
        lubCprType, applyCprTy, abstractCprTy, ensureCprTyArity, trimCprTy,
        CprSig (..), topCprSig, mkCprSigForArity, mkCprSig, seqCprSig
      ) where

    import GHC.Prelude

    import GHC.Types.Basic
    import GHC.Utils.Outputable
    import GHC.Utils.Binary

    --
    -- * CprResult
    --

    -- | The constructed product result lattice.
    --
    -- @
    --                    NoCPR
    --                      |
    --                 ConCPR ConTag
    --                      |
    --                    BotCPR
    -- @
    data CprResult = NoCPR          -- ^ Top of the lattice
                   | ConCPR !ConTag -- ^ Returns a constructor from a data type
                   | BotCPR         -- ^ Bottom of the lattice
                   deriving( Eq, Show )

    lubCpr :: CprResult -> CprResult -> CprResult
    lubCpr (ConCPR t1) (ConCPR t2)
      | t1 == t2               = ConCPR t1
    lubCpr BotCPR      cpr     = cpr
    lubCpr cpr         BotCPR  = cpr
    lubCpr _           _       = NoCPR

    topCpr :: CprResult
    topCpr = NoCPR

    botCpr :: CprResult
    botCpr = BotCPR

    conCpr :: ConTag -> CprResult
    conCpr = ConCPR

    trimCpr :: CprResult -> CprResult
    trimCpr ConCPR{} = NoCPR
    trimCpr cpr      = cpr

    asConCpr :: CprResult -> Maybe ConTag
    asConCpr (ConCPR t)  = Just t
    asConCpr NoCPR       = Nothing
    asConCpr BotCPR      = Nothing

    --
    -- * CprType
    --

    -- | The abstract domain \(A_t\) from the original 'CPR for Haskell' paper.
    data CprType
      = CprType
      { ct_arty :: !Arity     -- ^ Number of value arguments the denoted expression
                              --   eats before returning the 'ct_cpr'
      , ct_cpr  :: !CprResult -- ^ 'CprResult' eventually unleashed when applied to
                              --   'ct_arty' arguments
      }

    instance Eq CprType where
      a == b =  ct_cpr a == ct_cpr b
             && (ct_arty a == ct_arty b || ct_cpr a == topCpr)

    topCprType :: CprType
    topCprType = CprType 0 topCpr

    botCprType :: CprType
    botCprType = CprType 0 botCpr -- TODO: Figure out if arity 0 does what we want... Yes it does: arity zero means we may unleash it under any number of incoming arguments

    conCprType :: ConTag -> CprType
    conCprType con_tag = CprType 0 (conCpr con_tag)

    lubCprType :: CprType -> CprType -> CprType
    lubCprType ty1@(CprType n1 cpr1) ty2@(CprType n2 cpr2)
      -- The arity of bottom CPR types can be extended arbitrarily.
      | cpr1 == botCpr && n1 <= n2 = ty2
      | cpr2 == botCpr && n2 <= n1 = ty1
      -- There might be non-bottom CPR types with mismatching arities.
      -- Consider test DmdAnalGADTs. We want to return top in these cases.
      | n1 == n2                   = CprType n1 (lubCpr cpr1 cpr2)
      | otherwise                  = topCprType

    applyCprTy :: CprType -> Arity -> CprType
    applyCprTy (CprType n res) k
      | n >= k        = CprType (n-k) res
      | res == botCpr = botCprType
      | otherwise     = topCprType

    abstractCprTy :: CprType -> CprType
    abstractCprTy (CprType n res)
      | res == topCpr = topCprType
      | otherwise     = CprType (n+1) res

    ensureCprTyArity :: Arity -> CprType -> CprType
    ensureCprTyArity n ty@(CprType m _)
      | n == m    = ty
      | otherwise = topCprType

    trimCprTy :: CprType -> CprType
    trimCprTy (CprType arty res) = CprType arty (trimCpr res)

    -- | The arity of the wrapped 'CprType' is the arity at which it is safe
    -- to unleash. See Note [Understanding DmdType and StrictSig] in "GHC.Types.Demand"
    newtype CprSig = CprSig { getCprSig :: CprType }
      deriving (Eq, Binary)

    -- | Turns a 'CprType' computed for the particular 'Arity' into a 'CprSig'
    -- unleashable at that arity. See Note [Understanding DmdType and StrictSig] in
    -- "GHC.Types.Demand"
    mkCprSigForArity :: Arity -> CprType -> CprSig
    mkCprSigForArity arty ty = CprSig (ensureCprTyArity arty ty)

    topCprSig :: CprSig
    topCprSig = CprSig topCprType

    mkCprSig :: Arity -> CprResult -> CprSig
    mkCprSig arty cpr = CprSig (CprType arty cpr)

    seqCprSig :: CprSig -> ()
    seqCprSig sig = sig `seq` ()

    instance Outputable CprResult where
      ppr NoCPR        = empty
      ppr (ConCPR n)   = char 'm' <> int n
      ppr BotCPR       = char 'b'

    instance Outputable CprType where
      ppr (CprType arty res) = ppr arty <> ppr res

    -- | Only print the CPR result
    instance Outputable CprSig where
      ppr (CprSig ty) = ppr (ct_cpr ty)

    instance Binary CprResult where
      put_ bh (ConCPR n)   = do { putByte bh 0; put_ bh n }
      put_ bh NoCPR        = putByte bh 1
      put_ bh BotCPR       = putByte bh 2

      get  bh = do
              h <- getByte bh
              case h of
                0 -> do { n <- get bh; return (ConCPR n) }
                1 -> return NoCPR
                _ -> return BotCPR

    instance Binary CprType where
      put_ bh (CprType arty cpr) = do
        put_ bh arty
        put_ bh cpr
      get  bh = CprType <$> get bh <*> get bh