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 :: Usage -> SDoc
ppr Usage
Zero = String -> SDoc
text String
"0"
  ppr Usage
Bottom = String -> SDoc
text String
"Bottom"
  ppr (MUsage Mult
x) = 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 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
One 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 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

unitUE :: NamedThing n => n -> Mult -> UsageEnv
unitUE :: forall n. NamedThing n => n -> Mult -> UsageEnv
unitUE n
x Mult
w = NameEnv Mult -> Bool -> UsageEnv
UsageEnv (forall a. Name -> a -> NameEnv a
unitNameEnv (forall a. NamedThing a => a -> Name
getName n
x) Mult
w) Bool
False

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

bottomUE :: UsageEnv
bottomUE = NameEnv Mult -> Bool -> UsageEnv
UsageEnv 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 (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
One UsageEnv
ue = UsageEnv
ue
scaleUE Mult
w (UsageEnv NameEnv Mult
e Bool
_) =
  NameEnv Mult -> Bool -> UsageEnv
UsageEnv (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 (forall a.
(a -> a -> a) -> NameEnv a -> a -> NameEnv a -> a -> NameEnv a
plusNameEnv_CD Mult -> Mult -> Mult
mkMultSup NameEnv Mult
e1 Mult
Many NameEnv Mult
e2 Mult
Many) Bool
False
supUE (UsageEnv NameEnv Mult
e1 Bool
b1) (UsageEnv NameEnv Mult
e2 Bool
b2) = NameEnv Mult -> Bool -> UsageEnv
UsageEnv (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
Many
         combineUsage (Just Mult
x) Maybe Mult
Nothing  | Bool
b2        = Mult
x
                                        | Bool
otherwise = Mult
Many
         combineUsage Maybe Mult
Nothing  Maybe Mult
Nothing  = forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"supUE" (forall a. Outputable a => a -> SDoc
ppr NameEnv Mult
e1 SDoc -> SDoc -> 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 = 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 (forall a. NameEnv a -> Name -> NameEnv a
delFromNameEnv NameEnv Mult
e (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 forall a. NameEnv a -> Name -> Maybe a
lookupNameEnv NameEnv Mult
e (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

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