{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE GADTs             #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}
-- | Optimization passes for removing dead code, mainly dead let bindings.
module PlutusIR.Transform.DeadCode (removeDeadBindings) where

import PlutusIR
import PlutusIR.Analysis.Dependencies qualified as Deps
import PlutusIR.MkPir
import PlutusIR.Transform.Rename ()

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

import Control.Lens
import Control.Monad.Reader

import Data.Coerce
import Data.Set qualified as Set

import Algebra.Graph qualified as G
import Algebra.Graph.ToGraph qualified as T
import Data.List.NonEmpty qualified as NE
import PlutusCore.Quote (MonadQuote, freshTyName, liftQuote)
import PlutusCore.StdLib.Data.ScottUnit qualified as Unit
import Witherable (Witherable (wither))

-- | Remove all the dead let bindings in a term.
removeDeadBindings
    :: (PLC.HasUnique name PLC.TermUnique,
       PLC.ToBuiltinMeaning uni fun, PLC.MonadQuote m)
    => Term TyName name uni fun a
    -> m (Term TyName name uni fun a)
removeDeadBindings :: Term TyName name uni fun a -> m (Term TyName name uni fun a)
removeDeadBindings Term TyName name uni fun a
t = do
    Term TyName name uni fun a
tRen <- Term TyName name uni fun a -> m (Term TyName name uni fun a)
forall a (m :: * -> *). (Rename a, MonadQuote m) => a -> m a
PLC.rename Term TyName name uni fun a
t
    Quote (Term TyName name uni fun a)
-> m (Term TyName name uni fun a)
forall (m :: * -> *) a. MonadQuote m => Quote a -> m a
liftQuote (Quote (Term TyName name uni fun a)
 -> m (Term TyName name uni fun a))
-> Quote (Term TyName name uni fun a)
-> m (Term TyName name uni fun a)
forall a b. (a -> b) -> a -> b
$ ReaderT Liveness (QuoteT Identity) (Term TyName name uni fun a)
-> Liveness -> Quote (Term TyName name uni fun a)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (LensLike
  (WrappedMonad (ReaderT Liveness (QuoteT Identity)))
  (Term TyName name uni fun a)
  (Term TyName name uni fun a)
  (Term TyName name uni fun a)
  (Term TyName name uni fun a)
-> (Term TyName name uni fun a
    -> ReaderT Liveness (QuoteT Identity) (Term TyName name uni fun a))
-> Term TyName name uni fun a
-> ReaderT Liveness (QuoteT Identity) (Term TyName name uni fun a)
forall (m :: * -> *) a b.
Monad m =>
LensLike (WrappedMonad m) a b a b -> (b -> m b) -> a -> m b
transformMOf LensLike
  (WrappedMonad (ReaderT Liveness (QuoteT Identity)))
  (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 tyname name (uni :: * -> *) fun a.
Traversal'
  (Term tyname name uni fun a) (Term tyname name uni fun a)
termSubterms Term TyName name uni fun a
-> ReaderT Liveness (QuoteT Identity) (Term TyName name uni fun a)
forall (m :: * -> *) name (uni :: * -> *) fun a.
(MonadReader Liveness m, HasUnique name TermUnique,
 MonadQuote m) =>
Term TyName name uni fun a -> m (Term TyName name uni fun a)
processTerm Term TyName name uni fun a
tRen) (Term TyName name uni fun a -> Liveness
forall name tyname (uni :: * -> *) fun a.
(HasUnique name TermUnique, HasUnique tyname TypeUnique,
 ToBuiltinMeaning uni fun) =>
Term tyname name uni fun a -> Liveness
calculateLiveness Term TyName name uni fun a
tRen)

type Liveness = Set.Set Deps.Node

calculateLiveness
    :: (PLC.HasUnique name PLC.TermUnique, PLC.HasUnique tyname PLC.TypeUnique,
       PLC.ToBuiltinMeaning uni fun)
    => Term tyname name uni fun a
    -> Liveness
calculateLiveness :: Term tyname name uni fun a -> Liveness
calculateLiveness Term tyname name uni fun a
t =
    let
        depGraph :: G.Graph Deps.Node
        depGraph :: Graph Node
depGraph = (Graph Node, StrictnessMap) -> Graph Node
forall a b. (a, b) -> a
fst ((Graph Node, StrictnessMap) -> Graph Node)
-> (Graph Node, StrictnessMap) -> Graph Node
forall a b. (a -> b) -> a -> b
$ Term tyname name uni fun a -> (Graph Node, StrictnessMap)
forall g tyname name (uni :: * -> *) fun a.
(DepGraph g, HasUnique tyname TypeUnique,
 HasUnique name TermUnique, ToBuiltinMeaning uni fun) =>
Term tyname name uni fun a -> (g, StrictnessMap)
Deps.runTermDeps Term tyname name uni fun a
t
    in [Node] -> Liveness
forall a. Ord a => [a] -> Set a
Set.fromList ([Node] -> Liveness) -> [Node] -> Liveness
forall a b. (a -> b) -> a -> b
$ ToVertex (Graph Node) -> Graph Node -> [ToVertex (Graph Node)]
forall t.
(ToGraph t, Ord (ToVertex t)) =>
ToVertex t -> t -> [ToVertex t]
T.reachable ToVertex (Graph Node)
Node
Deps.Root Graph Node
depGraph

live :: (MonadReader Liveness m, PLC.HasUnique n unique) => n -> m Bool
live :: n -> m Bool
live n
n =
    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
    in (Liveness -> Bool) -> m Bool
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((Liveness -> Bool) -> m Bool) -> (Liveness -> Bool) -> m Bool
forall a b. (a -> b) -> a -> b
$ Node -> Liveness -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (Unique -> Node
Deps.Variable Unique
u)

liveBinding
    :: (MonadReader Liveness m, PLC.HasUnique name PLC.TermUnique, MonadQuote m)
    => Binding TyName name uni fun a
    -> m (Maybe (Binding TyName name uni fun a))
liveBinding :: Binding TyName name uni fun a
-> m (Maybe (Binding TyName name uni fun a))
liveBinding =
    let
        -- TODO: HasUnique instances for VarDecl and TyVarDecl?
        liveVarDecl :: VarDecl tyname n uni fun ann -> m Bool
liveVarDecl (VarDecl ann
_ n
n Type tyname uni ann
_) = n -> m Bool
forall (m :: * -> *) n unique.
(MonadReader Liveness m, HasUnique n unique) =>
n -> m Bool
live n
n
        liveTyVarDecl :: TyVarDecl n ann -> m Bool
liveTyVarDecl (TyVarDecl ann
_ n
n Kind ann
_) = n -> m Bool
forall (m :: * -> *) n unique.
(MonadReader Liveness m, HasUnique n unique) =>
n -> m Bool
live n
n
    in \case
        b :: Binding TyName name uni fun a
b@(TermBind a
_ Strictness
_ VarDecl TyName name uni fun a
d Term TyName name uni fun a
_) -> do
            Bool
l <- VarDecl TyName name uni fun a -> m Bool
forall (m :: * -> *) n unique tyname (uni :: * -> *) fun ann.
(MonadReader Liveness m, HasUnique n unique) =>
VarDecl tyname n uni fun ann -> m Bool
liveVarDecl VarDecl TyName name uni fun a
d
            Maybe (Binding TyName name uni fun a)
-> m (Maybe (Binding TyName name uni fun a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Binding TyName name uni fun a)
 -> m (Maybe (Binding TyName name uni fun a)))
-> Maybe (Binding TyName name uni fun a)
-> m (Maybe (Binding TyName name uni fun a))
forall a b. (a -> b) -> a -> b
$ if Bool
l then Binding TyName name uni fun a
-> Maybe (Binding TyName name uni fun a)
forall a. a -> Maybe a
Just Binding TyName name uni fun a
b else Maybe (Binding TyName name uni fun a)
forall a. Maybe a
Nothing
        b :: Binding TyName name uni fun a
b@(TypeBind a
_ TyVarDecl TyName a
d Type TyName uni a
_) -> do
            Bool
l <- TyVarDecl TyName a -> m Bool
forall (m :: * -> *) n unique ann.
(MonadReader Liveness m, HasUnique n unique) =>
TyVarDecl n ann -> m Bool
liveTyVarDecl TyVarDecl TyName a
d
            Maybe (Binding TyName name uni fun a)
-> m (Maybe (Binding TyName name uni fun a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Binding TyName name uni fun a)
 -> m (Maybe (Binding TyName name uni fun a)))
-> Maybe (Binding TyName name uni fun a)
-> m (Maybe (Binding TyName name uni fun a))
forall a b. (a -> b) -> a -> b
$ if Bool
l then Binding TyName name uni fun a
-> Maybe (Binding TyName name uni fun a)
forall a. a -> Maybe a
Just Binding TyName name uni fun a
b else Maybe (Binding TyName name uni fun a)
forall a. Maybe a
Nothing
        b :: Binding TyName name uni fun a
b@(DatatypeBind a
x (Datatype a
_ TyVarDecl TyName a
d [TyVarDecl TyName a]
_ name
destr [VarDecl TyName name uni fun a]
constrs)) -> do
            Bool
dtypeLive <- TyVarDecl TyName a -> m Bool
forall (m :: * -> *) n unique ann.
(MonadReader Liveness m, HasUnique n unique) =>
TyVarDecl n ann -> m Bool
liveTyVarDecl TyVarDecl TyName a
d
            Bool
destrLive <- name -> m Bool
forall (m :: * -> *) n unique.
(MonadReader Liveness m, HasUnique n unique) =>
n -> m Bool
live name
destr
            [Bool]
constrsLive <- (VarDecl TyName name uni fun a -> m Bool)
-> [VarDecl TyName name uni fun a] -> m [Bool]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse VarDecl TyName name uni fun a -> m Bool
forall (m :: * -> *) n unique tyname (uni :: * -> *) fun ann.
(MonadReader Liveness m, HasUnique n unique) =>
VarDecl tyname n uni fun ann -> m Bool
liveVarDecl [VarDecl TyName name uni fun a]
constrs
            let termLive :: Bool
termLive = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or (Bool
destrLive Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
: [Bool]
constrsLive)
            case (Bool
dtypeLive, Bool
termLive) of
                 -- At least one term-level part is live, keep the whole thing
                (Bool
_, Bool
True)      -> Maybe (Binding TyName name uni fun a)
-> m (Maybe (Binding TyName name uni fun a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Binding TyName name uni fun a)
 -> m (Maybe (Binding TyName name uni fun a)))
-> Maybe (Binding TyName name uni fun a)
-> m (Maybe (Binding TyName name uni fun a))
forall a b. (a -> b) -> a -> b
$ Binding TyName name uni fun a
-> Maybe (Binding TyName name uni fun a)
forall a. a -> Maybe a
Just Binding TyName name uni fun a
b
                -- Nothing is live, remove the whole thing
                (Bool
False, Bool
False) -> Maybe (Binding TyName name uni fun a)
-> m (Maybe (Binding TyName name uni fun a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Binding TyName name uni fun a)
forall a. Maybe a
Nothing
                -- See Note [Dependencies for datatype bindings, and pruning them]
                 -- Datatype is live but no term-level parts are, replace with a trivial type binding
                (Bool
True, Bool
False)  -> Binding TyName name uni fun a
-> Maybe (Binding TyName name uni fun a)
forall a. a -> Maybe a
Just (Binding TyName name uni fun a
 -> Maybe (Binding TyName name uni fun a))
-> (Type TyName uni a -> Binding TyName name uni fun a)
-> Type TyName uni a
-> Maybe (Binding TyName name uni fun a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a
-> TyVarDecl TyName a
-> Type TyName uni a
-> Binding TyName name uni fun a
forall tyname name (uni :: * -> *) fun a.
a
-> TyVarDecl tyname a
-> Type tyname uni a
-> Binding tyname name uni fun a
TypeBind a
x TyVarDecl TyName a
d (Type TyName uni a -> Maybe (Binding TyName name uni fun a))
-> m (Type TyName uni a)
-> m (Maybe (Binding TyName name uni fun a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Kind a -> m (Type TyName uni a)
forall (m :: * -> *) a (uni :: * -> *).
MonadQuote m =>
Kind a -> m (Type TyName uni a)
mkTypeOfKind (TyVarDecl TyName a -> Kind a
forall tyname ann. TyVarDecl tyname ann -> Kind ann
_tyVarDeclKind TyVarDecl TyName a
d)

-- | Given a kind, make a type (any type!) of that kind.
-- Generates things of the form 'unit -> unit -> ... -> unit'
mkTypeOfKind :: MonadQuote m => Kind a -> m (Type TyName uni a)
mkTypeOfKind :: Kind a -> m (Type TyName uni a)
mkTypeOfKind = \case
    -- The scott-encoded unit here is a little bulky but it continues to be the easiest
    -- way to get a type of kind Type without relying on builtins.
    Type a
a -> Type TyName uni a -> m (Type TyName uni a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type TyName uni a -> m (Type TyName uni a))
-> Type TyName uni a -> m (Type TyName uni a)
forall a b. (a -> b) -> a -> b
$ a
a a -> Type TyName uni () -> Type TyName uni a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Type TyName uni ()
forall (uni :: * -> *). Type TyName uni ()
Unit.unit
    KindArrow a
a Kind a
ki Kind a
ki' -> do
        TyName
n <- Text -> m TyName
forall (m :: * -> *). MonadQuote m => Text -> m TyName
freshTyName Text
"a"
        a -> TyName -> Kind a -> Type TyName uni a -> Type TyName uni a
forall tyname (uni :: * -> *) ann.
ann
-> tyname -> Kind ann -> Type tyname uni ann -> Type tyname uni ann
TyLam a
a TyName
n Kind a
ki (Type TyName uni a -> Type TyName uni a)
-> m (Type TyName uni a) -> m (Type TyName uni a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Kind a -> m (Type TyName uni a)
forall (m :: * -> *) a (uni :: * -> *).
MonadQuote m =>
Kind a -> m (Type TyName uni a)
mkTypeOfKind Kind a
ki'

processTerm
    :: (MonadReader Liveness m, PLC.HasUnique name PLC.TermUnique, MonadQuote m)
    => Term TyName name uni fun a
    -> m (Term TyName name uni fun a)
processTerm :: Term TyName name uni fun a -> m (Term TyName name uni fun a)
processTerm = \case
    -- throw away dead bindings
    Let a
x Recursivity
r NonEmpty (Binding TyName name uni fun a)
bs Term TyName name uni fun a
t -> a
-> Recursivity
-> [Binding TyName name uni fun a]
-> Term TyName name uni fun a
-> Term TyName name uni fun a
forall a tyname name (uni :: * -> *) fun.
a
-> Recursivity
-> [Binding tyname name uni fun a]
-> Term tyname name uni fun a
-> Term tyname name uni fun a
mkLet a
x Recursivity
r ([Binding TyName name uni fun a]
 -> Term TyName name uni fun a -> Term TyName name uni fun a)
-> m [Binding TyName name uni fun a]
-> m (Term TyName name uni fun a -> Term TyName name uni fun a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Binding TyName name uni fun a
 -> m (Maybe (Binding TyName name uni fun a)))
-> [Binding TyName name uni fun a]
-> m [Binding TyName name uni fun a]
forall (t :: * -> *) (f :: * -> *) a b.
(Witherable t, Applicative f) =>
(a -> f (Maybe b)) -> t a -> f (t b)
wither Binding TyName name uni fun a
-> m (Maybe (Binding TyName name uni fun a))
forall (m :: * -> *) name (uni :: * -> *) fun a.
(MonadReader Liveness m, HasUnique name TermUnique,
 MonadQuote m) =>
Binding TyName name uni fun a
-> m (Maybe (Binding TyName name uni fun a))
liveBinding (NonEmpty (Binding TyName name uni fun a)
-> [Binding TyName name uni fun a]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (Binding TyName name uni fun a)
bs) m (Term TyName name uni fun a -> Term TyName name uni fun a)
-> m (Term TyName name uni fun a) -> m (Term TyName name uni fun a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Term TyName name uni fun a -> m (Term TyName name uni fun a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term TyName name uni fun a
t
    Term TyName name uni fun a
x            -> Term TyName name uni fun a -> m (Term TyName name uni fun a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term TyName name uni fun a
x