Add simplifier tags in UPLC simplifier (#6540)

Co-authored-by: Ramsay Taylor <[email protected]>

plutus-core/executables/plutus/AnyProgram/Compile.hs

@@ -31,10 +31,8 @@ import Control.Lens hiding ((%~))
 import Control.Monad.Error.Lens
 import Control.Monad.Except (MonadError)
 import Control.Monad.Reader
-import Control.Monad.State (StateT (runStateT))
 import Data.Singletons.Decide
 import Data.Text
-import PlutusCore.Compiler.Types (initUPLCSimplifierTrace)
 import PlutusPrelude hiding ((%~))
 
 -- Note that we use for erroring the original term's annotation
@@ -113,7 +111,7 @@ compileProgram = curry $ \case
             -- first self-"compile" to plc (just for reusing code)
             compileProgram sng1 (SPlc n2 a2)
         -- PLC.compileProgram subsumes uplcOptimise
-        >=> (PLC.runQuoteT . PLC.evalCompile PLC.defaultCompilationOpts .
+        >=> (PLC.runQuoteT . flip runReaderT PLC.defaultCompilationOpts .
                 plcToUplcViaName n2 PLC.compileProgram)
         >=> pure . UPLC.UnrestrictedProgram
 
@@ -230,8 +228,7 @@ uplcOptimise =
                                 case safeOrUnsafe of
                                   SafeOptimise   -> set UPLC.soConservativeOpts True
                                   UnsafeOptimise -> id
-                in (fmap . fmap) fst
-                   . fmap (PLC.runQuoteT . flip runStateT initUPLCSimplifierTrace)
+                in fmap PLC.runQuoteT
                    . _Wrapped
                    . uplcViaName (UPLC.simplifyProgram sOpts def)
 

plutus-core/plutus-core.cabal

@@ -212,6 +212,7 @@ library
     UntypedPlutusCore.Rename
     UntypedPlutusCore.Size
     UntypedPlutusCore.Transform.CaseOfCase
+    UntypedPlutusCore.Transform.Simplifier
 
   other-modules:
     Data.Aeson.Flatten

plutus-core/plutus-core/src/PlutusCore/Compiler.hs

@@ -2,8 +2,6 @@ module PlutusCore.Compiler
   ( module Opts
   , compileTerm
   , compileProgram
-  , runCompile
-  , evalCompile
   ) where
 
 import PlutusCore.Compiler.Erase
@@ -16,14 +14,12 @@ import UntypedPlutusCore.Core.Type qualified as UPLC
 import UntypedPlutusCore.Simplify qualified as UPLC
 
 import Control.Lens (view)
-import Control.Monad.Reader (MonadReader, ReaderT (runReaderT))
-import Control.Monad.State (MonadState (..), StateT (runStateT))
+import Control.Monad.Reader (MonadReader)
 
 -- | Compile a PLC term to UPLC, and optimize it.
 compileTerm
   :: (Compiling m uni fun name a
   , MonadReader (CompilationOpts name fun a) m
-  , MonadState (UPLCSimplifierTrace name uni fun a) m
   )
   => Term tyname name uni fun a
   -> m (UPLC.Term name uni fun a)
@@ -38,31 +34,7 @@ compileTerm t = do
 compileProgram
   :: (Compiling m uni fun name a
   , MonadReader (CompilationOpts name fun a) m
-  , MonadState (UPLCSimplifierTrace name uni fun a) m
   )
   => Program tyname name uni fun a
   -> m (UPLC.Program name uni fun a)
 compileProgram (Program a v t) = UPLC.Program a v <$> compileTerm t
-
-type Compile m name uni fun a =
-  ReaderT
-    (CompilationOpts name fun a)
-    (StateT
-       (UPLCSimplifierTrace name uni fun a)
-       m
-    )
-
-runCompile
-  :: CompilationOpts name fun a
-  -> Compile m name uni fun a b
-  -> m (b, UPLCSimplifierTrace name uni fun a)
-runCompile opts =
-  flip runStateT initUPLCSimplifierTrace
-  . flip runReaderT opts
-
-evalCompile
-  :: Functor m
-  => CompilationOpts name fun a
-  -> Compile m name uni fun a b
-  -> m b
-evalCompile opts = fmap fst . runCompile opts

plutus-core/plutus-core/src/PlutusCore/Compiler/Types.hs

@@ -6,19 +6,6 @@ import Data.Hashable
 import PlutusCore.Builtin
 import PlutusCore.Name.Unique
 import PlutusCore.Quote
-import UntypedPlutusCore.Core.Type qualified as UPLC
-
--- TODO1: move somewhere more appropriate?
--- TODO2: we probably don't want this in memory so after MVP
--- we should consider serializing this to disk
-newtype UPLCSimplifierTrace name uni fun a =
-  UPLCSimplifierTrace
-    { uplcSimplifierTrace
-      :: [UPLC.Term name uni fun a]
-    }
-
-initUPLCSimplifierTrace :: UPLCSimplifierTrace name uni fun a
-initUPLCSimplifierTrace = UPLCSimplifierTrace []
 
 type Compiling m uni fun name a =
   ( ToBuiltinMeaning uni fun

plutus-core/plutus-ir/test/PlutusIR/Transform/StrictLetRec/Tests/Lib.hs

@@ -75,7 +75,7 @@ compileTplcProgramOrFail
 compileTplcProgramOrFail plcProgram =
   handlePirErrorByFailing @SrcSpan =<< do
     TPLC.compileProgram plcProgram
-      & TPLC.evalCompile TPLC.defaultCompilationOpts
+      & flip runReaderT TPLC.defaultCompilationOpts
       & runQuoteT
       & runExceptT
 

plutus-core/testlib/PlutusCore/Test.hs

@@ -193,7 +193,7 @@ instance
   toUPlc =
     pure
       . TPLC.runQuote
-      . TPLC.evalCompile TPLC.defaultCompilationOpts
+      . flip runReaderT TPLC.defaultCompilationOpts
       . TPLC.compileProgram
 
 instance ToUPlc (UPLC.Program UPLC.NamedDeBruijn uni fun ()) uni fun where

plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Simplify.hs

@@ -6,6 +6,7 @@ module UntypedPlutusCore.Simplify (
     module Opts,
     simplifyTerm,
     simplifyProgram,
+    simplifyProgramWithTrace,
     InlineHints (..),
 ) where
 
@@ -21,74 +22,96 @@ import UntypedPlutusCore.Transform.Cse
 import UntypedPlutusCore.Transform.FloatDelay (floatDelay)
 import UntypedPlutusCore.Transform.ForceDelay (forceDelay)
 import UntypedPlutusCore.Transform.Inline (InlineHints (..), inline)
+import UntypedPlutusCore.Transform.Simplifier
 
 import Control.Monad
-import Control.Monad.State.Class (MonadState)
-import Control.Monad.State.Class qualified as State
 import Data.List as List (foldl')
 import Data.Typeable
 
 simplifyProgram ::
     forall name uni fun m a.
-    (Compiling m uni fun name a
-    , MonadState (UPLCSimplifierTrace name uni fun a) m
-    ) =>
+    (Compiling m uni fun name a) =>
     SimplifyOpts name a ->
     BuiltinSemanticsVariant fun ->
     Program name uni fun a ->
     m (Program name uni fun a)
 simplifyProgram opts builtinSemanticsVariant (Program a v t) =
   Program a v <$> simplifyTerm opts builtinSemanticsVariant t
 
+simplifyProgramWithTrace ::
+    forall name uni fun m a.
+    (Compiling m uni fun name a) =>
+    SimplifyOpts name a ->
+    BuiltinSemanticsVariant fun ->
+    Program name uni fun a ->
+    m (Program name uni fun a, SimplifierTrace name uni fun a)
+simplifyProgramWithTrace opts builtinSemanticsVariant (Program a v t) = do
+  (result, trace) <-
+    runSimplifierT
+      $ termSimplifier opts builtinSemanticsVariant t
+  pure (Program a v result, trace)
+
 simplifyTerm ::
     forall name uni fun m a.
-    ( Compiling m uni fun name a
-    , MonadState (UPLCSimplifierTrace name uni fun a) m
-    ) =>
+    (Compiling m uni fun name a) =>
     SimplifyOpts name a ->
     BuiltinSemanticsVariant fun ->
     Term name uni fun a ->
     m (Term name uni fun a)
-simplifyTerm opts builtinSemanticsVariant =
+simplifyTerm opts builtinSemanticsVariant term =
+  evalSimplifierT $ termSimplifier opts builtinSemanticsVariant term
+
+termSimplifier ::
+    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 opts builtinSemanticsVariant =
     simplifyNTimes (_soMaxSimplifierIterations opts) >=> cseNTimes cseTimes
   where
     -- Run the simplifier @n@ times
-    simplifyNTimes :: Int -> Term name uni fun a -> m (Term name uni fun a)
+    simplifyNTimes ::
+      Int ->
+      Term name uni fun a ->
+      SimplifierT name uni fun a m (Term name uni fun a)
     simplifyNTimes n = List.foldl' (>=>) pure $ map simplifyStep [1..n]
 
     -- Run CSE @n@ times, interleaved with the simplifier.
     -- See Note [CSE]
-    cseNTimes :: Int -> Term name uni fun a -> m (Term name uni fun a)
+    cseNTimes ::
+      Int ->
+      Term name uni fun a ->
+      SimplifierT name uni fun a m (Term name uni fun a)
     cseNTimes n = foldl' (>=>) pure $ concatMap (\i -> [cseStep i, simplifyStep i]) [1..n]
 
     -- generate simplification step
-    simplifyStep :: Int -> Term name uni fun a -> m (Term name uni fun a)
+    simplifyStep ::
+      Int ->
+      Term name uni fun a ->
+      SimplifierT name uni fun a m (Term name uni fun a)
     simplifyStep _ =
-      traceAST
-        >=> floatDelay
-        >=> traceAST
-        >=> pure . forceDelay
-        >=> traceAST
-        >=> pure . caseOfCase'
-        >=> traceAST
-        >=> pure . caseReduce
-        >=> traceAST
+        floatDelay
+        >=> forceDelay
+        >=> caseOfCase'
+        >=> caseReduce
         >=> inline (_soInlineConstants opts) (_soInlineHints opts) builtinSemanticsVariant
-        >=> traceAST
 
-    caseOfCase' :: Term name uni fun a -> Term name uni fun a
+    caseOfCase' ::
+      Term name uni fun a ->
+      SimplifierT name uni fun a m (Term name uni fun a)
     caseOfCase' = case eqT @fun @DefaultFun of
       Just Refl -> caseOfCase
-      Nothing   -> id
+      Nothing   -> pure
 
-    cseStep :: Int -> Term name uni fun a -> m (Term name uni fun a)
+    cseStep ::
+      Int ->
+      Term name uni fun a ->
+      SimplifierT name uni fun a m (Term name uni fun a)
     cseStep _ =
       case (eqT @name @Name, eqT @uni @PLC.DefaultUni) of
         (Just Refl, Just Refl) -> cse builtinSemanticsVariant
         _                      -> pure
 
-    traceAST ast = do
-      State.modify' (\st -> st { uplcSimplifierTrace = uplcSimplifierTrace st ++ [ast] })
-      return ast
-
     cseTimes = if _soConservativeOpts opts then 0 else _soMaxCseIterations opts

plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseOfCase.hs

@@ -21,11 +21,23 @@ module UntypedPlutusCore.Transform.CaseOfCase (caseOfCase) where
 import PlutusCore qualified as PLC
 import PlutusCore.MkPlc
 import UntypedPlutusCore.Core
+import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (CaseOfCase), SimplifierT,
+                                               recordSimplification)
 
 import Control.Lens
 
-caseOfCase :: (fun ~ PLC.DefaultFun) => Term name uni fun a -> Term name uni fun a
-caseOfCase = transformOf termSubterms $ \case
+caseOfCase
+    :: fun ~ PLC.DefaultFun
+    => Monad m
+    => Term name uni fun a
+    -> SimplifierT name uni fun a m (Term name uni fun a)
+caseOfCase term = do
+  let result = transformOf termSubterms processTerm term
+  recordSimplification term CaseOfCase result
+  return result
+
+processTerm :: (fun ~ PLC.DefaultFun) => Term name uni fun a -> Term name uni fun a
+processTerm = \case
   Case ann scrut alts
     | ( ite@(Force a (Builtin _ PLC.IfThenElse))
         , [cond, (trueAnn, true@Constr{}), (falseAnn, false@Constr{})]

plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseReduce.hs

@@ -6,12 +6,20 @@ module UntypedPlutusCore.Transform.CaseReduce
 
 import PlutusCore.MkPlc
 import UntypedPlutusCore.Core
+import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (CaseReduce), SimplifierT,
+                                               recordSimplification)
 
 import Control.Lens (transformOf)
 import Data.Vector qualified as V
 
-caseReduce :: Term name uni fun a -> Term name uni fun a
-caseReduce = transformOf termSubterms processTerm
+caseReduce
+    :: Monad m
+    => Term name uni fun a
+    -> SimplifierT name uni fun a m (Term name uni fun a)
+caseReduce term = do
+    let result = transformOf termSubterms processTerm term
+    recordSimplification term CaseReduce result
+    return result
 
 processTerm :: Term name uni fun a -> Term name uni fun a
 processTerm = \case

plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Cse.hs

@@ -11,6 +11,8 @@ import PlutusCore.Builtin (ToBuiltinMeaning (BuiltinSemanticsVariant))
 import UntypedPlutusCore.Core
 import UntypedPlutusCore.Purity (isWorkFree)
 import UntypedPlutusCore.Size (termSize)
+import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (CSE), SimplifierT,
+                                               recordSimplification)
 
 import Control.Arrow ((>>>))
 import Control.Lens (foldrOf, transformOf)
@@ -215,7 +217,7 @@ cse ::
   ) =>
   BuiltinSemanticsVariant fun ->
   Term Name uni fun ann ->
-  m (Term Name uni fun ann)
+  SimplifierT Name uni fun ann m (Term Name uni fun ann)
 cse builtinSemanticsVariant t0 = do
   t <- rename t0
   let annotated = annotate t
@@ -229,7 +231,9 @@ cse builtinSemanticsVariant t0 = do
           . join
           . Map.elems
           $ countOccs builtinSemanticsVariant annotated
-  mkCseTerm commonSubexprs annotated
+  result <- mkCseTerm commonSubexprs annotated
+  recordSimplification t0 CSE result
+  return result
 
 -- | The second pass. See Note [CSE].
 annotate :: Term name uni fun ann -> Term name uni fun (Path, ann)

plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/FloatDelay.hs

@@ -61,18 +61,25 @@ import PlutusCore.Name.UniqueMap qualified as UMap
 import PlutusCore.Name.UniqueSet qualified as USet
 import UntypedPlutusCore.Core.Plated (termSubterms)
 import UntypedPlutusCore.Core.Type (Term (..))
+import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (FloatDelay), SimplifierT,
+                                               recordSimplification)
 
 import Control.Lens (forOf, forOf_, transformOf)
-import Control.Monad ((>=>))
 import Control.Monad.Trans.Writer.CPS (Writer, execWriter, runWriter, tell)
 
 floatDelay ::
-  (PLC.MonadQuote m, PLC.Rename (Term name uni fun a), PLC.HasUnique name PLC.TermUnique) =>
+  ( PLC.MonadQuote m
+  , PLC.Rename (Term name uni fun a)
+  , PLC.HasUnique name PLC.TermUnique
+  ) =>
   Term name uni fun a ->
-  m (Term name uni fun a)
-floatDelay =
-  PLC.rename >=> \t ->
-    pure . uncurry (flip simplifyBodies) $ simplifyArgs (unforcedVars t) t
+  SimplifierT name uni fun a m (Term name uni fun a)
+floatDelay term = do
+  result <-
+    PLC.rename term >>= \t ->
+        pure . uncurry (flip simplifyBodies) $ simplifyArgs (unforcedVars t) t
+  recordSimplification term FloatDelay result
+  return result
 
 {- | First pass. Returns the names of all variables, at least one occurrence
 of which is not under `Force`.