    module GHC.Core.UsageEnv
      ( Usage(..)
      , UsageEnv
      , addUE
      , addUsage
      , bottomUE
      , deleteUE
      , lookupUE
      , popUE
      , scaleUE
      , scaleUsage
      , supUE
      , supUEs
      , singleUsageUE
      , zeroUE
      ) where

    import Data.Foldable
    import GHC.Prelude
    import GHC.Core.Multiplicity
    import GHC.Types.Var
    import GHC.Types.Name
    import GHC.Types.Name.Env
    import GHC.Utils.Outputable
    import GHC.Utils.Panic

    --
    -- * Usage environments
    --

    -- The typechecker and the linter output usage environments. See Note [Usages]
    -- in Multiplicity. Every absent name being considered to map to 'Zero' of
    -- 'Bottom' depending on a flag. See Note [Zero as a usage] in Multiplicity, see
    -- Note [Bottom as a usage] in Multiplicity.

    data Usage = Zero | Bottom | MUsage Mult

    instance Outputable Usage where
      ppr :: Usage -> SDoc
    ppr Usage
    Zero = String -> SDoc
    forall doc. IsLine doc => String -> doc
    text String
    "0"
      ppr Usage
    Bottom = String -> SDoc
    forall doc. IsLine doc => String -> doc
    text String
    "Bottom"
      ppr (MUsage Mult
    x) = Mult -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Mult
    x

    addUsage :: Usage -> Usage -> Usage
    addUsage :: Usage -> Usage -> Usage
    addUsage Usage
    Zero Usage
    x = Usage
    x
    addUsage Usage
    x Usage
    Zero = Usage
    x
    addUsage Usage
    Bottom Usage
    x = Usage
    x
    addUsage Usage
    x Usage
    Bottom = Usage
    x
    addUsage (MUsage Mult
    x) (MUsage Mult
    y) = Mult -> Usage
    MUsage (Mult -> Usage) -> Mult -> Usage
    forall a b. (a -> b) -> a -> b
    $ Mult -> Mult -> Mult
    mkMultAdd Mult
    x Mult
    y

    scaleUsage :: Mult -> Usage -> Usage
    scaleUsage :: Mult -> Usage -> Usage
    scaleUsage Mult
    OneTy Usage
    Bottom     = Usage
    Bottom
    scaleUsage Mult
    _     Usage
    Zero       = Usage
    Zero
    scaleUsage Mult
    x     Usage
    Bottom     = Mult -> Usage
    MUsage Mult
    x
    scaleUsage Mult
    x     (MUsage Mult
    y) = Mult -> Usage
    MUsage (Mult -> Usage) -> Mult -> Usage
    forall a b. (a -> b) -> a -> b
    $ Mult -> Mult -> Mult
    mkMultMul Mult
    x Mult
    y

    -- For now, we use extra multiplicity Bottom for empty case.
    data UsageEnv = UsageEnv !(NameEnv Mult) Bool

    -- | Record a single usage of an Id, i.e. {n: 1}
    -- Exception: We do not record external names (both GlobalIds and top-level LocalIds)
    -- because they're not relevant to linearity checking.
    singleUsageUE :: Id -> UsageEnv
    singleUsageUE :: Id -> UsageEnv
    singleUsageUE Id
    x | Name -> Bool
    isExternalName Name
    n = UsageEnv
    zeroUE
                    | Bool
    otherwise = NameEnv Mult -> Bool -> UsageEnv
    UsageEnv (Name -> Mult -> NameEnv Mult
    forall a. Name -> a -> NameEnv a
    unitNameEnv Name
    n Mult
    OneTy) Bool
    False
      where n :: Name
    n = Id -> Name
    forall a. NamedThing a => a -> Name
    getName Id
    x

    zeroUE, bottomUE :: UsageEnv
    zeroUE :: UsageEnv
    zeroUE = NameEnv Mult -> Bool -> UsageEnv
    UsageEnv NameEnv Mult
    forall a. NameEnv a
    emptyNameEnv Bool
    False

    bottomUE :: UsageEnv
    bottomUE = NameEnv Mult -> Bool -> UsageEnv
    UsageEnv NameEnv Mult
    forall a. NameEnv a
    emptyNameEnv Bool
    True

    addUE :: UsageEnv -> UsageEnv -> UsageEnv
    addUE :: UsageEnv -> UsageEnv -> UsageEnv
    addUE (UsageEnv NameEnv Mult
    e1 Bool
    b1) (UsageEnv NameEnv Mult
    e2 Bool
    b2) =
      NameEnv Mult -> Bool -> UsageEnv
    UsageEnv ((Mult -> Mult -> Mult)
    -> NameEnv Mult -> NameEnv Mult -> NameEnv Mult
    forall a. (a -> a -> a) -> NameEnv a -> NameEnv a -> NameEnv a
    plusNameEnv_C Mult -> Mult -> Mult
    mkMultAdd NameEnv Mult
    e1 NameEnv Mult
    e2) (Bool
    b1 Bool -> Bool -> Bool
    || Bool
    b2)

    scaleUE :: Mult -> UsageEnv -> UsageEnv
    scaleUE :: Mult -> UsageEnv -> UsageEnv
    scaleUE Mult
    OneTy UsageEnv
    ue = UsageEnv
    ue
    scaleUE Mult
    w (UsageEnv NameEnv Mult
    e Bool
    _) =
      NameEnv Mult -> Bool -> UsageEnv
    UsageEnv ((Mult -> Mult) -> NameEnv Mult -> NameEnv Mult
    forall elt1 elt2. (elt1 -> elt2) -> NameEnv elt1 -> NameEnv elt2
    mapNameEnv (Mult -> Mult -> Mult
    mkMultMul Mult
    w) NameEnv Mult
    e) Bool
    False

    supUE :: UsageEnv -> UsageEnv -> UsageEnv
    supUE :: UsageEnv -> UsageEnv -> UsageEnv
    supUE (UsageEnv NameEnv Mult
    e1 Bool
    False) (UsageEnv NameEnv Mult
    e2 Bool
    False) =
      NameEnv Mult -> Bool -> UsageEnv
    UsageEnv ((Mult -> Mult -> Mult)
    -> NameEnv Mult -> Mult -> NameEnv Mult -> Mult -> NameEnv Mult
    forall a.
    (a -> a -> a) -> NameEnv a -> a -> NameEnv a -> a -> NameEnv a
    plusNameEnv_CD Mult -> Mult -> Mult
    mkMultSup NameEnv Mult
    e1 Mult
    ManyTy NameEnv Mult
    e2 Mult
    ManyTy) Bool
    False
    supUE (UsageEnv NameEnv Mult
    e1 Bool
    b1) (UsageEnv NameEnv Mult
    e2 Bool
    b2) = NameEnv Mult -> Bool -> UsageEnv
    UsageEnv ((Maybe Mult -> Maybe Mult -> Mult)
    -> NameEnv Mult -> NameEnv Mult -> NameEnv Mult
    forall a.
    (Maybe a -> Maybe a -> a) -> NameEnv a -> NameEnv a -> NameEnv a
    plusNameEnv_CD2 Maybe Mult -> Maybe Mult -> Mult
    combineUsage NameEnv Mult
    e1 NameEnv Mult
    e2) (Bool
    b1 Bool -> Bool -> Bool
    && Bool
    b2)
       where combineUsage :: Maybe Mult -> Maybe Mult -> Mult
    combineUsage (Just Mult
    x) (Just Mult
    y) = Mult -> Mult -> Mult
    mkMultSup Mult
    x Mult
    y
             combineUsage Maybe Mult
    Nothing  (Just Mult
    x) | Bool
    b1        = Mult
    x
                                            | Bool
    otherwise = Mult
    ManyTy
             combineUsage (Just Mult
    x) Maybe Mult
    Nothing  | Bool
    b2        = Mult
    x
                                            | Bool
    otherwise = Mult
    ManyTy
             combineUsage Maybe Mult
    Nothing  Maybe Mult
    Nothing  = String -> SDoc -> Mult
    forall a. HasCallStack => String -> SDoc -> a
    pprPanic String
    "supUE" (NameEnv Mult -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr NameEnv Mult
    e1 SDoc -> SDoc -> SDoc
    forall doc. IsLine doc => doc -> doc -> doc
    <+> NameEnv Mult -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr NameEnv Mult
    e2)
    -- Note: If you are changing this logic, check 'mkMultSup' in Multiplicity as well.

    supUEs :: [UsageEnv] -> UsageEnv
    supUEs :: [UsageEnv] -> UsageEnv
    supUEs = (UsageEnv -> UsageEnv -> UsageEnv)
    -> UsageEnv -> [UsageEnv] -> UsageEnv
    forall a b. (a -> b -> b) -> b -> [a] -> b
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr UsageEnv -> UsageEnv -> UsageEnv
    supUE UsageEnv
    bottomUE

    deleteUE :: NamedThing n => UsageEnv -> n -> UsageEnv
    deleteUE :: forall n. NamedThing n => UsageEnv -> n -> UsageEnv
    deleteUE (UsageEnv NameEnv Mult
    e Bool
    b) n
    x = NameEnv Mult -> Bool -> UsageEnv
    UsageEnv (NameEnv Mult -> Name -> NameEnv Mult
    forall a. NameEnv a -> Name -> NameEnv a
    delFromNameEnv NameEnv Mult
    e (n -> Name
    forall a. NamedThing a => a -> Name
    getName n
    x)) Bool
    b

    -- | |lookupUE x env| returns the multiplicity assigned to |x| in |env|, if |x| is not
    -- bound in |env|, then returns |Zero| or |Bottom|.
    lookupUE :: NamedThing n => UsageEnv -> n -> Usage
    lookupUE :: forall n. NamedThing n => UsageEnv -> n -> Usage
    lookupUE (UsageEnv NameEnv Mult
    e Bool
    has_bottom) n
    x =
      case NameEnv Mult -> Name -> Maybe Mult
    forall a. NameEnv a -> Name -> Maybe a
    lookupNameEnv NameEnv Mult
    e (n -> Name
    forall a. NamedThing a => a -> Name
    getName n
    x) of
        Just Mult
    w  -> Mult -> Usage
    MUsage Mult
    w
        Maybe Mult
    Nothing -> if Bool
    has_bottom then Usage
    Bottom else Usage
    Zero

    popUE :: NamedThing n => UsageEnv -> n -> (Usage, UsageEnv)
    popUE :: forall n. NamedThing n => UsageEnv -> n -> (Usage, UsageEnv)
    popUE UsageEnv
    ue n
    x = (UsageEnv -> n -> Usage
    forall n. NamedThing n => UsageEnv -> n -> Usage
    lookupUE UsageEnv
    ue n
    x, UsageEnv -> n -> UsageEnv
    forall n. NamedThing n => UsageEnv -> n -> UsageEnv
    deleteUE UsageEnv
    ue n
    x)

    instance Outputable UsageEnv where
      ppr :: UsageEnv -> SDoc
    ppr (UsageEnv NameEnv Mult
    ne Bool
    b) = String -> SDoc
    forall doc. IsLine doc => String -> doc
    text String
    "UsageEnv:" SDoc -> SDoc -> SDoc
    forall doc. IsLine doc => doc -> doc -> doc
    <+> NameEnv Mult -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr NameEnv Mult
    ne SDoc -> SDoc -> SDoc
    forall doc. IsLine doc => doc -> doc -> doc
    <+> Bool -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Bool
    b