{-# LANGUAGE PatternSynonyms #-}

module UntypedPlutusCore.Transform.Optimizer
  ( OptimizerT (..)
  , Trace.OptStage
  , pattern Trace.FloatDelayStage
  , pattern Trace.ForceDelayStage
  , pattern Trace.ForceCaseDelayStage
  , pattern Trace.CaseReduceStage
  , pattern Trace.InlineStage
  , pattern Trace.CseStage
  , pattern Trace.ApplyToCaseStage
  , pattern Trace.CaseOfCaseStage
  , pattern Trace.LetFloatOutStage
  , pattern Trace.ConstantFoldingStage
  , pattern Trace.PolyBuiltinStage
  , Trace.OptimizerTrace (..)
  , Trace.Optimization (..)
  , runOptimizerT
  , evalOptimizerT
  , execOptimizerT
  , Optimizer
  , runOptimizer
  , evalOptimizer
  , execOptimizer
  , Trace.initOptimizerTrace
  , recordOptimization
  , recordOptimizationWithHints
  ) where

import Control.Monad.State (MonadTrans, StateT)
import Control.Monad.State qualified as State

import Control.Monad.Identity (Identity, runIdentity)
import PlutusCore.Quote (MonadQuote)
import UntypedPlutusCore.Core.Type (Term)
import UntypedPlutusCore.Transform.Certify.Hints qualified as Hints
import UntypedPlutusCore.Transform.Certify.Trace qualified as Trace

newtype OptimizerT name uni fun ann m a
  = OptimizerT
  { forall name (uni :: * -> *) fun ann (m :: * -> *) a.
OptimizerT name uni fun ann m a
-> StateT (OptimizerTrace name uni fun ann) m a
getOptimizerT :: StateT (Trace.OptimizerTrace name uni fun ann) m a
  }
  deriving newtype ((forall a b.
 (a -> b)
 -> OptimizerT name uni fun ann m a
 -> OptimizerT name uni fun ann m b)
-> (forall a b.
    a
    -> OptimizerT name uni fun ann m b
    -> OptimizerT name uni fun ann m a)
-> Functor (OptimizerT name uni fun ann m)
forall a b.
a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m a
forall a b.
(a -> b)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Functor m =>
a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m a
forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Functor m =>
(a -> b)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Functor m =>
(a -> b)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
fmap :: forall a b.
(a -> b)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
$c<$ :: forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Functor m =>
a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m a
<$ :: forall a b.
a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m a
Functor, Functor (OptimizerT name uni fun ann m)
Functor (OptimizerT name uni fun ann m) =>
(forall a. a -> OptimizerT name uni fun ann m a)
-> (forall a b.
    OptimizerT name uni fun ann m (a -> b)
    -> OptimizerT name uni fun ann m a
    -> OptimizerT name uni fun ann m b)
-> (forall a b c.
    (a -> b -> c)
    -> OptimizerT name uni fun ann m a
    -> OptimizerT name uni fun ann m b
    -> OptimizerT name uni fun ann m c)
-> (forall a b.
    OptimizerT name uni fun ann m a
    -> OptimizerT name uni fun ann m b
    -> OptimizerT name uni fun ann m b)
-> (forall a b.
    OptimizerT name uni fun ann m a
    -> OptimizerT name uni fun ann m b
    -> OptimizerT name uni fun ann m a)
-> Applicative (OptimizerT name uni fun ann m)
forall a. a -> OptimizerT name uni fun ann m a
forall a b.
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m a
forall a b.
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m b
forall a b.
OptimizerT name uni fun ann m (a -> b)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
forall a b c.
(a -> b -> c)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m c
forall name (uni :: * -> *) fun ann (m :: * -> *).
Monad m =>
Functor (OptimizerT name uni fun ann m)
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
Monad m =>
a -> OptimizerT name uni fun ann m a
forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m a
forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m b
forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m (a -> b)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
forall name (uni :: * -> *) fun ann (m :: * -> *) a b c.
Monad m =>
(a -> b -> c)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall name (uni :: * -> *) fun ann (m :: * -> *) a.
Monad m =>
a -> OptimizerT name uni fun ann m a
pure :: forall a. a -> OptimizerT name uni fun ann m a
$c<*> :: forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m (a -> b)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
<*> :: forall a b.
OptimizerT name uni fun ann m (a -> b)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
$cliftA2 :: forall name (uni :: * -> *) fun ann (m :: * -> *) a b c.
Monad m =>
(a -> b -> c)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m c
liftA2 :: forall a b c.
(a -> b -> c)
-> OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m c
$c*> :: forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m b
*> :: forall a b.
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m b
$c<* :: forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m a
<* :: forall a b.
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m a
Applicative, Applicative (OptimizerT name uni fun ann m)
Applicative (OptimizerT name uni fun ann m) =>
(forall a b.
 OptimizerT name uni fun ann m a
 -> (a -> OptimizerT name uni fun ann m b)
 -> OptimizerT name uni fun ann m b)
-> (forall a b.
    OptimizerT name uni fun ann m a
    -> OptimizerT name uni fun ann m b
    -> OptimizerT name uni fun ann m b)
-> (forall a. a -> OptimizerT name uni fun ann m a)
-> Monad (OptimizerT name uni fun ann m)
forall a. a -> OptimizerT name uni fun ann m a
forall a b.
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m b
forall a b.
OptimizerT name uni fun ann m a
-> (a -> OptimizerT name uni fun ann m b)
-> OptimizerT name uni fun ann m b
forall name (uni :: * -> *) fun ann (m :: * -> *).
Monad m =>
Applicative (OptimizerT name uni fun ann m)
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
Monad m =>
a -> OptimizerT name uni fun ann m a
forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m b
forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m a
-> (a -> OptimizerT name uni fun ann m b)
-> OptimizerT name uni fun ann m b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m a
-> (a -> OptimizerT name uni fun ann m b)
-> OptimizerT name uni fun ann m b
>>= :: forall a b.
OptimizerT name uni fun ann m a
-> (a -> OptimizerT name uni fun ann m b)
-> OptimizerT name uni fun ann m b
$c>> :: forall name (uni :: * -> *) fun ann (m :: * -> *) a b.
Monad m =>
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m b
>> :: forall a b.
OptimizerT name uni fun ann m a
-> OptimizerT name uni fun ann m b
-> OptimizerT name uni fun ann m b
$creturn :: forall name (uni :: * -> *) fun ann (m :: * -> *) a.
Monad m =>
a -> OptimizerT name uni fun ann m a
return :: forall a. a -> OptimizerT name uni fun ann m a
Monad, (forall (m :: * -> *).
 Monad m =>
 Monad (OptimizerT name uni fun ann m)) =>
(forall (m :: * -> *) a.
 Monad m =>
 m a -> OptimizerT name uni fun ann m a)
-> MonadTrans (OptimizerT name uni fun ann)
forall name (uni :: * -> *) fun ann (m :: * -> *).
Monad m =>
Monad (OptimizerT name uni fun ann m)
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
Monad m =>
m a -> OptimizerT name uni fun ann m a
forall (m :: * -> *).
Monad m =>
Monad (OptimizerT name uni fun ann m)
forall (m :: * -> *) a.
Monad m =>
m a -> OptimizerT name uni fun ann m a
forall (t :: (* -> *) -> * -> *).
(forall (m :: * -> *). Monad m => Monad (t m)) =>
(forall (m :: * -> *) a. Monad m => m a -> t m a) -> MonadTrans t
$clift :: forall name (uni :: * -> *) fun ann (m :: * -> *) a.
Monad m =>
m a -> OptimizerT name uni fun ann m a
lift :: forall (m :: * -> *) a.
Monad m =>
m a -> OptimizerT name uni fun ann m a
MonadTrans)

instance MonadQuote m => MonadQuote (OptimizerT name uni fun ann m)

runOptimizerT
  :: OptimizerT name uni fun ann m a
  -> m (a, Trace.OptimizerTrace name uni fun ann)
runOptimizerT :: forall name (uni :: * -> *) fun ann (m :: * -> *) a.
OptimizerT name uni fun ann m a
-> m (a, OptimizerTrace name uni fun ann)
runOptimizerT = (StateT (OptimizerTrace name uni fun ann) m a
 -> OptimizerTrace name uni fun ann
 -> m (a, OptimizerTrace name uni fun ann))
-> OptimizerTrace name uni fun ann
-> StateT (OptimizerTrace name uni fun ann) m a
-> m (a, OptimizerTrace name uni fun ann)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (OptimizerTrace name uni fun ann) m a
-> OptimizerTrace name uni fun ann
-> m (a, OptimizerTrace name uni fun ann)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
State.runStateT OptimizerTrace name uni fun ann
forall name (uni :: * -> *) fun a. OptimizerTrace name uni fun a
Trace.initOptimizerTrace (StateT (OptimizerTrace name uni fun ann) m a
 -> m (a, OptimizerTrace name uni fun ann))
-> (OptimizerT name uni fun ann m a
    -> StateT (OptimizerTrace name uni fun ann) m a)
-> OptimizerT name uni fun ann m a
-> m (a, OptimizerTrace name uni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OptimizerT name uni fun ann m a
-> StateT (OptimizerTrace name uni fun ann) m a
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
OptimizerT name uni fun ann m a
-> StateT (OptimizerTrace name uni fun ann) m a
getOptimizerT

evalOptimizerT
  :: Monad m => OptimizerT name uni fun ann m a -> m a
evalOptimizerT :: forall (m :: * -> *) name (uni :: * -> *) fun ann a.
Monad m =>
OptimizerT name uni fun ann m a -> m a
evalOptimizerT = (StateT (OptimizerTrace name uni fun ann) m a
 -> OptimizerTrace name uni fun ann -> m a)
-> OptimizerTrace name uni fun ann
-> StateT (OptimizerTrace name uni fun ann) m a
-> m a
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (OptimizerTrace name uni fun ann) m a
-> OptimizerTrace name uni fun ann -> m a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
State.evalStateT OptimizerTrace name uni fun ann
forall name (uni :: * -> *) fun a. OptimizerTrace name uni fun a
Trace.initOptimizerTrace (StateT (OptimizerTrace name uni fun ann) m a -> m a)
-> (OptimizerT name uni fun ann m a
    -> StateT (OptimizerTrace name uni fun ann) m a)
-> OptimizerT name uni fun ann m a
-> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OptimizerT name uni fun ann m a
-> StateT (OptimizerTrace name uni fun ann) m a
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
OptimizerT name uni fun ann m a
-> StateT (OptimizerTrace name uni fun ann) m a
getOptimizerT

execOptimizerT
  :: Monad m => OptimizerT name uni fun ann m a -> m (Trace.OptimizerTrace name uni fun ann)
execOptimizerT :: forall (m :: * -> *) name (uni :: * -> *) fun ann a.
Monad m =>
OptimizerT name uni fun ann m a
-> m (OptimizerTrace name uni fun ann)
execOptimizerT = (StateT (OptimizerTrace name uni fun ann) m a
 -> OptimizerTrace name uni fun ann
 -> m (OptimizerTrace name uni fun ann))
-> OptimizerTrace name uni fun ann
-> StateT (OptimizerTrace name uni fun ann) m a
-> m (OptimizerTrace name uni fun ann)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (OptimizerTrace name uni fun ann) m a
-> OptimizerTrace name uni fun ann
-> m (OptimizerTrace name uni fun ann)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
State.execStateT OptimizerTrace name uni fun ann
forall name (uni :: * -> *) fun a. OptimizerTrace name uni fun a
Trace.initOptimizerTrace (StateT (OptimizerTrace name uni fun ann) m a
 -> m (OptimizerTrace name uni fun ann))
-> (OptimizerT name uni fun ann m a
    -> StateT (OptimizerTrace name uni fun ann) m a)
-> OptimizerT name uni fun ann m a
-> m (OptimizerTrace name uni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OptimizerT name uni fun ann m a
-> StateT (OptimizerTrace name uni fun ann) m a
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
OptimizerT name uni fun ann m a
-> StateT (OptimizerTrace name uni fun ann) m a
getOptimizerT

type Optimizer name uni fun ann = OptimizerT name uni fun ann Identity

runOptimizer :: Optimizer name uni fun ann a -> (a, Trace.OptimizerTrace name uni fun ann)
runOptimizer :: forall name (uni :: * -> *) fun ann a.
Optimizer name uni fun ann a
-> (a, OptimizerTrace name uni fun ann)
runOptimizer = Identity (a, OptimizerTrace name uni fun ann)
-> (a, OptimizerTrace name uni fun ann)
forall a. Identity a -> a
runIdentity (Identity (a, OptimizerTrace name uni fun ann)
 -> (a, OptimizerTrace name uni fun ann))
-> (Optimizer name uni fun ann a
    -> Identity (a, OptimizerTrace name uni fun ann))
-> Optimizer name uni fun ann a
-> (a, OptimizerTrace name uni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Optimizer name uni fun ann a
-> Identity (a, OptimizerTrace name uni fun ann)
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
OptimizerT name uni fun ann m a
-> m (a, OptimizerTrace name uni fun ann)
runOptimizerT

evalOptimizer :: Optimizer name uni fun ann a -> a
evalOptimizer :: forall name (uni :: * -> *) fun ann a.
Optimizer name uni fun ann a -> a
evalOptimizer = Identity a -> a
forall a. Identity a -> a
runIdentity (Identity a -> a)
-> (Optimizer name uni fun ann a -> Identity a)
-> Optimizer name uni fun ann a
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Optimizer name uni fun ann a -> Identity a
forall (m :: * -> *) name (uni :: * -> *) fun ann a.
Monad m =>
OptimizerT name uni fun ann m a -> m a
evalOptimizerT

execOptimizer :: Optimizer name uni fun ann a -> Trace.OptimizerTrace name uni fun ann
execOptimizer :: forall name (uni :: * -> *) fun ann a.
Optimizer name uni fun ann a -> OptimizerTrace name uni fun ann
execOptimizer = Identity (OptimizerTrace name uni fun ann)
-> OptimizerTrace name uni fun ann
forall a. Identity a -> a
runIdentity (Identity (OptimizerTrace name uni fun ann)
 -> OptimizerTrace name uni fun ann)
-> (Optimizer name uni fun ann a
    -> Identity (OptimizerTrace name uni fun ann))
-> Optimizer name uni fun ann a
-> OptimizerTrace name uni fun ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Optimizer name uni fun ann a
-> Identity (OptimizerTrace name uni fun ann)
forall (m :: * -> *) name (uni :: * -> *) fun ann a.
Monad m =>
OptimizerT name uni fun ann m a
-> m (OptimizerTrace name uni fun ann)
execOptimizerT

recordOptimization
  :: Monad m
  => Term name uni fun a
  -> Trace.OptStage
  -> Term name uni fun a
  -> OptimizerT name uni fun a m ()
recordOptimization :: forall (m :: * -> *) name (uni :: * -> *) fun a.
Monad m =>
Term name uni fun a
-> OptStage
-> Term name uni fun a
-> OptimizerT name uni fun a m ()
recordOptimization = Hints
-> Term name uni fun a
-> OptStage
-> Term name uni fun a
-> OptimizerT name uni fun a m ()
forall (m :: * -> *) name (uni :: * -> *) fun a.
Monad m =>
Hints
-> Term name uni fun a
-> OptStage
-> Term name uni fun a
-> OptimizerT name uni fun a m ()
recordOptimizationWithHints Hints
Hints.NoHints

recordOptimizationWithHints
  :: Monad m
  => Hints.Hints
  -> Term name uni fun a
  -> Trace.OptStage
  -> Term name uni fun a
  -> OptimizerT name uni fun a m ()
recordOptimizationWithHints :: forall (m :: * -> *) name (uni :: * -> *) fun a.
Monad m =>
Hints
-> Term name uni fun a
-> OptStage
-> Term name uni fun a
-> OptimizerT name uni fun a m ()
recordOptimizationWithHints Hints
hints Term name uni fun a
before OptStage
stage Term name uni fun a
after =
  let optimization :: Optimization name uni fun a
optimization = Term name uni fun a
-> OptStage
-> Hints
-> Term name uni fun a
-> Optimization name uni fun a
forall name (uni :: * -> *) fun a.
Term name uni fun a
-> OptStage
-> Hints
-> Term name uni fun a
-> Optimization name uni fun a
Trace.Optimization Term name uni fun a
before OptStage
stage Hints
hints Term name uni fun a
after
   in StateT (OptimizerTrace name uni fun a) m ()
-> OptimizerT name uni fun a m ()
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
StateT (OptimizerTrace name uni fun ann) m a
-> OptimizerT name uni fun ann m a
OptimizerT (StateT (OptimizerTrace name uni fun a) m ()
 -> OptimizerT name uni fun a m ())
-> ((OptimizerTrace name uni fun a
     -> OptimizerTrace name uni fun a)
    -> StateT (OptimizerTrace name uni fun a) m ())
-> (OptimizerTrace name uni fun a -> OptimizerTrace name uni fun a)
-> OptimizerT name uni fun a m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (OptimizerTrace name uni fun a -> OptimizerTrace name uni fun a)
-> StateT (OptimizerTrace name uni fun a) m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
State.modify' ((OptimizerTrace name uni fun a -> OptimizerTrace name uni fun a)
 -> OptimizerT name uni fun a m ())
-> (OptimizerTrace name uni fun a -> OptimizerTrace name uni fun a)
-> OptimizerT name uni fun a m ()
forall a b. (a -> b) -> a -> b
$ \OptimizerTrace name uni fun a
st ->
        OptimizerTrace name uni fun a
st {Trace.optimizerTrace = optimization : Trace.optimizerTrace st}