{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}

module UntypedPlutusCore.Transform.Certify.Trace where

import UntypedPlutusCore.Core.Type (Term)
import UntypedPlutusCore.Transform.Certify.Hints qualified as Certify

import Control.DeepSeq
import GHC.Generics

{-| Datatype which represents optimization passes which are also
certified.

This means that these passes are formalized as part of the certifier,
and adding a new pass constructor to this type means that it is expected
the pass will be also certified in the same PR.

WARNING: the order of the constructors MUST be the same as the order
of their counterparts in 'VerifiedCompilation.Trace'. -}
data CertifiedOptStage
  = FloatDelay
  | ForceDelay
  | ForceCaseDelay
  | Inline
  | CSE
  | ApplyToCase
  | CaseReduce
  | LetFloatOut
  deriving stock (Int -> CertifiedOptStage -> ShowS
[CertifiedOptStage] -> ShowS
CertifiedOptStage -> String
(Int -> CertifiedOptStage -> ShowS)
-> (CertifiedOptStage -> String)
-> ([CertifiedOptStage] -> ShowS)
-> Show CertifiedOptStage
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CertifiedOptStage -> ShowS
showsPrec :: Int -> CertifiedOptStage -> ShowS
$cshow :: CertifiedOptStage -> String
show :: CertifiedOptStage -> String
$cshowList :: [CertifiedOptStage] -> ShowS
showList :: [CertifiedOptStage] -> ShowS
Show, (forall x. CertifiedOptStage -> Rep CertifiedOptStage x)
-> (forall x. Rep CertifiedOptStage x -> CertifiedOptStage)
-> Generic CertifiedOptStage
forall x. Rep CertifiedOptStage x -> CertifiedOptStage
forall x. CertifiedOptStage -> Rep CertifiedOptStage x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CertifiedOptStage -> Rep CertifiedOptStage x
from :: forall x. CertifiedOptStage -> Rep CertifiedOptStage x
$cto :: forall x. Rep CertifiedOptStage x -> CertifiedOptStage
to :: forall x. Rep CertifiedOptStage x -> CertifiedOptStage
Generic)
  deriving anyclass (CertifiedOptStage -> ()
(CertifiedOptStage -> ()) -> NFData CertifiedOptStage
forall a. (a -> ()) -> NFData a
$crnf :: CertifiedOptStage -> ()
rnf :: CertifiedOptStage -> ()
NFData)

{-| Datatype which represents optimization passes which are not yet
certified.

WARNING: the order of the constructors MUST be the same as the order
of their counterparts in 'VerifiedCompilation.Trace'.

IMPORTANT: if you add a new pass, or modify an existing pass, without
also modifying the certifier in the same PR, you must add/move its
corresponding constructor to this type. Please also open an issue
at https://github.com/IntersectMBO/plutus/issues. -}
data UncertifiedOptStage
  = CaseOfCase
  | ConstantFolding
  | PolyBuiltin
  deriving stock (Int -> UncertifiedOptStage -> ShowS
[UncertifiedOptStage] -> ShowS
UncertifiedOptStage -> String
(Int -> UncertifiedOptStage -> ShowS)
-> (UncertifiedOptStage -> String)
-> ([UncertifiedOptStage] -> ShowS)
-> Show UncertifiedOptStage
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> UncertifiedOptStage -> ShowS
showsPrec :: Int -> UncertifiedOptStage -> ShowS
$cshow :: UncertifiedOptStage -> String
show :: UncertifiedOptStage -> String
$cshowList :: [UncertifiedOptStage] -> ShowS
showList :: [UncertifiedOptStage] -> ShowS
Show, (forall x. UncertifiedOptStage -> Rep UncertifiedOptStage x)
-> (forall x. Rep UncertifiedOptStage x -> UncertifiedOptStage)
-> Generic UncertifiedOptStage
forall x. Rep UncertifiedOptStage x -> UncertifiedOptStage
forall x. UncertifiedOptStage -> Rep UncertifiedOptStage x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. UncertifiedOptStage -> Rep UncertifiedOptStage x
from :: forall x. UncertifiedOptStage -> Rep UncertifiedOptStage x
$cto :: forall x. Rep UncertifiedOptStage x -> UncertifiedOptStage
to :: forall x. Rep UncertifiedOptStage x -> UncertifiedOptStage
Generic)
  deriving anyclass (UncertifiedOptStage -> ()
(UncertifiedOptStage -> ()) -> NFData UncertifiedOptStage
forall a. (a -> ()) -> NFData a
$crnf :: UncertifiedOptStage -> ()
rnf :: UncertifiedOptStage -> ()
NFData)

type OptStage = Either UncertifiedOptStage CertifiedOptStage

pattern FloatDelayStage :: OptStage
pattern $mFloatDelayStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bFloatDelayStage :: OptStage
FloatDelayStage = Right FloatDelay

pattern ForceDelayStage :: OptStage
pattern $mForceDelayStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bForceDelayStage :: OptStage
ForceDelayStage = Right ForceDelay

pattern ForceCaseDelayStage :: OptStage
pattern $mForceCaseDelayStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bForceCaseDelayStage :: OptStage
ForceCaseDelayStage = Right ForceCaseDelay

pattern CaseReduceStage :: OptStage
pattern $mCaseReduceStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bCaseReduceStage :: OptStage
CaseReduceStage = Right CaseReduce

pattern InlineStage :: OptStage
pattern $mInlineStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bInlineStage :: OptStage
InlineStage = Right Inline

pattern CseStage :: OptStage
pattern $mCseStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bCseStage :: OptStage
CseStage = Right CSE

pattern ApplyToCaseStage :: OptStage
pattern $mApplyToCaseStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bApplyToCaseStage :: OptStage
ApplyToCaseStage = Right ApplyToCase

pattern CaseOfCaseStage :: OptStage
pattern $mCaseOfCaseStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bCaseOfCaseStage :: OptStage
CaseOfCaseStage = Left CaseOfCase

pattern LetFloatOutStage :: OptStage
pattern $mLetFloatOutStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bLetFloatOutStage :: OptStage
LetFloatOutStage = Right LetFloatOut

pattern ConstantFoldingStage :: OptStage
pattern $mConstantFoldingStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bConstantFoldingStage :: OptStage
ConstantFoldingStage = Left ConstantFolding

pattern PolyBuiltinStage :: OptStage
pattern $mPolyBuiltinStage :: forall {r}. OptStage -> ((# #) -> r) -> ((# #) -> r) -> r
$bPolyBuiltinStage :: OptStage
PolyBuiltinStage = Left PolyBuiltin

{-# COMPLETE
  FloatDelayStage
  , ForceDelayStage
  , ForceCaseDelayStage
  , CaseReduceStage
  , InlineStage
  , CseStage
  , ApplyToCaseStage
  , CaseOfCaseStage
  , LetFloatOutStage
  , ConstantFoldingStage
  , PolyBuiltinStage
  #-}

data Optimization name uni fun a
  = Optimization
  { forall name (uni :: * -> *) fun a.
Optimization name uni fun a -> Term name uni fun a
beforeAST :: Term name uni fun a
  , forall name (uni :: * -> *) fun a.
Optimization name uni fun a -> OptStage
stage :: OptStage
  , forall name (uni :: * -> *) fun a.
Optimization name uni fun a -> Hints
hints :: Certify.Hints
  , forall name (uni :: * -> *) fun a.
Optimization name uni fun a -> Term name uni fun a
afterAST :: Term name uni fun a
  }

-- TODO2: we probably don't want this in memory so after MVP
-- we should consider serializing this to disk
newtype OptimizerTrace name uni fun a
  = OptimizerTrace
  { forall name (uni :: * -> *) fun a.
OptimizerTrace name uni fun a -> [Optimization name uni fun a]
optimizerTrace
      :: [Optimization name uni fun a]
  }

initOptimizerTrace :: OptimizerTrace name uni fun a
initOptimizerTrace :: forall name (uni :: * -> *) fun a. OptimizerTrace name uni fun a
initOptimizerTrace = [Optimization name uni fun a] -> OptimizerTrace name uni fun a
forall name (uni :: * -> *) fun a.
[Optimization name uni fun a] -> OptimizerTrace name uni fun a
OptimizerTrace []

allASTs :: OptimizerTrace name uni fun a -> [Term name uni fun a]
allASTs :: forall name (uni :: * -> *) fun a.
OptimizerTrace name uni fun a -> [Term name uni fun a]
allASTs = \case
  OptimizerTrace [] -> []
  OptimizerTrace xs :: [Optimization name uni fun a]
xs@(Optimization name uni fun a
x : [Optimization name uni fun a]
_) ->
    -- `OptimizerTrace` is in reverse order: the first item is the last pass run.
    Optimization name uni fun a -> Term name uni fun a
forall name (uni :: * -> *) fun a.
Optimization name uni fun a -> Term name uni fun a
afterAST Optimization name uni fun a
x Term name uni fun a
-> [Term name uni fun a] -> [Term name uni fun a]
forall a. a -> [a] -> [a]
: (Optimization name uni fun a -> Term name uni fun a)
-> [Optimization name uni fun a] -> [Term name uni fun a]
forall a b. (a -> b) -> [a] -> [b]
map Optimization name uni fun a -> Term name uni fun a
forall name (uni :: * -> *) fun a.
Optimization name uni fun a -> Term name uni fun a
beforeAST [Optimization name uni fun a]
xs