{-# LANGUAGE OverloadedStrings #-}

module Transform.Simplify.Lib where

import Control.Lens ((&), (.~))
import Data.ByteString.Lazy qualified as BSL
import Data.Text.Encoding (encodeUtf8)
import PlutusCore qualified as PLC
import PlutusCore.Builtin (BuiltinSemanticsVariant)
import PlutusCore.Pretty (PrettyPlc, Render (render), prettyPlcReadableSimple)
import PlutusPrelude (Default (def))
import Test.Tasty (TestTree)
import Test.Tasty.Golden (goldenVsString)
import UntypedPlutusCore (Name, SimplifierTrace, Term, defaultSimplifyOpts, runSimplifierT,
                          soInlineCallsiteGrowth, soMaxCseIterations, soMaxSimplifierIterations,
                          soPreserveLogging, termSimplifier)

-- TODO Fix duplication with other golden tests, quite annoying
goldenVsPretty :: (PrettyPlc a) => String -> String -> a -> TestTree
goldenVsPretty :: forall a. PrettyPlc a => String -> String -> a -> TestTree
goldenVsPretty String
extn String
name a
value =
  String -> String -> IO ByteString -> TestTree
goldenVsString String
name (String
"untyped-plutus-core/test/Transform/" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
extn) (IO ByteString -> TestTree) -> IO ByteString -> TestTree
forall a b. (a -> b) -> a -> b
$
    ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString)
-> (Doc Any -> ByteString) -> Doc Any -> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BSL.fromStrict (ByteString -> ByteString)
-> (Doc Any -> ByteString) -> Doc Any -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
encodeUtf8 (Text -> ByteString) -> (Doc Any -> Text) -> Doc Any -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Doc Any -> Text
forall ann. Doc ann -> Text
forall str ann. Render str => Doc ann -> str
render (Doc Any -> IO ByteString) -> Doc Any -> IO ByteString
forall a b. (a -> b) -> a -> b
$
      a -> Doc Any
forall a ann. PrettyPlc a => a -> Doc ann
prettyPlcReadableSimple a
value

goldenVsSimplified :: String -> Term Name PLC.DefaultUni PLC.DefaultFun () -> TestTree
goldenVsSimplified :: String -> Term Name DefaultUni DefaultFun () -> TestTree
goldenVsSimplified String
name =
  String -> String -> Term Name DefaultUni DefaultFun () -> TestTree
forall a. PrettyPlc a => String -> String -> a -> TestTree
goldenVsPretty String
".uplc.golden" String
name
    (Term Name DefaultUni DefaultFun () -> TestTree)
-> (Term Name DefaultUni DefaultFun ()
    -> Term Name DefaultUni DefaultFun ())
-> Term Name DefaultUni DefaultFun ()
-> TestTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Quote (Term Name DefaultUni DefaultFun ())
-> Term Name DefaultUni DefaultFun ()
forall a. Quote a -> a
PLC.runQuote
    (Quote (Term Name DefaultUni DefaultFun ())
 -> Term Name DefaultUni DefaultFun ())
-> (Term Name DefaultUni DefaultFun ()
    -> Quote (Term Name DefaultUni DefaultFun ()))
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Term Name DefaultUni DefaultFun (),
  SimplifierTrace Name DefaultUni DefaultFun ())
 -> Term Name DefaultUni DefaultFun ())
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
-> Quote (Term Name DefaultUni DefaultFun ())
forall a b. (a -> b) -> QuoteT Identity a -> QuoteT Identity b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Term Name DefaultUni DefaultFun (),
 SimplifierTrace Name DefaultUni DefaultFun ())
-> Term Name DefaultUni DefaultFun ()
forall a b. (a, b) -> a
fst
    (QuoteT
   Identity
   (Term Name DefaultUni DefaultFun (),
    SimplifierTrace Name DefaultUni DefaultFun ())
 -> Quote (Term Name DefaultUni DefaultFun ()))
-> (Term Name DefaultUni DefaultFun ()
    -> QuoteT
         Identity
         (Term Name DefaultUni DefaultFun (),
          SimplifierTrace Name DefaultUni DefaultFun ()))
-> Term Name DefaultUni DefaultFun ()
-> Quote (Term Name DefaultUni DefaultFun ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term Name DefaultUni DefaultFun ()
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
testSimplify

testSimplify
  :: Term Name PLC.DefaultUni PLC.DefaultFun ()
  -> PLC.Quote
      ( Term Name PLC.DefaultUni PLC.DefaultFun ()
      , SimplifierTrace Name PLC.DefaultUni PLC.DefaultFun ()
      )
testSimplify :: Term Name DefaultUni DefaultFun ()
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
testSimplify =
    SimplifierT
  Name
  DefaultUni
  DefaultFun
  ()
  (QuoteT Identity)
  (Term Name DefaultUni DefaultFun ())
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
SimplifierT name uni fun ann m a
-> m (a, SimplifierTrace name uni fun ann)
runSimplifierT
    (SimplifierT
   Name
   DefaultUni
   DefaultFun
   ()
   (QuoteT Identity)
   (Term Name DefaultUni DefaultFun ())
 -> QuoteT
      Identity
      (Term Name DefaultUni DefaultFun (),
       SimplifierTrace Name DefaultUni DefaultFun ()))
-> (Term Name DefaultUni DefaultFun ()
    -> SimplifierT
         Name
         DefaultUni
         DefaultFun
         ()
         (QuoteT Identity)
         (Term Name DefaultUni DefaultFun ()))
-> Term Name DefaultUni DefaultFun ()
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SimplifyOpts Name ()
-> BuiltinSemanticsVariant DefaultFun
-> Term Name DefaultUni DefaultFun ()
-> SimplifierT
     Name
     DefaultUni
     DefaultFun
     ()
     (QuoteT Identity)
     (Term Name DefaultUni DefaultFun ())
forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
termSimplifier
      ( SimplifyOpts Name ()
forall name a. SimplifyOpts name a
defaultSimplifyOpts
          -- Just run one iteration, to see what that does
          SimplifyOpts Name ()
-> (SimplifyOpts Name () -> SimplifyOpts Name ())
-> SimplifyOpts Name ()
forall a b. a -> (a -> b) -> b
& (Int -> Identity Int)
-> SimplifyOpts Name () -> Identity (SimplifyOpts Name ())
forall name a (f :: * -> *).
Functor f =>
(Int -> f Int) -> SimplifyOpts name a -> f (SimplifyOpts name a)
soMaxSimplifierIterations ((Int -> Identity Int)
 -> SimplifyOpts Name () -> Identity (SimplifyOpts Name ()))
-> Int -> SimplifyOpts Name () -> SimplifyOpts Name ()
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Int
1
          SimplifyOpts Name ()
-> (SimplifyOpts Name () -> SimplifyOpts Name ())
-> SimplifyOpts Name ()
forall a b. a -> (a -> b) -> b
& (Int -> Identity Int)
-> SimplifyOpts Name () -> Identity (SimplifyOpts Name ())
forall name a (f :: * -> *).
Functor f =>
(Int -> f Int) -> SimplifyOpts name a -> f (SimplifyOpts name a)
soMaxCseIterations ((Int -> Identity Int)
 -> SimplifyOpts Name () -> Identity (SimplifyOpts Name ()))
-> Int -> SimplifyOpts Name () -> SimplifyOpts Name ()
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Int
0
          SimplifyOpts Name ()
-> (SimplifyOpts Name () -> SimplifyOpts Name ())
-> SimplifyOpts Name ()
forall a b. a -> (a -> b) -> b
& (Size -> Identity Size)
-> SimplifyOpts Name () -> Identity (SimplifyOpts Name ())
forall name a (f :: * -> *).
Functor f =>
(Size -> f Size) -> SimplifyOpts name a -> f (SimplifyOpts name a)
soInlineCallsiteGrowth ((Size -> Identity Size)
 -> SimplifyOpts Name () -> Identity (SimplifyOpts Name ()))
-> Size -> SimplifyOpts Name () -> SimplifyOpts Name ()
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Size
0
          SimplifyOpts Name ()
-> (SimplifyOpts Name () -> SimplifyOpts Name ())
-> SimplifyOpts Name ()
forall a b. a -> (a -> b) -> b
& (Bool -> Identity Bool)
-> SimplifyOpts Name () -> Identity (SimplifyOpts Name ())
forall name a (f :: * -> *).
Functor f =>
(Bool -> f Bool) -> SimplifyOpts name a -> f (SimplifyOpts name a)
soPreserveLogging ((Bool -> Identity Bool)
 -> SimplifyOpts Name () -> Identity (SimplifyOpts Name ()))
-> Bool -> SimplifyOpts Name () -> SimplifyOpts Name ()
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Bool
False
      )
      (BuiltinSemanticsVariant DefaultFun
forall a. Default a => a
def :: BuiltinSemanticsVariant PLC.DefaultFun)

goldenVsCse :: String -> Term Name PLC.DefaultUni PLC.DefaultFun () -> TestTree
goldenVsCse :: String -> Term Name DefaultUni DefaultFun () -> TestTree
goldenVsCse String
name =
  String -> String -> Term Name DefaultUni DefaultFun () -> TestTree
forall a. PrettyPlc a => String -> String -> a -> TestTree
goldenVsPretty String
".uplc.golden" String
name
    (Term Name DefaultUni DefaultFun () -> TestTree)
-> (Term Name DefaultUni DefaultFun ()
    -> Term Name DefaultUni DefaultFun ())
-> Term Name DefaultUni DefaultFun ()
-> TestTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Quote (Term Name DefaultUni DefaultFun ())
-> Term Name DefaultUni DefaultFun ()
forall a. Quote a -> a
PLC.runQuote
    (Quote (Term Name DefaultUni DefaultFun ())
 -> Term Name DefaultUni DefaultFun ())
-> (Term Name DefaultUni DefaultFun ()
    -> Quote (Term Name DefaultUni DefaultFun ()))
-> Term Name DefaultUni DefaultFun ()
-> Term Name DefaultUni DefaultFun ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Term Name DefaultUni DefaultFun (),
  SimplifierTrace Name DefaultUni DefaultFun ())
 -> Term Name DefaultUni DefaultFun ())
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
-> Quote (Term Name DefaultUni DefaultFun ())
forall a b. (a -> b) -> QuoteT Identity a -> QuoteT Identity b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Term Name DefaultUni DefaultFun (),
 SimplifierTrace Name DefaultUni DefaultFun ())
-> Term Name DefaultUni DefaultFun ()
forall a b. (a, b) -> a
fst
    (QuoteT
   Identity
   (Term Name DefaultUni DefaultFun (),
    SimplifierTrace Name DefaultUni DefaultFun ())
 -> Quote (Term Name DefaultUni DefaultFun ()))
-> (Term Name DefaultUni DefaultFun ()
    -> QuoteT
         Identity
         (Term Name DefaultUni DefaultFun (),
          SimplifierTrace Name DefaultUni DefaultFun ()))
-> Term Name DefaultUni DefaultFun ()
-> Quote (Term Name DefaultUni DefaultFun ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term Name DefaultUni DefaultFun ()
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
testCse

testCse
  :: Term Name PLC.DefaultUni PLC.DefaultFun ()
  -> PLC.Quote
      ( Term Name PLC.DefaultUni PLC.DefaultFun ()
      , SimplifierTrace Name PLC.DefaultUni PLC.DefaultFun ()
      )
testCse :: Term Name DefaultUni DefaultFun ()
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
testCse =
    SimplifierT
  Name
  DefaultUni
  DefaultFun
  ()
  (QuoteT Identity)
  (Term Name DefaultUni DefaultFun ())
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
SimplifierT name uni fun ann m a
-> m (a, SimplifierTrace name uni fun ann)
runSimplifierT
    (SimplifierT
   Name
   DefaultUni
   DefaultFun
   ()
   (QuoteT Identity)
   (Term Name DefaultUni DefaultFun ())
 -> QuoteT
      Identity
      (Term Name DefaultUni DefaultFun (),
       SimplifierTrace Name DefaultUni DefaultFun ()))
-> (Term Name DefaultUni DefaultFun ()
    -> SimplifierT
         Name
         DefaultUni
         DefaultFun
         ()
         (QuoteT Identity)
         (Term Name DefaultUni DefaultFun ()))
-> Term Name DefaultUni DefaultFun ()
-> QuoteT
     Identity
     (Term Name DefaultUni DefaultFun (),
      SimplifierTrace Name DefaultUni DefaultFun ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SimplifyOpts Name ()
-> BuiltinSemanticsVariant DefaultFun
-> Term Name DefaultUni DefaultFun ()
-> SimplifierT
     Name
     DefaultUni
     DefaultFun
     ()
     (QuoteT Identity)
     (Term Name DefaultUni DefaultFun ())
forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
termSimplifier
      ( SimplifyOpts Name ()
forall name a. SimplifyOpts name a
defaultSimplifyOpts
          -- Just run one iteration, to see what that does
          SimplifyOpts Name ()
-> (SimplifyOpts Name () -> SimplifyOpts Name ())
-> SimplifyOpts Name ()
forall a b. a -> (a -> b) -> b
& (Int -> Identity Int)
-> SimplifyOpts Name () -> Identity (SimplifyOpts Name ())
forall name a (f :: * -> *).
Functor f =>
(Int -> f Int) -> SimplifyOpts name a -> f (SimplifyOpts name a)
soMaxSimplifierIterations ((Int -> Identity Int)
 -> SimplifyOpts Name () -> Identity (SimplifyOpts Name ()))
-> Int -> SimplifyOpts Name () -> SimplifyOpts Name ()
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Int
0
          SimplifyOpts Name ()
-> (SimplifyOpts Name () -> SimplifyOpts Name ())
-> SimplifyOpts Name ()
forall a b. a -> (a -> b) -> b
& (Int -> Identity Int)
-> SimplifyOpts Name () -> Identity (SimplifyOpts Name ())
forall name a (f :: * -> *).
Functor f =>
(Int -> f Int) -> SimplifyOpts name a -> f (SimplifyOpts name a)
soMaxCseIterations ((Int -> Identity Int)
 -> SimplifyOpts Name () -> Identity (SimplifyOpts Name ()))
-> Int -> SimplifyOpts Name () -> SimplifyOpts Name ()
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Int
1
          SimplifyOpts Name ()
-> (SimplifyOpts Name () -> SimplifyOpts Name ())
-> SimplifyOpts Name ()
forall a b. a -> (a -> b) -> b
& (Size -> Identity Size)
-> SimplifyOpts Name () -> Identity (SimplifyOpts Name ())
forall name a (f :: * -> *).
Functor f =>
(Size -> f Size) -> SimplifyOpts name a -> f (SimplifyOpts name a)
soInlineCallsiteGrowth ((Size -> Identity Size)
 -> SimplifyOpts Name () -> Identity (SimplifyOpts Name ()))
-> Size -> SimplifyOpts Name () -> SimplifyOpts Name ()
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Size
0
          SimplifyOpts Name ()
-> (SimplifyOpts Name () -> SimplifyOpts Name ())
-> SimplifyOpts Name ()
forall a b. a -> (a -> b) -> b
& (Bool -> Identity Bool)
-> SimplifyOpts Name () -> Identity (SimplifyOpts Name ())
forall name a (f :: * -> *).
Functor f =>
(Bool -> f Bool) -> SimplifyOpts name a -> f (SimplifyOpts name a)
soPreserveLogging ((Bool -> Identity Bool)
 -> SimplifyOpts Name () -> Identity (SimplifyOpts Name ()))
-> Bool -> SimplifyOpts Name () -> SimplifyOpts Name ()
forall s t a b. ASetter s t a b -> b -> s -> t
.~ Bool
False
      )
      (BuiltinSemanticsVariant DefaultFun
forall a. Default a => a
def :: BuiltinSemanticsVariant PLC.DefaultFun)