GHC/Core/UsageEnv.hs

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

    import Data.Foldable
    import GHC.Prelude
    import GHC.Core.Multiplicity
    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 Zero = text "0"
      ppr Bottom = text "Bottom"
      ppr (MUsage x) = ppr x

    addUsage :: Usage -> Usage -> Usage
    addUsage Zero x = x
    addUsage x Zero = x
    addUsage Bottom x = x
    addUsage x Bottom = x
    addUsage (MUsage x) (MUsage y) = MUsage $ mkMultAdd x y

    scaleUsage :: Mult -> Usage -> Usage
    scaleUsage One Bottom     = Bottom
    scaleUsage _   Zero       = Zero
    scaleUsage x   Bottom     = MUsage x
    scaleUsage x   (MUsage y) = MUsage $ mkMultMul x y

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

    unitUE :: NamedThing n => n -> Mult -> UsageEnv
    unitUE x w = UsageEnv (unitNameEnv (getName x) w) False

    zeroUE, bottomUE :: UsageEnv
    zeroUE = UsageEnv emptyNameEnv False

    bottomUE = UsageEnv emptyNameEnv True

    addUE :: UsageEnv -> UsageEnv -> UsageEnv
    addUE (UsageEnv e1 b1) (UsageEnv e2 b2) =
      UsageEnv (plusNameEnv_C mkMultAdd e1 e2) (b1 || b2)

    scaleUE :: Mult -> UsageEnv -> UsageEnv
    scaleUE One ue = ue
    scaleUE w (UsageEnv e _) =
      UsageEnv (mapNameEnv (mkMultMul w) e) False

    supUE :: UsageEnv -> UsageEnv -> UsageEnv
    supUE (UsageEnv e1 False) (UsageEnv e2 False) =
      UsageEnv (plusNameEnv_CD mkMultSup e1 Many e2 Many) False
    supUE (UsageEnv e1 b1) (UsageEnv e2 b2) = UsageEnv (plusNameEnv_CD2 combineUsage e1 e2) (b1 && b2)
       where combineUsage (Just x) (Just y) = mkMultSup x y
             combineUsage Nothing  (Just x) | b1        = x
                                            | otherwise = Many
             combineUsage (Just x) Nothing  | b2        = x
                                            | otherwise = Many
             combineUsage Nothing  Nothing  = pprPanic "supUE" (ppr e1 <+> ppr e2)
    -- Note: If you are changing this logic, check 'mkMultSup' in Multiplicity as well.

    supUEs :: [UsageEnv] -> UsageEnv
    supUEs = foldr supUE bottomUE

    deleteUE :: NamedThing n => UsageEnv -> n -> UsageEnv
    deleteUE (UsageEnv e b) x = UsageEnv (delFromNameEnv e (getName x)) 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 (UsageEnv e has_bottom) x =
      case lookupNameEnv e (getName x) of
        Just w  -> MUsage w
        Nothing -> if has_bottom then Bottom else Zero

    instance Outputable UsageEnv where
      ppr (UsageEnv ne b) = text "UsageEnv:" <+> ppr ne <+> ppr b