{-# LANGUAGE FlexibleContexts #-}

-- | Functions for computing variable usage inside terms and types.
module PlutusIR.Analysis.Usages (termUsages, typeUsages, Usages, getUsageCount, allUsed) where

import PlutusPrelude ((<^>))

import PlutusIR
import PlutusIR.Subst

import PlutusCore qualified as PLC
import PlutusCore.Name.Unique qualified as PLC

import Control.Lens

import Data.MultiSet qualified as MSet
import Data.MultiSet.Lens
import Data.Set qualified as Set

type Usages = MSet.MultiSet PLC.Unique

-- | Get the usage count of @n@.
getUsageCount :: (PLC.HasUnique n unique) => n -> Usages -> Int
getUsageCount :: forall n unique. HasUnique n unique => n -> Usages -> Int
getUsageCount n
n = Unique -> Usages -> Int
forall a. Ord a => a -> MultiSet a -> Int
MSet.occur (n
n n -> Getting Unique n Unique -> Unique
forall s a. s -> Getting a s a -> a
^. (unique -> Const Unique unique) -> n -> Const Unique n
forall a unique. HasUnique a unique => Lens' a unique
Lens' n unique
PLC.unique ((unique -> Const Unique unique) -> n -> Const Unique n)
-> ((Unique -> Const Unique Unique)
    -> unique -> Const Unique unique)
-> Getting Unique n Unique
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Unique -> Const Unique Unique) -> unique -> Const Unique unique
forall s t a b. (Coercible s a, Coercible t b) => Iso s t a b
Iso unique unique Unique Unique
coerced)

-- | Get a set of @n@s which are used at least once.
allUsed :: Usages -> Set.Set PLC.Unique
allUsed :: Usages -> Set Unique
allUsed = Usages -> Set Unique
forall a. MultiSet a -> Set a
MSet.toSet

termUsages ::
  (PLC.HasUnique name PLC.TermUnique, PLC.HasUnique tyname PLC.TypeUnique) =>
  Term tyname name uni fun a ->
  Usages
termUsages :: forall name tyname (uni :: * -> *) fun a.
(HasUnique name TermUnique, HasUnique tyname TypeUnique) =>
Term tyname name uni fun a -> Usages
termUsages = Getting Usages (Term tyname name uni fun a) Unique
-> Term tyname name uni fun a -> Usages
forall a s. Getting (MultiSet a) s a -> s -> MultiSet a
multiSetOf ((name -> f name)
-> Term tyname name uni fun a -> f (Term tyname name uni fun a)
forall tyname name (uni :: * -> *) fun ann (f :: * -> *).
(Contravariant f, Applicative f) =>
(name -> f name)
-> Term tyname name uni fun ann -> f (Term tyname name uni fun ann)
vTerm ((name -> f name)
 -> Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> ((Unique -> f Unique) -> name -> f name)
-> (Unique -> f Unique)
-> Term tyname name uni fun a
-> f (Term tyname name uni fun a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Unique -> f Unique) -> name -> f name
forall name unique. HasUnique name unique => Lens' name Unique
Lens' name Unique
PLC.theUnique (forall {f :: * -> *}.
 (Contravariant f, Applicative f) =>
 (Unique -> f Unique)
 -> Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> (forall {f :: * -> *}.
    (Contravariant f, Applicative f) =>
    (Unique -> f Unique)
    -> Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> forall {f :: * -> *}.
   (Contravariant f, Applicative f) =>
   (Unique -> f Unique)
   -> Term tyname name uni fun a -> f (Term tyname name uni fun a)
forall s a. Fold s a -> Fold s a -> Fold s a
<^> (tyname -> f tyname)
-> Term tyname name uni fun a -> f (Term tyname name uni fun a)
forall tyname name (uni :: * -> *) fun ann (f :: * -> *).
(Contravariant f, Applicative f) =>
(tyname -> f tyname)
-> Term tyname name uni fun ann -> f (Term tyname name uni fun ann)
tvTerm ((tyname -> f tyname)
 -> Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> ((Unique -> f Unique) -> tyname -> f tyname)
-> (Unique -> f Unique)
-> Term tyname name uni fun a
-> f (Term tyname name uni fun a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Unique -> f Unique) -> tyname -> f tyname
forall name unique. HasUnique name unique => Lens' name Unique
Lens' tyname Unique
PLC.theUnique)

-- TODO: move to plutus-core
typeUsages ::
  (PLC.HasUnique tyname PLC.TypeUnique) =>
  Type tyname uni a ->
  Usages
typeUsages :: forall tyname (uni :: * -> *) a.
HasUnique tyname TypeUnique =>
Type tyname uni a -> Usages
typeUsages = Getting Usages (Type tyname uni a) Unique
-> Type tyname uni a -> Usages
forall a s. Getting (MultiSet a) s a -> s -> MultiSet a
multiSetOf ((tyname -> Const Usages tyname)
-> Type tyname uni a -> Const Usages (Type tyname uni a)
forall tyname (uni :: * -> *) ann (f :: * -> *).
(Contravariant f, Applicative f) =>
(tyname -> f tyname)
-> Type tyname uni ann -> f (Type tyname uni ann)
tvTy ((tyname -> Const Usages tyname)
 -> Type tyname uni a -> Const Usages (Type tyname uni a))
-> ((Unique -> Const Usages Unique)
    -> tyname -> Const Usages tyname)
-> Getting Usages (Type tyname uni a) Unique
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Unique -> Const Usages Unique) -> tyname -> Const Usages tyname
forall name unique. HasUnique name unique => Lens' name Unique
Lens' tyname Unique
PLC.theUnique)