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