content: {-# LANGUAGE GeneralisedNewtypeDeriving #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# LANGUAGE PatternSynonyms #-}
  
  -- | 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.idCprSig'.
  module GHC.Types.Cpr (
      Cpr (ConCpr), topCpr, botCpr, flatConCpr, asConCpr,
      CprType (..), topCprType, botCprType, flatConCprType,
      lubCprType, applyCprTy, abstractCprTy, trimCprTy,
      UnpackConFieldsResult (..), unpackConFieldsCpr,
      CprSig (..), topCprSig, isTopCprSig, mkCprSigForArity, mkCprSig,
      seqCprSig, prependArgsCprSig
    ) where
  
  import GHC.Prelude
  
  import GHC.Core.DataCon
  import GHC.Types.Basic
  import GHC.Utils.Binary
  import GHC.Utils.Misc
  import GHC.Utils.Outputable
  import GHC.Utils.Panic
  
  --
  -- * Cpr
  --
  
  data Cpr
    = BotCpr
    | ConCpr_ !ConTag ![Cpr]
    -- ^ The number of field Cprs equals 'dataConRepArity'.
    -- If all of them are top, better use 'FlatConCpr', as ensured by the pattern
    -- synonym 'ConCpr'.
    | FlatConCpr !ConTag
    -- ^ @FlatConCpr tag@ is an efficient encoding for @'ConCpr_' tag [TopCpr..]@.
    -- Purely for compiler perf. Can be constructed with 'ConCpr'.
    | TopCpr
    deriving Eq
  
  pattern ConCpr :: ConTag -> [Cpr] -> Cpr
  pattern ConCpr t cs <- ConCpr_ t cs where
    ConCpr t cs
      | all (== TopCpr) cs = FlatConCpr t
      | otherwise          = ConCpr_ t cs
  {-# COMPLETE BotCpr, TopCpr, FlatConCpr, ConCpr #-}
  
  viewConTag :: Cpr -> Maybe ConTag
  viewConTag (FlatConCpr t) = Just t
  viewConTag (ConCpr t _)   = Just t
  viewConTag _              = Nothing
  {-# INLINE viewConTag #-}
  
  lubCpr :: Cpr -> Cpr -> Cpr
  lubCpr BotCpr      cpr     = cpr
  lubCpr cpr         BotCpr  = cpr
  lubCpr (FlatConCpr t1) (viewConTag -> Just t2)
    | t1 == t2 = FlatConCpr t1
  lubCpr (viewConTag -> Just t1) (FlatConCpr t2)
    | t1 == t2 = FlatConCpr t2
  lubCpr (ConCpr t1 cs1) (ConCpr t2 cs2)
    | t1 == t2 = ConCpr t1 (lubFieldCprs cs1 cs2)
  lubCpr _           _       = TopCpr
  
  lubFieldCprs :: [Cpr] -> [Cpr] -> [Cpr]
  lubFieldCprs as bs
    | as `equalLength` bs = zipWith lubCpr as bs
    | otherwise           = []
  
  topCpr :: Cpr
  topCpr = TopCpr
  
  botCpr :: Cpr
  botCpr = BotCpr
  
  flatConCpr :: ConTag -> Cpr
  flatConCpr t = FlatConCpr t
  
  trimCpr :: Cpr -> Cpr
  trimCpr BotCpr = botCpr
  trimCpr _      = topCpr
  
  asConCpr :: Cpr -> Maybe (ConTag, [Cpr])
  asConCpr (ConCpr t cs)  = Just (t, cs)
  asConCpr (FlatConCpr t) = Just (t, [])
  asConCpr TopCpr         = Nothing
  asConCpr BotCpr         = Nothing
  
  seqCpr :: Cpr -> ()
  seqCpr (ConCpr _ cs) = foldr (seq . seqCpr) () cs
  seqCpr _             = ()
  
  --
  -- * 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  :: !Cpr   -- ^ 'Cpr' 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
  
  flatConCprType :: ConTag -> CprType
  flatConCprType con_tag = CprType { ct_arty = 0, ct_cpr = flatConCpr 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
  
  trimCprTy :: CprType -> CprType
  trimCprTy (CprType arty res) = CprType arty (trimCpr res)
  
  -- | The result of 'unpackConFieldsCpr'.
  data UnpackConFieldsResult
    = AllFieldsSame !Cpr
    | ForeachField ![Cpr]
  
  -- | Unpacks a 'ConCpr'-shaped 'Cpr' and returns the field 'Cpr's wrapped in a
  -- 'ForeachField'. Otherwise, it returns 'AllFieldsSame' with the appropriate
  -- 'Cpr' to assume for each field.
  --
  -- The use of 'UnpackConFieldsResult' allows O(1) space for the common,
  -- non-'ConCpr' case.
  unpackConFieldsCpr :: DataCon -> Cpr -> UnpackConFieldsResult
  unpackConFieldsCpr dc (ConCpr t cs)
    | t == dataConTag dc, cs `lengthIs` dataConRepArity dc
    = ForeachField cs
  unpackConFieldsCpr _  BotCpr = AllFieldsSame BotCpr
  unpackConFieldsCpr _  _      = AllFieldsSame TopCpr
  {-# INLINE unpackConFieldsCpr #-}
  
  seqCprTy :: CprType -> ()
  seqCprTy (CprType _ cpr) = seqCpr cpr
  
  -- | The arity of the wrapped 'CprType' is the arity at which it is safe
  -- to unleash. See Note [Understanding DmdType and DmdSig] 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 DmdSig] in
  -- "GHC.Types.Demand"
  mkCprSigForArity :: Arity -> CprType -> CprSig
  mkCprSigForArity arty ty@(CprType n _)
    | arty /= n = topCprSig -- Trim on arity mismatch
    | otherwise = CprSig ty
  
  topCprSig :: CprSig
  topCprSig = CprSig topCprType
  
  isTopCprSig :: CprSig -> Bool
  isTopCprSig (CprSig ty) = ct_cpr ty == topCpr
  
  mkCprSig :: Arity -> Cpr -> CprSig
  mkCprSig arty cpr = CprSig (CprType arty cpr)
  
  seqCprSig :: CprSig -> ()
  seqCprSig (CprSig ty) = seqCprTy ty
  
  prependArgsCprSig :: Arity -> CprSig -> CprSig
  -- ^ Add extra value args to CprSig
  prependArgsCprSig n_extra cpr_sig@(CprSig (CprType arity cpr))
    | n_extra == 0 = cpr_sig
    | otherwise    = assertPpr (n_extra > 0) (ppr n_extra) $
                     CprSig (CprType (arity + n_extra) cpr)
  
  -- | BNF:
  --
  -- > cpr ::= ''                               -- TopCpr
  -- >      |  n                                -- FlatConCpr n
  -- >      |  n '(' cpr1 ',' cpr2 ',' ... ')'  -- ConCpr n [cpr1,cpr2,...]
  -- >      |  'b'                              -- BotCpr
  --
  -- Examples:
  --   * `f x = f x` has result CPR `b`
  --   * `1(1,)` is a valid (nested) 'Cpr' denotation for `(I# 42#, f 42)`.
  instance Outputable Cpr where
    ppr TopCpr         = empty
    ppr (FlatConCpr n) = int n
    ppr (ConCpr n cs)  = int n <> parens (pprWithCommas ppr cs)
    ppr BotCpr         = char 'b'
  
  -- | BNF:
  --
  -- > cpr_ty ::= cpr               -- short form if arty == 0
  -- >         |  '\' arty '.' cpr  -- if arty > 0
  --
  -- Examples:
  --   * `f x y z = f x y z` has denotation `\3.b`
  --   * `g !x = (x+1, x+2)` has denotation `\1.1(1,1)`.
  instance Outputable CprType where
    ppr (CprType arty res)
      | 0 <- arty = ppr res
      | otherwise = char '\\' <> ppr arty <> char '.' <> ppr res
  
  -- | Only print the CPR result
  instance Outputable CprSig where
    ppr (CprSig ty) = ppr (ct_cpr ty)
  
  instance Binary Cpr where
    put_ bh TopCpr         = putByte bh 0
    put_ bh BotCpr         = putByte bh 1
    put_ bh (FlatConCpr n) = putByte bh 2 *> put_ bh n
    put_ bh (ConCpr n cs)  = putByte bh 3 *> put_ bh n *> put_ bh cs
    get  bh = do
      h <- getByte bh
      case h of
        0 -> return TopCpr
        1 -> return BotCpr
        2 -> FlatConCpr <$> get bh
        3 -> ConCpr <$> get bh <*> get bh
        _ -> pprPanic "Binary Cpr: Invalid tag" (int (fromIntegral h))
  
  instance Binary CprType where
    put_ bh (CprType arty cpr) = put_ bh arty *> put_ bh cpr
    get  bh                    = CprType <$> get bh <*> get bh
  