From 7ff6cee109161631e888ad575fb501ba3a15f4bd Mon Sep 17 00:00:00 2001
From: Ana Pantilie <ana.pantilie95@gmail.com>
Date: Mon, 18 Aug 2025 16:00:08 +0300
Subject: [PATCH 1/5] Tmp: notes

---
 .../src/UntypedPlutusCore/Transform/Inline.hs | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
index 9bffa04a9e7..6f290484f8c 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
@@ -283,16 +283,29 @@ processTerm = handleTerm
     :: Term name uni fun a
     -> InlineM name uni fun a (Term name uni fun a)
   handleTerm = \case
+    -- If term = single variable, try to substitute it
+    -- using the substitution map from the InlineM state
     v@(Var _ n) -> fromMaybe v <$> substName n
     -- See Note [Differences from PIR inliner] 3
+
+    -- this is like "let bs (list of bindings) in t"
     (extractApps -> Just (bs, t)) -> do
+      -- it tries to find all of the substitutions which
+      -- we allow to be applied in the term, inside
+      -- processSingleBinding the state is modified to keep
+      -- track of them
       bs' <- wither (processSingleBinding t) bs
       t' <- processTerm t
+      -- !!the substitutions are not applied here!!
+      -- we just return the term from the recursive call,
+      -- with the applications which can be substituted
       pure $ restoreApps bs' t'
     t -> inlineSaturatedApp =<< forMOf termSubterms t processTerm
 
   -- See Note [Renaming strategy]
   substName :: name -> InlineM name uni fun a (Maybe (Term name uni fun a))
+  -- renaming should not matter for the certifier since there we
+  -- use DeBruijn
   substName name = gets (lookupTerm name) >>= traverse renameTerm
 
   -- See Note [Inlining approach and 'Secrets of the GHC Inliner']
@@ -310,9 +323,13 @@ processSingleBinding
   -> UTermDef name uni fun a
   -> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
 processSingleBinding body (Def vd@(UVarDecl a n) rhs0) = do
+  -- maybeAddSubst is basically just "processTerm"!
+  -- after running "processTerm" it checks whether we can really
+  -- inline the UTermDef
   maybeAddSubst body a n rhs0 >>= \case
     Just rhs -> do
       let (binders, rhsBody) = UPLC.splitParams rhs
+      -- here is where the state is modified with the new substitution
       modify' . extendVarInfo n $
         VarInfo
           { _varBinders = binders
@@ -561,6 +578,9 @@ fullyApplyAndBetaReduce info args0 = do
           if safe
             then do
               acc' <-
+                -- this needs to be changed, such that every time
+                -- a variable is substituted, we store the resulting
+                -- subterm with an incremented attribute
                 termSubstNamesM
                   ( \n ->
                       if n == param

From cafb62a30cfa5468eaa48e25afed05c246519f44 Mon Sep 17 00:00:00 2001
From: Ana Pantilie <ana.pantilie95@gmail.com>
Date: Tue, 19 Aug 2025 11:35:36 +0300
Subject: [PATCH 2/5] Tmp: add Integer to annotations

---
 .../src/UntypedPlutusCore/Transform/Inline.hs | 72 ++++++++++---------
 1 file changed, 40 insertions(+), 32 deletions(-)

diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
index 6f290484f8c..37feb2b4029 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
@@ -191,6 +191,12 @@ extendVarInfo
   -> S name uni fun a
 extendVarInfo n info s = s & vars %~ UMap.insertByName n info
 
+addAnnotationsTODO :: Term name uni fun a -> Term name uni fun (Integer, a)
+addAnnotationsTODO = undefined
+
+removeAnnotationsTODO :: Term name uni fun (Integer, a) -> Term name uni fun a
+removeAnnotationsTODO = undefined
+
 {-| Inline simple bindings. Relies on global uniqueness, and preserves it.
 See Note [Inlining and global uniqueness]
 -}
@@ -203,10 +209,10 @@ inline
   -- ^ inline constants
   -> Bool
   -- ^ preserve logging
-  -> InlineHints name a
+  -> InlineHints name (Integer, a)
   -> PLC.BuiltinSemanticsVariant fun
   -> Term name uni fun a
-  -> SimplifierT name uni fun a m (Term name uni fun a)
+  -> SimplifierT name uni fun (Integer, a) m (Term name uni fun a)
 inline
   callsiteGrowth
   inlineConstants
@@ -214,21 +220,23 @@ inline
   hints
   builtinSemanticsVariant
   t = do
+    let t' = addAnnotationsTODO t
     result <-
       liftQuote $
         flip evalStateT mempty $
           runReaderT
-            (processTerm t)
+            (processTerm t')
             InlineInfo
-              { _iiUsages = Usages.termUsages t
+              { _iiUsages = Usages.termUsages t'
               , _iiHints = hints
               , _iiBuiltinSemanticsVariant = builtinSemanticsVariant
               , _iiInlineConstants = inlineConstants
               , _iiInlineCallsiteGrowth = callsiteGrowth
               , _iiPreserveLogging = preserveLogging
               }
-    recordSimplification t Inline result
-    return result
+    recordSimplification t' Inline result
+    let result' = removeAnnotationsTODO result
+    return result'
 
 -- See Note [Differences from PIR inliner] 3
 
@@ -275,13 +283,13 @@ restoreApps defs t = makeLams [] t (reverse defs)
 processTerm
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun a
-  -> InlineM name uni fun a (Term name uni fun a)
+  => Term name uni fun (Integer, a)
+  -> InlineM name uni fun (Integer, a) (Term name uni fun (Integer, a))
 processTerm = handleTerm
  where
   handleTerm
-    :: Term name uni fun a
-    -> InlineM name uni fun a (Term name uni fun a)
+    :: Term name uni fun (Integer, a)
+    -> InlineM name uni fun (Integer, a) (Term name uni fun (Integer, a))
   handleTerm = \case
     -- If term = single variable, try to substitute it
     -- using the substitution map from the InlineM state
@@ -303,15 +311,15 @@ processTerm = handleTerm
     t -> inlineSaturatedApp =<< forMOf termSubterms t processTerm
 
   -- See Note [Renaming strategy]
-  substName :: name -> InlineM name uni fun a (Maybe (Term name uni fun a))
+  substName :: name -> InlineM name uni fun (Integer, a) (Maybe (Term name uni fun (Integer, a)))
   -- renaming should not matter for the certifier since there we
   -- use DeBruijn
   substName name = gets (lookupTerm name) >>= traverse renameTerm
 
   -- See Note [Inlining approach and 'Secrets of the GHC Inliner']
   renameTerm
-    :: InlineTerm name uni fun a
-    -> InlineM name uni fun a (Term name uni fun a)
+    :: InlineTerm name uni fun (Integer, a)
+    -> InlineM name uni fun (Integer, a) (Term name uni fun (Integer, a))
   renameTerm = \case
     -- Already processed term, just rename and put it in, don't do any
     -- further optimization here.
@@ -319,9 +327,9 @@ processTerm = handleTerm
 
 processSingleBinding
   :: (InliningConstraints name uni fun)
-  => Term name uni fun a
-  -> UTermDef name uni fun a
-  -> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
+  => Term name uni fun (Integer, a)
+  -> UTermDef name uni fun (Integer, a)
+  -> InlineM name uni fun (Integer, a) (Maybe (UTermDef name uni fun (Integer, a)))
 processSingleBinding body (Def vd@(UVarDecl a n) rhs0) = do
   -- maybeAddSubst is basically just "processTerm"!
   -- after running "processTerm" it checks whether we can really
@@ -349,11 +357,11 @@ Nothing means that we are inlining the term:
 maybeAddSubst
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun a
-  -> a
+  => Term name uni fun (Integer, a)
+  -> (Integer, a)
   -> name
-  -> Term name uni fun a
-  -> InlineM name uni fun a (Maybe (Term name uni fun a))
+  -> Term name uni fun (Integer, a)
+  -> InlineM name uni fun (Integer, a) (Maybe (Term name uni fun (Integer, a)))
 maybeAddSubst body a n rhs0 = do
   rhs <- processTerm rhs0
 
@@ -372,8 +380,8 @@ maybeAddSubst body a n rhs0 = do
  where
   extendAndDrop
     :: forall b
-     . InlineTerm name uni fun a
-    -> InlineM name uni fun a (Maybe b)
+     . InlineTerm name uni fun (Integer, a)
+    -> InlineM name uni fun (Integer, a) (Maybe b)
   extendAndDrop t = modify' (extendTerm n t) >> pure Nothing
 
 shouldUnconditionallyInline
@@ -561,16 +569,16 @@ sizeIsAcceptable inlineConstants = \case
 fullyApplyAndBetaReduce
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => VarInfo name uni fun a
-  -> [(a, Term name uni fun a)]
-  -> InlineM name uni fun a (Maybe (Term name uni fun a))
+  => VarInfo name uni fun (Integer, a)
+  -> [((Integer, a), Term name uni fun (Integer, a))]
+  -> InlineM name uni fun (Integer, a) (Maybe (Term name uni fun (Integer, a)))
 fullyApplyAndBetaReduce info args0 = do
   rhsBody <- liftDupable (let Done rhsBody = info ^. varRhsBody in rhsBody)
   let go
-        :: Term name uni fun a
+        :: Term name uni fun (Integer, a)
         -> [name]
-        -> [(a, Term name uni fun a)]
-        -> InlineM name uni fun a (Maybe (Term name uni fun a))
+        -> [((Integer, a), Term name uni fun (Integer, a))]
+        -> InlineM name uni fun (Integer, a) (Maybe (Term name uni fun (Integer, a)))
       go acc bs args = case (bs, args) of
         ([], _) -> pure . Just $ mkIterApp acc args
         (param : params, (_ann, arg) : args') -> do
@@ -597,8 +605,8 @@ fullyApplyAndBetaReduce info args0 = do
       -- inlining `a`, since inlining is the same as beta reduction.
       safeToBetaReduce
         :: name
-        -> Term name uni fun a
-        -> InlineM name uni fun a Bool
+        -> Term name uni fun (Integer, a)
+        -> InlineM name uni fun (Integer, a) Bool
       safeToBetaReduce a arg = shouldUnconditionallyInline False a arg rhsBody
   go rhsBody (info ^. varBinders) args0
 
@@ -609,8 +617,8 @@ See Note [Inlining and beta reduction of functions].
 inlineSaturatedApp
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun a
-  -> InlineM name uni fun a (Term name uni fun a)
+  => Term name uni fun (Integer, a)
+  -> InlineM name uni fun (Integer, a) (Term name uni fun (Integer, a))
 inlineSaturatedApp t
   | (Var _ann name, args) <- UPLC.splitApplication t =
       gets (lookupVarInfo name) >>= \case

From dc0aa1c2421f53de2010c777d3bfa2d940f85953 Mon Sep 17 00:00:00 2001
From: Ana Pantilie <ana.pantilie95@gmail.com>
Date: Tue, 19 Aug 2025 12:28:45 +0300
Subject: [PATCH 3/5] WIP: add SimplifierAnn

---
 .../src/UntypedPlutusCore/Transform/Inline.hs | 92 +++++++------------
 .../UntypedPlutusCore/Transform/Simplifier.hs | 15 ++-
 2 files changed, 43 insertions(+), 64 deletions(-)

diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
index 37feb2b4029..9bffa04a9e7 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
@@ -191,12 +191,6 @@ extendVarInfo
   -> S name uni fun a
 extendVarInfo n info s = s & vars %~ UMap.insertByName n info
 
-addAnnotationsTODO :: Term name uni fun a -> Term name uni fun (Integer, a)
-addAnnotationsTODO = undefined
-
-removeAnnotationsTODO :: Term name uni fun (Integer, a) -> Term name uni fun a
-removeAnnotationsTODO = undefined
-
 {-| Inline simple bindings. Relies on global uniqueness, and preserves it.
 See Note [Inlining and global uniqueness]
 -}
@@ -209,10 +203,10 @@ inline
   -- ^ inline constants
   -> Bool
   -- ^ preserve logging
-  -> InlineHints name (Integer, a)
+  -> InlineHints name a
   -> PLC.BuiltinSemanticsVariant fun
   -> Term name uni fun a
-  -> SimplifierT name uni fun (Integer, a) m (Term name uni fun a)
+  -> SimplifierT name uni fun a m (Term name uni fun a)
 inline
   callsiteGrowth
   inlineConstants
@@ -220,23 +214,21 @@ inline
   hints
   builtinSemanticsVariant
   t = do
-    let t' = addAnnotationsTODO t
     result <-
       liftQuote $
         flip evalStateT mempty $
           runReaderT
-            (processTerm t')
+            (processTerm t)
             InlineInfo
-              { _iiUsages = Usages.termUsages t'
+              { _iiUsages = Usages.termUsages t
               , _iiHints = hints
               , _iiBuiltinSemanticsVariant = builtinSemanticsVariant
               , _iiInlineConstants = inlineConstants
               , _iiInlineCallsiteGrowth = callsiteGrowth
               , _iiPreserveLogging = preserveLogging
               }
-    recordSimplification t' Inline result
-    let result' = removeAnnotationsTODO result
-    return result'
+    recordSimplification t Inline result
+    return result
 
 -- See Note [Differences from PIR inliner] 3
 
@@ -283,43 +275,30 @@ restoreApps defs t = makeLams [] t (reverse defs)
 processTerm
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun (Integer, a)
-  -> InlineM name uni fun (Integer, a) (Term name uni fun (Integer, a))
+  => Term name uni fun a
+  -> InlineM name uni fun a (Term name uni fun a)
 processTerm = handleTerm
  where
   handleTerm
-    :: Term name uni fun (Integer, a)
-    -> InlineM name uni fun (Integer, a) (Term name uni fun (Integer, a))
+    :: Term name uni fun a
+    -> InlineM name uni fun a (Term name uni fun a)
   handleTerm = \case
-    -- If term = single variable, try to substitute it
-    -- using the substitution map from the InlineM state
     v@(Var _ n) -> fromMaybe v <$> substName n
     -- See Note [Differences from PIR inliner] 3
-
-    -- this is like "let bs (list of bindings) in t"
     (extractApps -> Just (bs, t)) -> do
-      -- it tries to find all of the substitutions which
-      -- we allow to be applied in the term, inside
-      -- processSingleBinding the state is modified to keep
-      -- track of them
       bs' <- wither (processSingleBinding t) bs
       t' <- processTerm t
-      -- !!the substitutions are not applied here!!
-      -- we just return the term from the recursive call,
-      -- with the applications which can be substituted
       pure $ restoreApps bs' t'
     t -> inlineSaturatedApp =<< forMOf termSubterms t processTerm
 
   -- See Note [Renaming strategy]
-  substName :: name -> InlineM name uni fun (Integer, a) (Maybe (Term name uni fun (Integer, a)))
-  -- renaming should not matter for the certifier since there we
-  -- use DeBruijn
+  substName :: name -> InlineM name uni fun a (Maybe (Term name uni fun a))
   substName name = gets (lookupTerm name) >>= traverse renameTerm
 
   -- See Note [Inlining approach and 'Secrets of the GHC Inliner']
   renameTerm
-    :: InlineTerm name uni fun (Integer, a)
-    -> InlineM name uni fun (Integer, a) (Term name uni fun (Integer, a))
+    :: InlineTerm name uni fun a
+    -> InlineM name uni fun a (Term name uni fun a)
   renameTerm = \case
     -- Already processed term, just rename and put it in, don't do any
     -- further optimization here.
@@ -327,17 +306,13 @@ processTerm = handleTerm
 
 processSingleBinding
   :: (InliningConstraints name uni fun)
-  => Term name uni fun (Integer, a)
-  -> UTermDef name uni fun (Integer, a)
-  -> InlineM name uni fun (Integer, a) (Maybe (UTermDef name uni fun (Integer, a)))
+  => Term name uni fun a
+  -> UTermDef name uni fun a
+  -> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
 processSingleBinding body (Def vd@(UVarDecl a n) rhs0) = do
-  -- maybeAddSubst is basically just "processTerm"!
-  -- after running "processTerm" it checks whether we can really
-  -- inline the UTermDef
   maybeAddSubst body a n rhs0 >>= \case
     Just rhs -> do
       let (binders, rhsBody) = UPLC.splitParams rhs
-      -- here is where the state is modified with the new substitution
       modify' . extendVarInfo n $
         VarInfo
           { _varBinders = binders
@@ -357,11 +332,11 @@ Nothing means that we are inlining the term:
 maybeAddSubst
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun (Integer, a)
-  -> (Integer, a)
+  => Term name uni fun a
+  -> a
   -> name
-  -> Term name uni fun (Integer, a)
-  -> InlineM name uni fun (Integer, a) (Maybe (Term name uni fun (Integer, a)))
+  -> Term name uni fun a
+  -> InlineM name uni fun a (Maybe (Term name uni fun a))
 maybeAddSubst body a n rhs0 = do
   rhs <- processTerm rhs0
 
@@ -380,8 +355,8 @@ maybeAddSubst body a n rhs0 = do
  where
   extendAndDrop
     :: forall b
-     . InlineTerm name uni fun (Integer, a)
-    -> InlineM name uni fun (Integer, a) (Maybe b)
+     . InlineTerm name uni fun a
+    -> InlineM name uni fun a (Maybe b)
   extendAndDrop t = modify' (extendTerm n t) >> pure Nothing
 
 shouldUnconditionallyInline
@@ -569,16 +544,16 @@ sizeIsAcceptable inlineConstants = \case
 fullyApplyAndBetaReduce
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => VarInfo name uni fun (Integer, a)
-  -> [((Integer, a), Term name uni fun (Integer, a))]
-  -> InlineM name uni fun (Integer, a) (Maybe (Term name uni fun (Integer, a)))
+  => VarInfo name uni fun a
+  -> [(a, Term name uni fun a)]
+  -> InlineM name uni fun a (Maybe (Term name uni fun a))
 fullyApplyAndBetaReduce info args0 = do
   rhsBody <- liftDupable (let Done rhsBody = info ^. varRhsBody in rhsBody)
   let go
-        :: Term name uni fun (Integer, a)
+        :: Term name uni fun a
         -> [name]
-        -> [((Integer, a), Term name uni fun (Integer, a))]
-        -> InlineM name uni fun (Integer, a) (Maybe (Term name uni fun (Integer, a)))
+        -> [(a, Term name uni fun a)]
+        -> InlineM name uni fun a (Maybe (Term name uni fun a))
       go acc bs args = case (bs, args) of
         ([], _) -> pure . Just $ mkIterApp acc args
         (param : params, (_ann, arg) : args') -> do
@@ -586,9 +561,6 @@ fullyApplyAndBetaReduce info args0 = do
           if safe
             then do
               acc' <-
-                -- this needs to be changed, such that every time
-                -- a variable is substituted, we store the resulting
-                -- subterm with an incremented attribute
                 termSubstNamesM
                   ( \n ->
                       if n == param
@@ -605,8 +577,8 @@ fullyApplyAndBetaReduce info args0 = do
       -- inlining `a`, since inlining is the same as beta reduction.
       safeToBetaReduce
         :: name
-        -> Term name uni fun (Integer, a)
-        -> InlineM name uni fun (Integer, a) Bool
+        -> Term name uni fun a
+        -> InlineM name uni fun a Bool
       safeToBetaReduce a arg = shouldUnconditionallyInline False a arg rhsBody
   go rhsBody (info ^. varBinders) args0
 
@@ -617,8 +589,8 @@ See Note [Inlining and beta reduction of functions].
 inlineSaturatedApp
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun (Integer, a)
-  -> InlineM name uni fun (Integer, a) (Term name uni fun (Integer, a))
+  => Term name uni fun a
+  -> InlineM name uni fun a (Term name uni fun a)
 inlineSaturatedApp t
   | (Var _ann name, args) <- UPLC.splitApplication t =
       gets (lookupVarInfo name) >>= \case
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Simplifier.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Simplifier.hs
index 1bc15c5a97e..cf5f62bbffb 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Simplifier.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Simplifier.hs
@@ -4,6 +4,7 @@ module UntypedPlutusCore.Transform.Simplifier (
     SimplifierT (..),
     SimplifierTrace (..),
     SimplifierStage (..),
+    SimplifierAnn (..),
     Simplification (..),
     runSimplifierT,
     evalSimplifierT,
@@ -64,11 +65,17 @@ data SimplifierStage
   | Inline
   | CSE
 
+data SimplifierAnn a =
+  SimplifierAnn
+    { inlineCounter :: Integer
+    , ann           :: a
+    }
+
 data Simplification name uni fun a =
   Simplification
-    { beforeAST :: Term name uni fun a
+    { beforeAST :: Term name uni fun (SimplifierAnn a)
     , stage     :: SimplifierStage
-    , afterAST  :: Term name uni fun a
+    , afterAST  :: Term name uni fun (SimplifierAnn a)
     }
 
 -- TODO2: we probably don't want this in memory so after MVP
@@ -84,9 +91,9 @@ initSimplifierTrace = SimplifierTrace []
 
 recordSimplification
   :: Monad m
-  => Term name uni fun a
+  => Term name uni fun (SimplifierAnn a)
   -> SimplifierStage
-  -> Term name uni fun a
+  -> Term name uni fun (SimplifierAnn a)
   -> SimplifierT name uni fun a m ()
 recordSimplification beforeAST stage afterAST =
   let simplification = Simplification { beforeAST, stage, afterAST }

From 900a1de1f417468b6ae089ea0e3d3a247908afa9 Mon Sep 17 00:00:00 2001
From: Ana Pantilie <ana.pantilie95@gmail.com>
Date: Tue, 19 Aug 2025 18:46:31 +0300
Subject: [PATCH 4/5] Add SimplifierAnn to simplification code

---
 .../src/UntypedPlutusCore/Core/Type.hs        |  14 +++
 .../UntypedPlutusCore/Transform/CaseOfCase.hs |   7 +-
 .../UntypedPlutusCore/Transform/CaseReduce.hs |   7 +-
 .../src/UntypedPlutusCore/Transform/Cse.hs    |   7 +-
 .../UntypedPlutusCore/Transform/FloatDelay.hs |   7 +-
 .../Transform/ForceCaseDelay.hs               |   7 +-
 .../UntypedPlutusCore/Transform/ForceDelay.hs |   7 +-
 .../src/UntypedPlutusCore/Transform/Inline.hs | 103 ++++++++++--------
 .../UntypedPlutusCore/Transform/Simplifier.hs |  24 +++-
 9 files changed, 121 insertions(+), 62 deletions(-)

diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Core/Type.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Core/Type.hs
index 55e4d89d512..36801df504f 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Core/Type.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Core/Type.hs
@@ -19,6 +19,7 @@ module UntypedPlutusCore.Core.Type
     , bindFunM
     , bindFun
     , mapFun
+    , mapAnn
     , termAnn
     , UVarDecl(..)
     , uvarDeclName
@@ -194,3 +195,16 @@ bindFun f = runIdentity . bindFunM (coerce f)
 
 mapFun :: (ann -> fun -> fun') -> Term name uni fun ann -> Term name uni fun' ann
 mapFun f = bindFun $ \ann fun -> Builtin ann (f ann fun)
+
+mapAnn :: (ann -> ann') -> Term name uni fun ann -> Term name uni fun ann'
+mapAnn f = go where
+    go (Constant ann val)     = Constant (f ann) val
+    go (Builtin ann fun)      = Builtin (f ann) fun
+    go (Var ann name)         = Var (f ann) name
+    go (LamAbs ann name body) = LamAbs (f ann) name (go body)
+    go (Apply ann fun arg)    = Apply (f ann) (go fun) (go arg)
+    go (Delay ann term)       = Delay (f ann) (go term)
+    go (Force ann term)       = Force (f ann) (go term)
+    go (Error ann)            = Error (f ann)
+    go (Constr ann i args)    = Constr (f ann) i $ fmap go args
+    go (Case ann arg cs)      = Case (f ann) (go arg) $ fmap go cs
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseOfCase.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseOfCase.hs
index a2261d35a93..b4b85d6f6a3 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseOfCase.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseOfCase.hs
@@ -44,7 +44,7 @@ import PlutusCore.MkPlc (mkIterApp)
 import UntypedPlutusCore.Core
 import UntypedPlutusCore.Transform.CaseReduce qualified as CaseReduce
 import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (CaseOfCase), SimplifierT,
-                                               recordSimplification)
+                                               initSimplifierTerm, recordSimplification)
 
 import Control.Lens
 import Data.List (nub)
@@ -57,7 +57,10 @@ caseOfCase
     -> SimplifierT name uni fun a m (Term name uni fun a)
 caseOfCase term = do
   let result = transformOf termSubterms processTerm term
-  recordSimplification term CaseOfCase result
+  recordSimplification
+    (initSimplifierTerm term)
+    CaseOfCase
+    (initSimplifierTerm result)
   return result
 
 processTerm
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseReduce.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseReduce.hs
index 4d9d04643c2..c9c5e8a8d24 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseReduce.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseReduce.hs
@@ -12,7 +12,7 @@ import PlutusCore.Builtin (CaseBuiltin (..))
 import PlutusCore.MkPlc
 import UntypedPlutusCore.Core
 import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (CaseReduce), SimplifierT,
-                                               recordSimplification)
+                                               initSimplifierTerm, recordSimplification)
 
 caseReduce
     :: (Monad m, CaseBuiltin uni)
@@ -20,7 +20,10 @@ caseReduce
     -> SimplifierT name uni fun a m (Term name uni fun a)
 caseReduce term = do
     let result = transformOf termSubterms processTerm term
-    recordSimplification term CaseReduce result
+    recordSimplification
+       (initSimplifierTerm term)
+       CaseReduce
+       (initSimplifierTerm result)
     return result
 
 processTerm :: CaseBuiltin uni => Term name uni fun a -> Term name uni fun a
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Cse.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Cse.hs
index 6663f314484..69c7a41fe7d 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Cse.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Cse.hs
@@ -12,7 +12,7 @@ import UntypedPlutusCore.Core
 import UntypedPlutusCore.Purity (isWorkFree)
 import UntypedPlutusCore.Size (termSize)
 import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (CSE), SimplifierT,
-                                               recordSimplification)
+                                               initSimplifierTerm, recordSimplification)
 
 import Control.Arrow ((>>>))
 import Control.Lens (foldrOf, transformOf)
@@ -232,7 +232,10 @@ cse builtinSemanticsVariant t0 = do
           . Map.elems
           $ countOccs builtinSemanticsVariant annotated
   result <- mkCseTerm commonSubexprs annotated
-  recordSimplification t0 CSE result
+  recordSimplification
+    (initSimplifierTerm t0)
+    CSE
+    (initSimplifierTerm result)
   return result
 
 -- | The second pass. See Note [CSE].
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/FloatDelay.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/FloatDelay.hs
index 4070d3d0fed..4af6f038a79 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/FloatDelay.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/FloatDelay.hs
@@ -62,7 +62,7 @@ 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)
+                                               initSimplifierTerm, recordSimplification)
 
 import Control.Lens (forOf, forOf_, transformOf)
 import Control.Monad.Trans.Writer.CPS (Writer, execWriter, runWriter, tell)
@@ -78,7 +78,10 @@ floatDelay term = do
   result <-
     PLC.rename term >>= \t ->
         pure . uncurry (flip simplifyBodies) $ simplifyArgs (unforcedVars t) t
-  recordSimplification term FloatDelay result
+  recordSimplification
+    (initSimplifierTerm term)
+    FloatDelay
+    (initSimplifierTerm result)
   return result
 
 {- | First pass. Returns the names of all variables, at least one occurrence
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/ForceCaseDelay.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/ForceCaseDelay.hs
index 746385cf093..47c13f470e5 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/ForceCaseDelay.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/ForceCaseDelay.hs
@@ -35,7 +35,7 @@ where
 
 import UntypedPlutusCore.Core
 import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (ForceCaseDelay), SimplifierT,
-                                               recordSimplification)
+                                               initSimplifierTerm, recordSimplification)
 
 import Control.Lens
 
@@ -45,7 +45,10 @@ forceCaseDelay
     -> SimplifierT name uni fun a m (Term name uni fun a)
 forceCaseDelay term = do
   let result = transformOf termSubterms processTerm term
-  recordSimplification term ForceCaseDelay result
+  recordSimplification
+    (initSimplifierTerm term)
+    ForceCaseDelay
+    (initSimplifierTerm result)
   return result
 
 processTerm :: Term name uni fun a -> Term name uni fun a
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/ForceDelay.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/ForceDelay.hs
index 8199af10b38..e37c2a311e8 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/ForceDelay.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/ForceDelay.hs
@@ -154,7 +154,7 @@ import PlutusCore.MkPlc (mkIterApp)
 import UntypedPlutusCore.Core
 import UntypedPlutusCore.Purity (isPure, isWorkFree)
 import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (ForceDelay), SimplifierT,
-                                               recordSimplification)
+                                               initSimplifierTerm, recordSimplification)
 
 import Control.Lens (transformOf)
 import Control.Monad (guard)
@@ -170,7 +170,10 @@ forceDelay
     -> SimplifierT name uni fun a m (Term name uni fun a)
 forceDelay semVar term = do
     let result = transformOf termSubterms (processTerm semVar) term
-    recordSimplification term ForceDelay result
+    recordSimplification
+        (initSimplifierTerm term)
+        ForceDelay
+        (initSimplifierTerm result)
     return result
 
 {- | Checks whether the term is of the right form, and "pushes"
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
index 9bffa04a9e7..4b629be2ff5 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Inline.hs
@@ -59,8 +59,9 @@ import UntypedPlutusCore.Purity (EvalTerm (EvalTerm, Unknown), Purity (MaybeImpu
 import UntypedPlutusCore.Rename ()
 import UntypedPlutusCore.Size (Size, termSize)
 import UntypedPlutusCore.Subst (termSubstNamesM)
-import UntypedPlutusCore.Transform.Simplifier (SimplifierStage (Inline), SimplifierT,
-                                               recordSimplification)
+import UntypedPlutusCore.Transform.Simplifier (SimplifierAnn (..), SimplifierStage (Inline),
+                                               SimplifierT, SimplifierTerm, eraseSimplifierAnn,
+                                               initSimplifierTerm, recordSimplification)
 import Witherable (wither)
 
 {- Note [Differences from PIR inliner]
@@ -77,7 +78,9 @@ the PIR inliner.
 -}
 
 -- | Substitution range, 'SubstRng' in the paper.
-newtype InlineTerm name uni fun a = Done (Dupable (Term name uni fun a))
+newtype InlineTerm name uni fun a = Done (Dupable (SimplifierTerm name uni fun a))
+
+type IUTermDef name uni fun a = UTermDef name uni fun (SimplifierAnn a)
 
 {-| Term substitution, 'Subst' in the paper.
 A map of unprocessed variable and its substitution range.
@@ -99,7 +102,7 @@ makeLenses ''Subst
 data VarInfo name uni fun ann = VarInfo
   { _varBinders :: [name]
   -- ^ Lambda binders in the RHS (definition) of the variable.
-  , _varRhs     :: Term name uni fun ann
+  , _varRhs     :: SimplifierTerm name uni fun ann
   -- ^ The RHS (definition) of the variable.
   , _varRhsBody :: InlineTerm name uni fun ann
   {- ^ The body of the RHS of the variable (i.e., RHS minus the binders).
@@ -214,21 +217,23 @@ inline
   hints
   builtinSemanticsVariant
   t = do
+    let simplTerm :: SimplifierTerm name uni fun a
+        simplTerm = initSimplifierTerm t
     result <-
       liftQuote $
         flip evalStateT mempty $
           runReaderT
-            (processTerm t)
+            (processTerm simplTerm)
             InlineInfo
-              { _iiUsages = Usages.termUsages t
+              { _iiUsages = Usages.termUsages simplTerm
               , _iiHints = hints
               , _iiBuiltinSemanticsVariant = builtinSemanticsVariant
               , _iiInlineConstants = inlineConstants
               , _iiInlineCallsiteGrowth = callsiteGrowth
               , _iiPreserveLogging = preserveLogging
               }
-    recordSimplification t Inline result
-    return result
+    recordSimplification simplTerm Inline result
+    return (eraseSimplifierAnn result)
 
 -- See Note [Differences from PIR inliner] 3
 
@@ -243,12 +248,22 @@ Some examples will help:
 [[(\x . t) a] b] -> Nothing
 -}
 extractApps
-  :: Term name uni fun a
-  -> Maybe ([UTermDef name uni fun a], Term name uni fun a)
+  :: SimplifierTerm name uni fun a
+  -> Maybe ([IUTermDef name uni fun a], SimplifierTerm name uni fun a)
 extractApps = go []
  where
+  go
+    :: [SimplifierTerm name uni fun a]
+    -> SimplifierTerm name uni fun a
+    -> Maybe ([IUTermDef name uni fun a], SimplifierTerm name uni fun a)
   go argStack (Apply _ f arg) = go (arg : argStack) f
   go argStack t               = matchArgs argStack [] t
+
+  matchArgs
+    :: [SimplifierTerm name uni fun a]
+    -> [IUTermDef name uni fun a]
+    -> SimplifierTerm name uni fun a
+    -> Maybe ([IUTermDef name uni fun a], SimplifierTerm name uni fun a)
   matchArgs (arg : rest) acc (LamAbs a n body) =
     matchArgs rest (Def (UVarDecl a n) arg : acc) body
   matchArgs [] acc t =
@@ -257,9 +272,9 @@ extractApps = go []
 
 -- | The inverse of 'extractApps'.
 restoreApps
-  :: [UTermDef name uni fun a]
-  -> Term name uni fun a
-  -> Term name uni fun a
+  :: [IUTermDef name uni fun a]
+  -> SimplifierTerm name uni fun a
+  -> SimplifierTerm name uni fun a
 restoreApps defs t = makeLams [] t (reverse defs)
  where
   makeLams args acc (Def (UVarDecl a n) rhs : rest) =
@@ -275,13 +290,13 @@ restoreApps defs t = makeLams [] t (reverse defs)
 processTerm
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun a
-  -> InlineM name uni fun a (Term name uni fun a)
+  => SimplifierTerm name uni fun a
+  -> InlineM name uni fun a (SimplifierTerm name uni fun a)
 processTerm = handleTerm
  where
   handleTerm
-    :: Term name uni fun a
-    -> InlineM name uni fun a (Term name uni fun a)
+    :: SimplifierTerm name uni fun a
+    -> InlineM name uni fun a (SimplifierTerm name uni fun a)
   handleTerm = \case
     v@(Var _ n) -> fromMaybe v <$> substName n
     -- See Note [Differences from PIR inliner] 3
@@ -292,13 +307,13 @@ processTerm = handleTerm
     t -> inlineSaturatedApp =<< forMOf termSubterms t processTerm
 
   -- See Note [Renaming strategy]
-  substName :: name -> InlineM name uni fun a (Maybe (Term name uni fun a))
+  substName :: name -> InlineM name uni fun a (Maybe (SimplifierTerm name uni fun a))
   substName name = gets (lookupTerm name) >>= traverse renameTerm
 
   -- See Note [Inlining approach and 'Secrets of the GHC Inliner']
   renameTerm
     :: InlineTerm name uni fun a
-    -> InlineM name uni fun a (Term name uni fun a)
+    -> InlineM name uni fun a (SimplifierTerm name uni fun a)
   renameTerm = \case
     -- Already processed term, just rename and put it in, don't do any
     -- further optimization here.
@@ -306,9 +321,9 @@ processTerm = handleTerm
 
 processSingleBinding
   :: (InliningConstraints name uni fun)
-  => Term name uni fun a
-  -> UTermDef name uni fun a
-  -> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
+  => SimplifierTerm name uni fun a
+  -> IUTermDef name uni fun a
+  -> InlineM name uni fun a (Maybe (IUTermDef name uni fun a))
 processSingleBinding body (Def vd@(UVarDecl a n) rhs0) = do
   maybeAddSubst body a n rhs0 >>= \case
     Just rhs -> do
@@ -332,17 +347,17 @@ Nothing means that we are inlining the term:
 maybeAddSubst
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun a
-  -> a
+  => SimplifierTerm name uni fun a
+  -> SimplifierAnn a
   -> name
-  -> Term name uni fun a
-  -> InlineM name uni fun a (Maybe (Term name uni fun a))
+  -> SimplifierTerm name uni fun a
+  -> InlineM name uni fun a (Maybe (SimplifierTerm name uni fun a))
 maybeAddSubst body a n rhs0 = do
   rhs <- processTerm rhs0
 
   -- Check whether we've been told specifically to inline this
   hints <- view iiHints
-  case shouldInline hints a n of
+  case shouldInline hints (otherAnn a) n of
     AlwaysInline ->
       -- if we've been told specifically, then do it right away
       extendAndDrop (Done $ dupable rhs)
@@ -366,8 +381,8 @@ shouldUnconditionallyInline
   If so, bypass the purity check.
   -}
   -> name
-  -> Term name uni fun a
-  -> Term name uni fun a
+  -> SimplifierTerm name uni fun a
+  -> SimplifierTerm name uni fun a
   -> InlineM name uni fun a Bool
 shouldUnconditionallyInline safe n rhs body = do
   isTermPure <- checkPurity rhs
@@ -391,7 +406,7 @@ shouldUnconditionallyInline safe n rhs body = do
 -- | Check if term is pure. See Note [Inlining and purity]
 checkPurity
   :: (PLC.ToBuiltinMeaning uni fun)
-  => Term name uni fun a
+  => SimplifierTerm name uni fun a
   -> InlineM name uni fun a Bool
 checkPurity t = do
   builtinSemanticsVariant <- view iiBuiltinSemanticsVariant
@@ -413,7 +428,7 @@ isFirstVarBeforeEffects
    . (InliningConstraints name uni fun)
   => BuiltinSemanticsVariant fun
   -> name
-  -> Term name uni fun ann
+  -> SimplifierTerm name uni fun ann
   -> Bool
 isFirstVarBeforeEffects builtinSemanticsVariant n t =
   -- This can in the worst case traverse a lot of the term, which could lead to
@@ -442,11 +457,11 @@ isStrictIn
   :: forall name uni fun a
    . (Eq name)
   => name
-  -> Term name uni fun a
+  -> SimplifierTerm name uni fun a
   -> Bool
 isStrictIn name = go
  where
-  go :: Term name uni fun a -> Bool
+  go :: SimplifierTerm name uni fun a -> Bool
   go = \case
     Var _ann name' -> name == name'
     LamAbs _ann _paramName _body -> False
@@ -462,7 +477,7 @@ isStrictIn name = go
 effectSafe
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun a
+  => SimplifierTerm name uni fun a
   -> name
   -> Bool
   -- ^ is it pure? See Note [Inlining and purity]
@@ -481,7 +496,7 @@ or code.  See Note [Inlining approach and 'Secrets of the GHC Inliner']
 acceptable
   :: Bool
   -- ^ inline constants
-  -> Term name uni fun a
+  -> SimplifierTerm name uni fun a
   -> InlineM name uni fun a Bool
 acceptable inlineConstants t =
   -- See Note [Inlining criteria]
@@ -518,7 +533,7 @@ the given term acceptable?
 sizeIsAcceptable
   :: Bool
   -- ^ inline constants
-  -> Term name uni fun a
+  -> SimplifierTerm name uni fun a
   -> Bool
 sizeIsAcceptable inlineConstants = \case
   Builtin{} -> True
@@ -545,15 +560,15 @@ fullyApplyAndBetaReduce
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
   => VarInfo name uni fun a
-  -> [(a, Term name uni fun a)]
-  -> InlineM name uni fun a (Maybe (Term name uni fun a))
+  -> [(SimplifierAnn a, SimplifierTerm name uni fun a)]
+  -> InlineM name uni fun a (Maybe (SimplifierTerm name uni fun a))
 fullyApplyAndBetaReduce info args0 = do
   rhsBody <- liftDupable (let Done rhsBody = info ^. varRhsBody in rhsBody)
   let go
-        :: Term name uni fun a
+        :: SimplifierTerm name uni fun a
         -> [name]
-        -> [(a, Term name uni fun a)]
-        -> InlineM name uni fun a (Maybe (Term name uni fun a))
+        -> [(SimplifierAnn a, SimplifierTerm name uni fun a)]
+        -> InlineM name uni fun a (Maybe (SimplifierTerm name uni fun a))
       go acc bs args = case (bs, args) of
         ([], _) -> pure . Just $ mkIterApp acc args
         (param : params, (_ann, arg) : args') -> do
@@ -577,7 +592,7 @@ fullyApplyAndBetaReduce info args0 = do
       -- inlining `a`, since inlining is the same as beta reduction.
       safeToBetaReduce
         :: name
-        -> Term name uni fun a
+        -> SimplifierTerm name uni fun a
         -> InlineM name uni fun a Bool
       safeToBetaReduce a arg = shouldUnconditionallyInline False a arg rhsBody
   go rhsBody (info ^. varBinders) args0
@@ -589,8 +604,8 @@ See Note [Inlining and beta reduction of functions].
 inlineSaturatedApp
   :: forall name uni fun a
    . (InliningConstraints name uni fun)
-  => Term name uni fun a
-  -> InlineM name uni fun a (Term name uni fun a)
+  => SimplifierTerm name uni fun a
+  -> InlineM name uni fun a (SimplifierTerm name uni fun a)
 inlineSaturatedApp t
   | (Var _ann name, args) <- UPLC.splitApplication t =
       gets (lookupVarInfo name) >>= \case
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Simplifier.hs b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Simplifier.hs
index cf5f62bbffb..cd6fd42ce66 100644
--- a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Simplifier.hs
+++ b/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Simplifier.hs
@@ -15,6 +15,9 @@ module UntypedPlutusCore.Transform.Simplifier (
     execSimplifier,
     initSimplifierTrace,
     recordSimplification,
+    SimplifierTerm,
+    initSimplifierTerm,
+    eraseSimplifierAnn,
 ) where
 
 import Control.Monad.State (MonadTrans, StateT)
@@ -22,7 +25,7 @@ import Control.Monad.State qualified as State
 
 import Control.Monad.Identity (Identity, runIdentity)
 import PlutusCore.Quote (MonadQuote)
-import UntypedPlutusCore.Core.Type (Term)
+import UntypedPlutusCore.Core.Type (Term, mapAnn)
 
 newtype SimplifierT name uni fun ann m a =
   SimplifierT
@@ -68,14 +71,23 @@ data SimplifierStage
 data SimplifierAnn a =
   SimplifierAnn
     { inlineCounter :: Integer
-    , ann           :: a
+    , otherAnn      :: a
     }
 
+type SimplifierTerm name uni fun a =
+  Term name uni fun (SimplifierAnn a)
+
+initSimplifierTerm :: Term name uni fun a -> SimplifierTerm name uni fun a
+initSimplifierTerm = mapAnn (\otherAnn -> SimplifierAnn { inlineCounter = 0, otherAnn })
+
+eraseSimplifierAnn :: SimplifierTerm name uni fun a -> Term name uni fun a
+eraseSimplifierAnn = mapAnn (\SimplifierAnn { otherAnn } -> otherAnn)
+
 data Simplification name uni fun a =
   Simplification
-    { beforeAST :: Term name uni fun (SimplifierAnn a)
+    { beforeAST :: SimplifierTerm name uni fun a
     , stage     :: SimplifierStage
-    , afterAST  :: Term name uni fun (SimplifierAnn a)
+    , afterAST  :: SimplifierTerm name uni fun a
     }
 
 -- TODO2: we probably don't want this in memory so after MVP
@@ -91,9 +103,9 @@ initSimplifierTrace = SimplifierTrace []
 
 recordSimplification
   :: Monad m
-  => Term name uni fun (SimplifierAnn a)
+  => SimplifierTerm name uni fun a
   -> SimplifierStage
-  -> Term name uni fun (SimplifierAnn a)
+  -> SimplifierTerm name uni fun a
   -> SimplifierT name uni fun a m ()
 recordSimplification beforeAST stage afterAST =
   let simplification = Simplification { beforeAST, stage, afterAST }

From 3cd6dfe72178dd38115ad9c61ceba3f67f51e9f8 Mon Sep 17 00:00:00 2001
From: Ana Pantilie <ana.pantilie95@gmail.com>
Date: Wed, 20 Aug 2025 12:25:57 +0300
Subject: [PATCH 5/5] Fix compilation errors

---
 .../testlib/Transform/Inline/Spec.hs          | 36 ++++++++++++-------
 .../test/certifier/Test/Certifier/AST.hs      |  9 ++---
 2 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/plutus-core/untyped-plutus-core/testlib/Transform/Inline/Spec.hs b/plutus-core/untyped-plutus-core/testlib/Transform/Inline/Spec.hs
index 9d6a3b6de5c..48b9859f5c5 100644
--- a/plutus-core/untyped-plutus-core/testlib/Transform/Inline/Spec.hs
+++ b/plutus-core/untyped-plutus-core/testlib/Transform/Inline/Spec.hs
@@ -16,6 +16,7 @@ import PlutusCore.Quote (runQuote)
 import PlutusPrelude (def)
 import Test.Tasty (TestTree, testGroup)
 import Test.Tasty.HUnit (Assertion, assertBool, testCase)
+import UntypedPlutusCore (SimplifierTerm, initSimplifierTerm)
 import UntypedPlutusCore.Core (Term (..))
 import UntypedPlutusCore.Size (Size (..))
 import UntypedPlutusCore.Transform.Inline (InlineHints (..), InlineInfo (..), InlineM, S (..),
@@ -57,7 +58,7 @@ testVarBeforeAfterEffects = do
   assertBool "c is not evaluated before effects" $ not do
     isFirstVarBeforeEffects def c term
  where
-  term :: Term Name DefaultUni DefaultFun ()
+  term :: SimplifierTerm Name DefaultUni DefaultFun ()
   term =
     {- Evaluation order:
 
@@ -67,7 +68,8 @@ testVarBeforeAfterEffects = do
       4. pure work-free: c
       5. impure? maybe work?: addInteger (addInteger a b) c
     -}
-    addInteger (addInteger (var a) (var b)) (var c)
+    initSimplifierTerm
+    $ addInteger (addInteger (var a) (var b)) (var c)
   (a, b, c, _) = makeUniqueNames
 
 testVarIsEventuallyEvaluatedDelay :: Assertion
@@ -79,8 +81,10 @@ testVarIsEventuallyEvaluatedDelay = do
   assertBool "it's not known if var 'c' is eventually evaluated" $
     not (isStrictIn c term)
  where
-  term :: Term Name DefaultUni DefaultFun ()
-  term = delay (var a `addInteger` var b) `addInteger` var b
+  term :: SimplifierTerm Name DefaultUni DefaultFun ()
+  term =
+    initSimplifierTerm
+    $ delay (var a `addInteger` var b) `addInteger` var b
 
   (a, b, c, _) = makeUniqueNames
 
@@ -93,8 +97,10 @@ testVarIsEventuallyEvaluatedLambda = do
   assertBool "it's not known if var 'd' is eventually evaluated" $
     not (isStrictIn d term)
  where
-  term :: Term Name DefaultUni DefaultFun ()
-  term = lam b (var a `addInteger` var c) `app` var c
+  term :: SimplifierTerm Name DefaultUni DefaultFun ()
+  term =
+    initSimplifierTerm
+    $ lam b (var a `addInteger` var c) `app` var c
 
   (a, b, c, d) = makeUniqueNames
 
@@ -107,8 +113,10 @@ testVarIsEventuallyEvaluatedCaseBranch = do
   assertBool "it is not known if var 'd' is eventually evaluated" $
     not (isStrictIn d term)
  where
-  term :: Term Name DefaultUni DefaultFun ()
-  term = case_ (var b) [var a, var b, var c]
+  term :: SimplifierTerm Name DefaultUni DefaultFun ()
+  term =
+    initSimplifierTerm
+    $ case_ (var b) [var a, var b, var c]
 
   (a, b, c, d) = makeUniqueNames
 
@@ -119,8 +127,10 @@ testEffectSafePreservedLogs = do
   assertBool "a var before effects is \"effect safe\"" $
     runInlineWithLogging (effectSafe term a False)
  where
-  term :: Term Name DefaultUni DefaultFun ()
-  term = (var a `addInteger` var b) `addInteger` var c
+  term :: SimplifierTerm Name DefaultUni DefaultFun ()
+  term =
+    initSimplifierTerm
+    $ (var a `addInteger` var b) `addInteger` var c
 
   (a, b, c, _) = makeUniqueNames
 
@@ -131,8 +141,10 @@ testEffectSafeWithoutPreservedLogs = do
   assertBool "a var before effects is \"effect safe\"" $
     runInlineWithoutLogging (effectSafe term a False)
  where
-  term :: Term Name DefaultUni DefaultFun ()
-  term = (var a `addInteger` var b) `addInteger` var c
+  term :: SimplifierTerm Name DefaultUni DefaultFun ()
+  term =
+    initSimplifierTerm
+    $ (var a `addInteger` var b) `addInteger` var c
 
   (a, b, c, _) = makeUniqueNames
 
diff --git a/plutus-metatheory/test/certifier/Test/Certifier/AST.hs b/plutus-metatheory/test/certifier/Test/Certifier/AST.hs
index cefdbd8433b..e70723cb513 100644
--- a/plutus-metatheory/test/certifier/Test/Certifier/AST.hs
+++ b/plutus-metatheory/test/certifier/Test/Certifier/AST.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE ViewPatterns #-}
 module Test.Certifier.AST where
 
 import PlutusCore qualified as PLC
@@ -13,8 +14,8 @@ import Test.Tasty.HUnit
 
 mkMockTracePair
   :: SimplifierStage
-  -> Term Name DefaultUni DefaultFun ()
-  -> Term Name DefaultUni DefaultFun ()
+  -> SimplifierTerm Name DefaultUni DefaultFun ()
+  -> SimplifierTerm Name DefaultUni DefaultFun ()
   -> SimplifierTrace Name DefaultUni DefaultFun ()
 mkMockTracePair stage before' after' =
   SimplifierTrace
@@ -43,7 +44,7 @@ testSuccess
   -> Term Name PLC.DefaultUni PLC.DefaultFun ()
   -> Term Name PLC.DefaultUni PLC.DefaultFun ()
   -> TestTree
-testSuccess testName st bf af =
+testSuccess testName st (initSimplifierTerm -> bf) (initSimplifierTerm -> af) =
   testCase testName $ do
     let trace = mkMockTracePair st bf af
     result <- runCertifierWithMockTrace trace
@@ -57,7 +58,7 @@ testFailure
   -> Term Name PLC.DefaultUni PLC.DefaultFun ()
   -> Term Name PLC.DefaultUni PLC.DefaultFun ()
   -> TestTree
-testFailure testName st bf af =
+testFailure testName st (initSimplifierTerm -> bf) (initSimplifierTerm -> af) =
   testCase testName $ do
     let trace = mkMockTracePair st bf af
     result <- runCertifierWithMockTrace trace