{-# OPTIONS_GHC -Wall #-}

module FFI.OptimizerTrace
  ( Trace
  , mkFfiOptimizerTrace
  , toEvalResult
  ) where

import FFI.CostInfo
import FFI.Untyped qualified as FFI
import PlutusCore.Evaluation.Machine.ExBudget
import PlutusCore.Evaluation.Machine.ExMemory
import UntypedPlutusCore qualified as UPLC
import UntypedPlutusCore.Evaluation.Machine.Cek
import UntypedPlutusCore.Transform.Certify.Hints (Hints)
import UntypedPlutusCore.Transform.Optimizer
import Prelude hiding (head)

import Data.Coerce
import Data.Functor
import Data.List.NonEmptySep
import Data.SatInt
import Data.Text qualified as T

-- A certifier trace is a non-empty list of asts of type `a`, separated by the
-- optimizer pass that ran and the hints that were emitted
type Trace a = NonEmptySep (OptStage, Hints a) a

mkFfiOptimizerTrace
  :: OptimizerTrace UPLC.Name UPLC.DefaultUni UPLC.DefaultFun a
  -> Trace FFI.UTerm
mkFfiOptimizerTrace :: forall a.
OptimizerTrace Name DefaultUni DefaultFun a -> Trace UTerm
mkFfiOptimizerTrace (OptimizerTrace [Optimization Name DefaultUni DefaultFun a]
simplNonEmptySep) = [Optimization Name DefaultUni DefaultFun a] -> Trace UTerm
forall a.
[Optimization Name DefaultUni DefaultFun a] -> Trace UTerm
go ([Optimization Name DefaultUni DefaultFun a]
-> [Optimization Name DefaultUni DefaultFun a]
forall a. [a] -> [a]
reverse [Optimization Name DefaultUni DefaultFun a]
simplNonEmptySep)
  where
    -- Convert a term, which may be pre-term, post-term, or intermediate terms inside hints.
    convTerm :: UPLC.Term UPLC.Name UPLC.DefaultUni UPLC.DefaultFun a -> FFI.UTerm
    convTerm :: forall a. Term Name DefaultUni DefaultFun a -> UTerm
convTerm Term Name DefaultUni DefaultFun a
t = case Term Name DefaultUni DefaultFun a
-> Either
     FreeVariableError (Term NamedDeBruijn DefaultUni DefaultFun a)
forall (m :: * -> *) (uni :: * -> *) fun ann.
MonadError FreeVariableError m =>
Term Name uni fun ann -> m (Term NamedDeBruijn uni fun ann)
UPLC.deBruijnTerm Term Name DefaultUni DefaultFun a
t of
      Right Term NamedDeBruijn DefaultUni DefaultFun a
t' -> Term NamedDeBruijn DefaultUni DefaultFun () -> UTerm
forall a. Term NamedDeBruijn DefaultUni DefaultFun a -> UTerm
FFI.conv (Term NamedDeBruijn DefaultUni DefaultFun a
-> Term NamedDeBruijn DefaultUni DefaultFun ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void Term NamedDeBruijn DefaultUni DefaultFun a
t')
      Left (FreeVariableError
err :: UPLC.FreeVariableError) -> [Char] -> UTerm
forall a. HasCallStack => [Char] -> a
error ([Char] -> UTerm) -> [Char] -> UTerm
forall a b. (a -> b) -> a -> b
$ FreeVariableError -> [Char]
forall a. Show a => a -> [Char]
show FreeVariableError
err

    go
      :: [Optimization UPLC.Name UPLC.DefaultUni UPLC.DefaultFun a]
      -> Trace FFI.UTerm
    go :: forall a.
[Optimization Name DefaultUni DefaultFun a] -> Trace UTerm
go [] = [Char] -> Trace UTerm
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty trace"
    go [Optimization Term Name DefaultUni DefaultFun a
before OptStage
stage Hints (Term Name DefaultUni DefaultFun a)
hints Term Name DefaultUni DefaultFun a
after] =
      UTerm -> (OptStage, Hints UTerm) -> Trace UTerm -> Trace UTerm
forall sep a. a -> sep -> NonEmptySep sep a -> NonEmptySep sep a
Cons
        (Term Name DefaultUni DefaultFun a -> UTerm
forall a. Term Name DefaultUni DefaultFun a -> UTerm
convTerm Term Name DefaultUni DefaultFun a
before)
        (OptStage
stage, Term Name DefaultUni DefaultFun a -> UTerm
forall a. Term Name DefaultUni DefaultFun a -> UTerm
convTerm (Term Name DefaultUni DefaultFun a -> UTerm)
-> Hints (Term Name DefaultUni DefaultFun a) -> Hints UTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Hints (Term Name DefaultUni DefaultFun a)
hints)
        (UTerm -> Trace UTerm
forall sep a. a -> NonEmptySep sep a
Singleton (Term Name DefaultUni DefaultFun a -> UTerm
forall a. Term Name DefaultUni DefaultFun a -> UTerm
convTerm Term Name DefaultUni DefaultFun a
after))
    -- ignore _after, it should be equal to subsequent before
    go (Optimization Term Name DefaultUni DefaultFun a
before OptStage
stage Hints (Term Name DefaultUni DefaultFun a)
hints Term Name DefaultUni DefaultFun a
_after : [Optimization Name DefaultUni DefaultFun a]
xs) =
      UTerm -> (OptStage, Hints UTerm) -> Trace UTerm -> Trace UTerm
forall sep a. a -> sep -> NonEmptySep sep a -> NonEmptySep sep a
Cons (Term Name DefaultUni DefaultFun a -> UTerm
forall a. Term Name DefaultUni DefaultFun a -> UTerm
convTerm Term Name DefaultUni DefaultFun a
before) (OptStage
stage, Term Name DefaultUni DefaultFun a -> UTerm
forall a. Term Name DefaultUni DefaultFun a -> UTerm
convTerm (Term Name DefaultUni DefaultFun a -> UTerm)
-> Hints (Term Name DefaultUni DefaultFun a) -> Hints UTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Hints (Term Name DefaultUni DefaultFun a)
hints) ([Optimization Name DefaultUni DefaultFun a] -> Trace UTerm
forall a.
[Optimization Name DefaultUni DefaultFun a] -> Trace UTerm
go [Optimization Name DefaultUni DefaultFun a]
xs)

toEvalResult
  :: Maybe (CekEvaluationException UPLC.NamedDeBruijn UPLC.DefaultUni UPLC.DefaultFun)
  -> ExBudget
  -> EvalResult
toEvalResult :: Maybe (CekEvaluationException NamedDeBruijn DefaultUni DefaultFun)
-> ExBudget -> EvalResult
toEvalResult Maybe (CekEvaluationException NamedDeBruijn DefaultUni DefaultFun)
res ExBudget
budget = case Maybe (CekEvaluationException NamedDeBruijn DefaultUni DefaultFun)
res of
  Just CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err -> Text -> Integer -> Integer -> EvalResult
EvalFailure ([Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
-> [Char]
forall a. Show a => a -> [Char]
show CekEvaluationException NamedDeBruijn DefaultUni DefaultFun
err) Integer
cpu Integer
mem
  Maybe (CekEvaluationException NamedDeBruijn DefaultUni DefaultFun)
Nothing -> Integer -> Integer -> EvalResult
EvalSuccess Integer
cpu Integer
mem
  where
    cpu :: Integer
cpu = SatInt -> Integer
forall a. Num a => SatInt -> a
fromSatInt (SatInt -> Integer) -> SatInt -> Integer
forall a b. (a -> b) -> a -> b
$ ExCPU -> SatInt
forall a b. Coercible a b => a -> b
coerce (ExBudget -> ExCPU
exBudgetCPU ExBudget
budget)
    mem :: Integer
mem = SatInt -> Integer
forall a. Num a => SatInt -> a
fromSatInt (SatInt -> Integer) -> SatInt -> Integer
forall a b. (a -> b) -> a -> b
$ ExMemory -> SatInt
forall a b. Coercible a b => a -> b
coerce (ExBudget -> ExMemory
exBudgetMemory ExBudget
budget)