{-# LANGUAGE FlexibleContexts #-}
-- | Functions for computing variable usage inside terms and types.
module PlutusIR.Analysis.Usages (runTermUsages, runTypeUsages, Usages, getUsageCount, allUsed) where

import PlutusIR

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

import Control.Lens
import Control.Monad.State

import Data.Coerce
import Data.Foldable
import Data.Map qualified as Map
import Data.Set qualified as Set

-- | Variable uses, as a map from the 'PLC.Unique' to its usage count. Unused variables may be missing
-- or have usage count 0.
type Usages = Map.Map PLC.Unique Int

addUsage :: (PLC.HasUnique n unique) => n -> Usages -> Usages
addUsage :: n -> Usages -> Usages
addUsage n
n Usages
usages =
    let
        u :: Unique
u = unique -> Unique
coerce (unique -> Unique) -> unique -> Unique
forall a b. (a -> b) -> a -> b
$ n
n n -> Getting unique n unique -> unique
forall s a. s -> Getting a s a -> a
^. Getting unique n unique
forall a unique. HasUnique a unique => Lens' a unique
PLC.unique
        old :: Int
old = Int -> Unique -> Usages -> Int
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Int
0 Unique
u Usages
usages
    in Unique -> Int -> Usages -> Usages
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Unique
u (Int
oldInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Usages
usages

-- | Get the usage count of @n@.
getUsageCount :: (PLC.HasUnique n unique) => n -> Usages -> Int
getUsageCount :: n -> Usages -> Int
getUsageCount n
n Usages
usages = Int -> Unique -> Usages -> Int
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Int
0 (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
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
coerced) Usages
usages

-- | Get a set of @n@s which are used at least once.
allUsed :: Usages -> Set.Set PLC.Unique
allUsed :: Usages -> Set Unique
allUsed Usages
usages = Usages -> Set Unique
forall k a. Map k a -> Set k
Map.keysSet (Usages -> Set Unique) -> Usages -> Set Unique
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> Usages -> Usages
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) Usages
usages

-- | Compute the 'Usages' for a 'Term'.
runTermUsages
    :: (PLC.HasUnique name PLC.TermUnique, PLC.HasUnique tyname PLC.TypeUnique)
    => Term tyname name uni fun a
    -> Usages
runTermUsages :: Term tyname name uni fun a -> Usages
runTermUsages Term tyname name uni fun a
term = State Usages () -> Usages -> Usages
forall s a. State s a -> s -> s
execState (Term tyname name uni fun a -> State Usages ()
forall (m :: * -> *) name tyname (uni :: * -> *) fun a.
(MonadState Usages m, HasUnique name TermUnique,
 HasUnique tyname TypeUnique) =>
Term tyname name uni fun a -> m ()
termUsages Term tyname name uni fun a
term) Usages
forall a. Monoid a => a
mempty

-- | Compute the 'Usages' for a 'Type'.
runTypeUsages
    ::(PLC.HasUnique tyname PLC.TypeUnique)
    => Type tyname uni a
    -> Usages
runTypeUsages :: Type tyname uni a -> Usages
runTypeUsages Type tyname uni a
ty = State Usages () -> Usages -> Usages
forall s a. State s a -> s -> s
execState (Type tyname uni a -> State Usages ()
forall (m :: * -> *) tyname (uni :: * -> *) a.
(MonadState Usages m, HasUnique tyname TypeUnique) =>
Type tyname uni a -> m ()
typeUsages Type tyname uni a
ty) Usages
forall a. Monoid a => a
mempty

termUsages
    :: (MonadState Usages m, PLC.HasUnique name PLC.TermUnique, PLC.HasUnique tyname PLC.TypeUnique)
    => Term tyname name uni fun a
    -> m ()
termUsages :: Term tyname name uni fun a -> m ()
termUsages (Var a
_ name
n) = (Usages -> Usages) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (name -> Usages -> Usages
forall n unique. HasUnique n unique => n -> Usages -> Usages
addUsage name
n)
termUsages Term tyname name uni fun a
term      = (Term tyname name uni fun a -> m ())
-> [Term tyname name uni fun a] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Term tyname name uni fun a -> m ()
forall (m :: * -> *) name tyname (uni :: * -> *) fun a.
(MonadState Usages m, HasUnique name TermUnique,
 HasUnique tyname TypeUnique) =>
Term tyname name uni fun a -> m ()
termUsages (Term tyname name uni fun a
term Term tyname name uni fun a
-> Getting
     (Endo [Term tyname name uni fun a])
     (Term tyname name uni fun a)
     (Term tyname name uni fun a)
-> [Term tyname name uni fun a]
forall s a. s -> Getting (Endo [a]) s a -> [a]
^.. Getting
  (Endo [Term tyname name uni fun a])
  (Term tyname name uni fun a)
  (Term tyname name uni fun a)
forall tyname name (uni :: * -> *) fun a.
Traversal'
  (Term tyname name uni fun a) (Term tyname name uni fun a)
termSubterms) m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Type tyname uni a -> m ()) -> [Type tyname uni a] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Type tyname uni a -> m ()
forall (m :: * -> *) tyname (uni :: * -> *) a.
(MonadState Usages m, HasUnique tyname TypeUnique) =>
Type tyname uni a -> m ()
typeUsages (Term tyname name uni fun a
term Term tyname name uni fun a
-> Getting
     (Endo [Type tyname uni a])
     (Term tyname name uni fun a)
     (Type tyname uni a)
-> [Type tyname uni a]
forall s a. s -> Getting (Endo [a]) s a -> [a]
^.. Getting
  (Endo [Type tyname uni a])
  (Term tyname name uni fun a)
  (Type tyname uni a)
forall tyname name (uni :: * -> *) fun a.
Traversal' (Term tyname name uni fun a) (Type tyname uni a)
termSubtypes)

-- TODO: move to plutus-core
typeUsages
    :: (MonadState Usages m, PLC.HasUnique tyname PLC.TypeUnique)
    => Type tyname uni a
    -> m ()
typeUsages :: Type tyname uni a -> m ()
typeUsages (TyVar a
_ tyname
n) = (Usages -> Usages) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (tyname -> Usages -> Usages
forall n unique. HasUnique n unique => n -> Usages -> Usages
addUsage tyname
n)
typeUsages Type tyname uni a
ty          = (Type tyname uni a -> m ()) -> [Type tyname uni a] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Type tyname uni a -> m ()
forall (m :: * -> *) tyname (uni :: * -> *) a.
(MonadState Usages m, HasUnique tyname TypeUnique) =>
Type tyname uni a -> m ()
typeUsages (Type tyname uni a
ty Type tyname uni a
-> Getting
     (Endo [Type tyname uni a]) (Type tyname uni a) (Type tyname uni a)
-> [Type tyname uni a]
forall s a. s -> Getting (Endo [a]) s a -> [a]
^.. Getting
  (Endo [Type tyname uni a]) (Type tyname uni a) (Type tyname uni a)
forall tyname (uni :: * -> *) ann.
Traversal' (Type tyname uni ann) (Type tyname uni ann)
typeSubtypes)