{-# LANGUAGE ViewPatterns #-}

-- Produce an optimization report.
module PlutusCore.Executable.OptimizerReport
  ( ReportEntry (..)
  , OptimizerReport
  , buildReport
  , printReport
  ) where

import PlutusCore.Evaluation.Machine.ExBudget (ExBudget (..))
import PlutusCore.Evaluation.Machine.ExMemory (ExCPU (..), ExMemory (..))
import UntypedPlutusCore qualified as UPLC
import UntypedPlutusCore.Transform.Certify.Trace
  ( Optimization (stage)
  , OptimizerTrace (optimizerTrace)
  , allASTs
  )

import Data.Foldable
import Data.Maybe
import Data.SatInt
import System.IO (Handle, hPutStrLn)
import Text.Printf (printf)

data ReportEntry = ReportEntry
  { ReportEntry -> Maybe String
rePassName :: Maybe String
  -- ^ `Nothing` for the first entry, `Just` for all other entries.
  , ReportEntry -> Integer
reAstSize :: Integer
  , ReportEntry -> Maybe (Either String ExBudget)
reCost :: Maybe (Either String ExBudget)
  -- ^ `Nothing` means evaluation wasn't requested.  `Left` means evaluation failed.
  }

type OptimizerReport = [ReportEntry]

-- FIXME: we need a lot of `reverse` here because `OptimizerTrace` contains the passes
-- in reverse order.
buildReport
  :: OptimizerTrace UPLC.Name uni fun a
  -> [(Maybe err, ExBudget)]
  -> OptimizerReport
buildReport :: forall (uni :: * -> *) fun a err.
OptimizerTrace Name uni fun a
-> [(Maybe err, ExBudget)] -> OptimizerReport
buildReport OptimizerTrace Name uni fun a
trace ([(Maybe err, ExBudget)] -> [(Maybe err, ExBudget)]
forall a. [a] -> [a]
reverse -> [(Maybe err, ExBudget)]
costs) =
  (Maybe String
 -> Integer -> Maybe (Either String ExBudget) -> ReportEntry)
-> [Maybe String]
-> [Integer]
-> [Maybe (Either String ExBudget)]
-> OptimizerReport
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Maybe String
-> Integer -> Maybe (Either String ExBudget) -> ReportEntry
ReportEntry [Maybe String]
passNames [Integer]
sizes [Maybe (Either String ExBudget)]
mCosts
  where
    asts :: [Term Name uni fun a]
asts = [Term Name uni fun a] -> [Term Name uni fun a]
forall a. [a] -> [a]
reverse (OptimizerTrace Name uni fun a -> [Term Name uni fun a]
forall name (uni :: * -> *) fun a.
OptimizerTrace name uni fun a -> [Term name uni fun a]
allASTs OptimizerTrace Name uni fun a
trace)
    passNames :: [Maybe String]
passNames =
      Maybe String
forall a. Maybe a
Nothing Maybe String -> [Maybe String] -> [Maybe String]
forall a. a -> [a] -> [a]
: (String -> Maybe String
forall a. a -> Maybe a
Just (String -> Maybe String)
-> (Optimization Name uni fun a -> String)
-> Optimization Name uni fun a
-> Maybe String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UncertifiedOptStage -> String)
-> (CertifiedOptStage -> String)
-> Either UncertifiedOptStage CertifiedOptStage
-> String
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either UncertifiedOptStage -> String
forall a. Show a => a -> String
show CertifiedOptStage -> String
forall a. Show a => a -> String
show (Either UncertifiedOptStage CertifiedOptStage -> String)
-> (Optimization Name uni fun a
    -> Either UncertifiedOptStage CertifiedOptStage)
-> Optimization Name uni fun a
-> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Optimization Name uni fun a
-> Either UncertifiedOptStage CertifiedOptStage
forall name (uni :: * -> *) fun a.
Optimization name uni fun a
-> Either UncertifiedOptStage CertifiedOptStage
stage (Optimization Name uni fun a -> Maybe String)
-> [Optimization Name uni fun a] -> [Maybe String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Optimization Name uni fun a] -> [Optimization Name uni fun a]
forall a. [a] -> [a]
reverse (OptimizerTrace Name uni fun a -> [Optimization Name uni fun a]
forall name (uni :: * -> *) fun a.
OptimizerTrace name uni fun a -> [Optimization name uni fun a]
optimizerTrace OptimizerTrace Name uni fun a
trace))
    sizes :: [Integer]
sizes = AstSize -> Integer
UPLC.unAstSize (AstSize -> Integer)
-> (Term Name uni fun a -> AstSize)
-> Term Name uni fun a
-> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term Name uni fun a -> AstSize
forall name (uni :: * -> *) fun ann.
Term name uni fun ann -> AstSize
UPLC.termAstSize (Term Name uni fun a -> Integer)
-> [Term Name uni fun a] -> [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Term Name uni fun a]
asts
    mCosts :: [Maybe (Either String ExBudget)]
mCosts = case [(Maybe err, ExBudget)]
costs of
      [] -> Maybe (Either String ExBudget) -> [Maybe (Either String ExBudget)]
forall a. a -> [a]
repeat Maybe (Either String ExBudget)
forall a. Maybe a
Nothing -- evaluation wasn't requested
      [(Maybe err, ExBudget)]
cs -> Either String ExBudget -> Maybe (Either String ExBudget)
forall a. a -> Maybe a
Just (Either String ExBudget -> Maybe (Either String ExBudget))
-> ((Maybe err, ExBudget) -> Either String ExBudget)
-> (Maybe err, ExBudget)
-> Maybe (Either String ExBudget)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe err, ExBudget) -> Either String ExBudget
forall {a} {b}. (Maybe a, b) -> Either String b
classifyCost ((Maybe err, ExBudget) -> Maybe (Either String ExBudget))
-> [(Maybe err, ExBudget)] -> [Maybe (Either String ExBudget)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Maybe err, ExBudget)]
cs
    classifyCost :: (Maybe a, b) -> Either String b
classifyCost (Just a
_, b
_) = String -> Either String b
forall a b. a -> Either a b
Left String
"<failed>"
    classifyCost (Maybe a
Nothing, b
b) = b -> Either String b
forall a b. b -> Either a b
Right b
b

printReport :: Handle -> OptimizerReport -> IO ()
printReport :: Handle -> OptimizerReport -> IO ()
printReport Handle
h OptimizerReport
entries = do
  let withCosts :: Bool
withCosts = (ReportEntry -> Bool) -> OptimizerReport -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Maybe (Either String ExBudget) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Either String ExBudget) -> Bool)
-> (ReportEntry -> Maybe (Either String ExBudget))
-> ReportEntry
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ReportEntry -> Maybe (Either String ExBudget)
reCost) OptimizerReport
entries
      put :: String -> IO ()
put = Handle -> String -> IO ()
hPutStrLn Handle
h
  String -> IO ()
put String
"=== UPLC optimization report ==="
  String -> IO ()
put (Bool -> String
header Bool
withCosts)
  String -> IO ()
put (Int -> Char -> String
forall a. Int -> a -> [a]
replicate (String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Bool -> String
header Bool
withCosts)) Char
'-')
  (String -> IO ()) -> [String] -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ String -> IO ()
put ([String] -> IO ()) -> [String] -> IO ()
forall a b. (a -> b) -> a -> b
$
    (Maybe ReportEntry -> ReportEntry -> String)
-> [Maybe ReportEntry] -> OptimizerReport -> [String]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Bool -> Maybe ReportEntry -> ReportEntry -> String
formatEntry Bool
withCosts) (Maybe ReportEntry
forall a. Maybe a
Nothing Maybe ReportEntry -> [Maybe ReportEntry] -> [Maybe ReportEntry]
forall a. a -> [a] -> [a]
: (ReportEntry -> Maybe ReportEntry
forall a. a -> Maybe a
Just (ReportEntry -> Maybe ReportEntry)
-> OptimizerReport -> [Maybe ReportEntry]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OptimizerReport
entries)) OptimizerReport
entries
  String -> IO ()
put String
""

-- Column widths.
passColW, sizeColW, cpuColW, memColW :: Int
passColW :: Int
passColW = Int
20
sizeColW :: Int
sizeColW = Int
8
cpuColW :: Int
cpuColW = Int
15
memColW :: Int
memColW = Int
15

header :: Bool -> String
header :: Bool -> String
header Bool
False =
  String -> Int -> String -> Int -> String -> Int -> String -> String
forall r. PrintfType r => String -> r
printf
    String
"%-*s %*s %*s"
    Int
passColW
    String
"Pass"
    Int
sizeColW
    String
"Size"
    Int
sizeColW
    String
"ΔSize"
header Bool
True =
  String
-> Int
-> String
-> Int
-> String
-> Int
-> String
-> Int
-> String
-> Int
-> String
-> Int
-> String
-> Int
-> String
-> String
forall r. PrintfType r => String -> r
printf
    String
"%-*s %*s %*s %*s %*s %*s %*s"
    Int
passColW
    String
"Pass"
    Int
sizeColW
    String
"Size"
    Int
sizeColW
    String
"ΔSize"
    Int
cpuColW
    String
"CPU"
    Int
cpuColW
    String
"ΔCPU"
    Int
memColW
    String
"Mem"
    Int
memColW
    String
"ΔMem"

formatEntry :: Bool -> Maybe ReportEntry -> ReportEntry -> String
formatEntry :: Bool -> Maybe ReportEntry -> ReportEntry -> String
formatEntry Bool
withCosts Maybe ReportEntry
mPrev ReportEntry
cur =
  let name :: String
name = String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe String
"<initial>" (ReportEntry -> Maybe String
rePassName ReportEntry
cur)
      sz :: Integer
sz = ReportEntry -> Integer
reAstSize ReportEntry
cur
      dSz :: String
dSz = case Maybe ReportEntry
mPrev of
        Just ReportEntry
p -> Integer -> String
forall a. (Ord a, Num a, Show a) => a -> String
showSigned (Integer
sz Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- ReportEntry -> Integer
reAstSize ReportEntry
p)
        Maybe ReportEntry
Nothing -> String
""
      sizeBlock :: String
sizeBlock =
        String
-> Int -> String -> Int -> Integer -> Int -> String -> String
forall r. PrintfType r => String -> r
printf
          String
"%-*s %*d %*s"
          Int
passColW
          String
name
          Int
sizeColW
          Integer
sz
          Int
sizeColW
          String
dSz
          :: String
   in if Bool
withCosts
        then String
sizeBlock String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Maybe ReportEntry -> ReportEntry -> String
costBlock Maybe ReportEntry
mPrev ReportEntry
cur
        else String
sizeBlock

costBlock :: Maybe ReportEntry -> ReportEntry -> String
costBlock :: Maybe ReportEntry -> ReportEntry -> String
costBlock Maybe ReportEntry
mPrev ReportEntry
cur = case ReportEntry -> Maybe (Either String ExBudget)
reCost ReportEntry
cur of
  Maybe (Either String ExBudget)
Nothing -> String -> String -> String -> String -> String
forall {t} {t} {t} {t} {t}.
(PrintfArg t, PrintfArg t, PrintfArg t, PrintfArg t,
 PrintfType t) =>
t -> t -> t -> t -> t
spread String
"-" String
"" String
"-" String
""
  Just (Left String
tag) -> String -> String -> String -> String -> String
forall {t} {t} {t} {t} {t}.
(PrintfArg t, PrintfArg t, PrintfArg t, PrintfArg t,
 PrintfType t) =>
t -> t -> t -> t -> t
spread String
tag String
"" String
"" String
""
  Just (Right (ExBudget (ExCPU CostingInteger
cpu) (ExMemory CostingInteger
mem))) ->
    let mPrevB :: Maybe ExBudget
mPrevB =
          Maybe ReportEntry
mPrev Maybe ReportEntry
-> (ReportEntry -> Maybe ExBudget) -> Maybe ExBudget
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ReportEntry
r -> case ReportEntry -> Maybe (Either String ExBudget)
reCost ReportEntry
r of
            Just (Right ExBudget
b) -> ExBudget -> Maybe ExBudget
forall a. a -> Maybe a
Just ExBudget
b
            Maybe (Either String ExBudget)
_ -> Maybe ExBudget
forall a. Maybe a
Nothing
        (String
dCpu, String
dMem) = case Maybe ExBudget
mPrevB of
          Just (ExBudget (ExCPU CostingInteger
pc) (ExMemory CostingInteger
pm)) ->
            ( Int -> String
forall a. (Ord a, Num a, Show a) => a -> String
showSigned (CostingInteger -> Int
forall a. Num a => CostingInteger -> a
fromSatInt CostingInteger
cpu Int -> Int -> Int
forall a. Num a => a -> a -> a
- CostingInteger -> Int
forall a. Num a => CostingInteger -> a
fromSatInt CostingInteger
pc :: Int)
            , Int -> String
forall a. (Ord a, Num a, Show a) => a -> String
showSigned (CostingInteger -> Int
forall a. Num a => CostingInteger -> a
fromSatInt CostingInteger
mem Int -> Int -> Int
forall a. Num a => a -> a -> a
- CostingInteger -> Int
forall a. Num a => CostingInteger -> a
fromSatInt CostingInteger
pm :: Int)
            )
          Maybe ExBudget
Nothing -> (String
"", String
"")
     in String
-> Int
-> Int
-> Int
-> String
-> Int
-> Int
-> Int
-> String
-> String
forall r. PrintfType r => String -> r
printf
          String
"%*d %*s %*d %*s"
          Int
cpuColW
          (CostingInteger -> Int
forall a. Num a => CostingInteger -> a
fromSatInt CostingInteger
cpu :: Int)
          Int
cpuColW
          String
dCpu
          Int
memColW
          (CostingInteger -> Int
forall a. Num a => CostingInteger -> a
fromSatInt CostingInteger
mem :: Int)
          Int
memColW
          String
dMem
  where
    spread :: t -> t -> t -> t -> t
spread t
c t
dC t
m t
dM =
      String -> Int -> t -> Int -> t -> Int -> t -> Int -> t -> t
forall r. PrintfType r => String -> r
printf
        String
"%*s %*s %*s %*s"
        Int
cpuColW
        t
c
        Int
cpuColW
        t
dC
        Int
memColW
        t
m
        Int
memColW
        t
dM

showSigned :: (Ord a, Num a, Show a) => a -> String
showSigned :: forall a. (Ord a, Num a, Show a) => a -> String
showSigned a
n
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 = Char
'+' Char -> String -> String
forall a. a -> [a] -> [a]
: a -> String
forall a. Show a => a -> String
show a
n
  | Bool
otherwise = a -> String
forall a. Show a => a -> String
show a
n