Fix firstEffectfulTerm (#5361)

plutus-benchmark/script-contexts/test/Spec.hs

@@ -20,7 +20,7 @@ testCheckSc1 = testGroup "checkScriptContext1"
         runTermCek $ compiledCodeToTerm $ mkCheckScriptContext1Code (mkScriptContext 4)
     , testCase "fails on 5" $ assertBool "evaluation succeeded" $ isEvaluationFailure $
         runTermCek $ compiledCodeToTerm $ mkCheckScriptContext1Code (mkScriptContext 5)
-    , Tx.fitsInto "checkScriptContext1 (size)" (mkCheckScriptContext1Code (mkScriptContext 1)) 2054
+    , Tx.fitsInto "checkScriptContext1 (size)" (mkCheckScriptContext1Code (mkScriptContext 1)) 1985
     , runTestNested $ Tx.goldenBudget "checkScriptContext1-4" $
         mkCheckScriptContext1Code (mkScriptContext 4)
     , runTestNested $ Tx.goldenBudget "checkScriptContext1-20" $
@@ -33,7 +33,7 @@ testCheckSc2 = testGroup "checkScriptContext2"
           runTermCek $ compiledCodeToTerm $ mkCheckScriptContext2Code (mkScriptContext 4)
     , testCase "succeed on 5" $ assertBool "evaluation failed" $ isEvaluationSuccess $
           runTermCek $ compiledCodeToTerm $ mkCheckScriptContext2Code (mkScriptContext 5)
-    , Tx.fitsInto "checkScriptContext2 (size)" (mkCheckScriptContext2Code (mkScriptContext 1)) 1985
+    , Tx.fitsInto "checkScriptContext2 (size)" (mkCheckScriptContext2Code (mkScriptContext 1)) 1919
     , runTestNested $ Tx.goldenBudget "checkScriptContext2-4" $
           mkCheckScriptContext2Code (mkScriptContext 4)
     , runTestNested $ Tx.goldenBudget "checkScriptContext2-20" $

plutus-benchmark/script-contexts/test/checkScriptContext1-20.budget.golden

@@ -1,2 +1,2 @@
-({cpu: 438235997
-| mem: 1377569})
+({cpu: 432991997
+| mem: 1354769})

plutus-benchmark/script-contexts/test/checkScriptContext1-4.budget.golden

@@ -1,2 +1,2 @@
-({cpu: 126991245
-| mem: 401793})
+({cpu: 125059245
+| mem: 393393})

plutus-benchmark/script-contexts/test/checkScriptContext2-20.budget.golden

@@ -1,2 +1,2 @@
-({cpu: 418965508
-| mem: 1314926})
+({cpu: 413790508
+| mem: 1292426})

plutus-benchmark/script-contexts/test/checkScriptContext2-4.budget.golden

@@ -1,2 +1,2 @@
-({cpu: 121328388
-| mem: 383982})
+({cpu: 119465388
+| mem: 375882})

plutus-core/changelog.d/20230525_144849_unsafeFixIO_firstEffectfulTerm.md

@@ -0,0 +1,4 @@
+### Fixed
+
+- Fixed `PlutusIR.Purity.firstEffectfulTerm` and `UntypedPlutusCore.Transform.Inline.firstEffectfulTerm`,
+  which were sometimes too conservative and sometimes incorrect.

plutus-core/plutus-ir/src/PlutusIR/Purity.hs

@@ -5,17 +5,19 @@
 
 module PlutusIR.Purity
     ( isPure
+    , FirstEffectfulTerm (..)
     , firstEffectfulTerm
     , asBuiltinApp
     , isSaturated
     , BuiltinApp (..)
     , Arg (..)
     ) where
 
+import PlutusCore.Builtin
 import PlutusIR
 
+import Control.Applicative
 import Data.List.NonEmpty qualified as NE
-import PlutusCore.Builtin
 
 -- | An argument taken by a builtin: could be a term of a type.
 data Arg tyname name uni fun a = TypeArg (Type tyname uni a) | TermArg (Term tyname name uni fun a)
@@ -127,39 +129,49 @@ isPure ver varStrictness = go
             TermBind _ Strict _ rhs -> go rhs
             _                       -> True
 
+-- | Isomorphic to @Maybe (Term tyname name uni fun a)@. Used to represent the first
+-- subterm which will be evaluated in a term and which could have an effect.
+data FirstEffectfulTerm tyname name uni fun a
+  = EffectfulTerm (Term tyname name uni fun a)
+  -- | `Uncertain` indicates that we don't know for sure.
+  | Uncertain
+
 {- |
 Try to identify the first sub term which will be evaluated in the given term and
-which could have an effect. 'Nothing' indicates that we don't know, this function
-is conservative.
+which could have an effect. 'Nothing' indicates that there's no term to evaluate.
 -}
-firstEffectfulTerm :: Term tyname name uni fun a -> Maybe (Term tyname name uni fun a)
+firstEffectfulTerm ::
+  forall tyname name uni fun a.
+  Term tyname name uni fun a ->
+  Maybe (FirstEffectfulTerm tyname name uni fun a)
 firstEffectfulTerm = goTerm
     where
+      goTerm :: Term tyname name uni fun a -> Maybe (FirstEffectfulTerm tyname name uni fun a)
       goTerm = \case
-        Let _ NonRec bs b -> case goBindings (NE.toList bs) of
-            Just t' -> Just t'
-            Nothing -> goTerm b
-
-        Apply _ l _ -> goTerm l
-        TyInst _ t _ -> goTerm t
-        IWrap _ _ _ t -> goTerm t
-        Unwrap _ t -> goTerm t
-        Constr _ _ _ [] -> Nothing
-        Constr _ _ _ (t:_) -> goTerm t
-        Case _ _ t _ -> goTerm t
-
-        t@Var{} -> Just t
-        t@Error{} -> Just t
-        t@Builtin{} -> Just t
+        Let _ NonRec bs b -> goBindings (NE.toList bs) <|> goTerm b
 
-        -- Hard to know what gets evaluated first in a recursive let-binding,
-        -- just give up and say nothing
-        (Let _ Rec _ _) -> Nothing
-        TyAbs{} -> Nothing
-        LamAbs{} -> Nothing
-        Constant{} -> Nothing
+        Apply _ fun args  -> goTerm fun <|> goTerm args
+        TyInst _ t _      -> goTerm t
+        IWrap _ _ _ t     -> goTerm t
+        Unwrap _ t        -> goTerm t
+        Constr _ _ _ []   -> Nothing
+        Constr _ _ _ ts   -> asum $ goTerm <$> ts
+        Case _ _ t _      -> goTerm t
+
+        t@Var{}           -> Just (EffectfulTerm t)
+        t@Error{}         -> Just (EffectfulTerm t)
+        Builtin{}         -> Nothing
 
-      goBindings :: [Binding tyname name uni fun a] -> Maybe (Term tyname name uni fun a)
+        -- Hard to know what gets evaluated first in a recursive let-binding,
+        -- just give up and return `Uncertain`.
+        (Let _ Rec _ _)   -> Just Uncertain
+        TyAbs{}           -> Nothing
+        LamAbs{}          -> Nothing
+        Constant{}        -> Nothing
+
+      goBindings ::
+        [Binding tyname name uni fun a] ->
+        Maybe (FirstEffectfulTerm tyname name uni fun a)
       goBindings [] = Nothing
       goBindings (b:bs) = case b of
         -- Only strict term bindings can cause effects

plutus-core/plutus-ir/src/PlutusIR/Transform/Inline/Utils.hs

@@ -18,7 +18,7 @@ import PlutusCore.Subst (typeSubstTyNamesM)
 import PlutusIR
 import PlutusIR.Analysis.Dependencies qualified as Deps
 import PlutusIR.Analysis.Usages qualified as Usages
-import PlutusIR.Purity (firstEffectfulTerm, isPure)
+import PlutusIR.Purity (FirstEffectfulTerm (..), firstEffectfulTerm, isPure)
 import PlutusIR.Transform.Rename ()
 import PlutusPrelude
 
@@ -301,8 +301,8 @@ effectSafe body s n purity = do
     -- doing ~quadratic work as we process the program. However in practice most term
     -- types will make it give up, so it's not too bad.
     let immediatelyEvaluated = case firstEffectfulTerm body of
-            Just (Var _ n') -> n == n'
-            _               -> False
+            Just (EffectfulTerm (Var _ n')) -> n == n'
+            _                               -> False
     pure $ case s of
         Strict    -> purity || immediatelyEvaluated
         NonStrict -> True

plutus-core/plutus-ir/test/TransformSpec.hs

@@ -238,6 +238,8 @@ inline =
             , "single"
             , "immediateVar"
             , "immediateApp"
+            , "firstEffectfulTerm1"
+            , "firstEffectfulTerm2"
             -- these tests are all let bindings of functions
             , "letFunConstInt" -- const fn fully applied (integer)
             , "letFunConstBool" -- const fn fully applied (bool)

plutus-core/plutus-ir/test/transform/inline/firstEffectfulTerm1

@@ -0,0 +1,10 @@
+-- `a` should be inlined since it is the first effectful term in the body.
+(let
+  (nonrec)
+  (termbind
+    (strict)
+    (vardecl a (con integer))
+    { error (con integer) }
+  )
+  [ (lam x (con integer) x) a ]
+)

plutus-core/plutus-ir/test/transform/inline/firstEffectfulTerm1.golden

@@ -0,0 +1 @@
+[ (lam x (con integer) x) { error (con integer) } ]

plutus-core/plutus-ir/test/transform/inline/firstEffectfulTerm2

@@ -0,0 +1,18 @@
+-- `a` should not be inlined since it changes the semantics of the program.
+(let
+  (nonrec)
+  (termbind
+    (strict)
+    (vardecl a (fun (con integer) (con integer)))
+    { error (fun (con integer) (con integer)) }
+  )
+  (let
+    (nonrec)
+    (termbind
+      (strict)
+      (vardecl b (con integer))
+      [ (lam x (con integer) x) { error (con integer) } ]
+    )
+    [ a b ]
+  )
+)

plutus-core/plutus-ir/test/transform/inline/firstEffectfulTerm2.golden

@@ -0,0 +1,17 @@
+(let
+  (nonrec)
+  (termbind
+    (strict)
+    (vardecl a (fun (con integer) (con integer)))
+    { error (fun (con integer) (con integer)) }
+  )
+  (let
+    (nonrec)
+    (termbind
+      (strict)
+      (vardecl b (con integer))
+      [ (lam x (con integer) x) { error (con integer) } ]
+    )
+    [ a b ]
+  )
+)

plutus-core/plutus-ir/test/transform/inline/letNonPureMulti.golden

@@ -7,7 +7,6 @@
       [ { (builtin trace) (con integer) } (con string "hello") ] (con integer 1)
     ]
   )
-  (termbind (strict) (vardecl z (con integer)) y)
   (termbind (strict) (vardecl x (con integer)) y)
-  [ [ (builtin addInteger) z ] [ [ (builtin addInteger) x ] x ] ]
+  [ [ (builtin addInteger) y ] [ [ (builtin addInteger) x ] x ] ]
 )

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

@@ -36,6 +36,7 @@ import PlutusCore.Builtin qualified as PLC
 import PlutusCore.Name
 import PlutusCore.Quote
 
+import Control.Applicative
 import Control.Lens hiding (Strict)
 import Control.Monad (liftM2)
 import Control.Monad.Reader (ReaderT, asks, runReaderT)
@@ -270,23 +271,22 @@ acceptable t =
 
 {- |
 Try to identify the first sub term which will be evaluated in the given term and
-which could have an effect. 'Nothing' indicates that we don't know, this function
-is conservative.
+which could have an effect. 'Nothing' indicates that there's no term to evaluate.
 -}
 firstEffectfulTerm :: Term name uni fun a -> Maybe (Term name uni fun a)
 firstEffectfulTerm = goTerm
     where
       goTerm = \case
 
-        Apply _ l _      -> goTerm l
+        Apply _ fun args -> goTerm fun <|> goTerm args
         Force _ t        -> goTerm t
         Constr _ _ []    -> Nothing
-        Constr _ _ (t:_) -> goTerm t
+        Constr _ _ ts    -> asum $ goTerm <$> ts
         Case _ t _       -> goTerm t
 
         t@Var{}          -> Just t
         t@Error{}        -> Just t
-        t@Builtin{}      -> Just t
+        Builtin{}        -> Nothing
 
         LamAbs{}         -> Nothing
         Constant{}       -> Nothing

plutus-tx-plugin/test/Budget/allCheap.budget.golden

@@ -1,2 +1,2 @@
-({cpu: 3115053
-| mem: 12402})
+({cpu: 3184053
+| mem: 12702})

plutus-tx-plugin/test/Budget/allCheap.uplc-readable.golden

@@ -1,5 +1,27 @@
-program 1.1.0 ((\eta ->
-                  (\s ->
+program 1.1.0 ((\go ->
+                  (\eta -> go eta)
+                    ((\c ->
+                        c
+                          1
+                          (c
+                             2
+                             (c
+                                3
+                                (c
+                                   4
+                                   (c
+                                      5
+                                      (c
+                                         6
+                                         (c
+                                            7
+                                            (c
+                                               8
+                                               (c
+                                                  9
+                                                  (c 10 (constr 0 [])))))))))))
+                       (\ds ds -> constr 1 [ds, ds])))
+                 ((\s ->
                      s s)
                     (\s
                       x ->
@@ -15,22 +37,4 @@ program 1.1.0 ((\eta ->
                                                                    (constr 0 [  ])) [ (delay (go
                                                                                                 xs))
                                                                                     , (delay (constr 1 [  ])) ]))) ]))
-                         (s s))
-                    eta)
-                 ((\c ->
-                     c
-                       1
-                       (c
-                          2
-                          (c
-                             3
-                             (c
-                                4
-                                (c
-                                   5
-                                   (c
-                                      6
-                                      (c
-                                         7
-                                         (c 8 (c 9 (c 10 (constr 0 [])))))))))))
-                    (\ds ds -> constr 1 [ds, ds])))
+                         (s s))))

plutus-tx-plugin/test/Budget/allExpensive.budget.golden

@@ -1,2 +1,2 @@
-({cpu: 12197630
-| mem: 40720})
+({cpu: 12266630
+| mem: 41020})

plutus-tx-plugin/test/Budget/allExpensive.uplc-readable.golden

@@ -1,5 +1,27 @@
-program 1.1.0 ((\eta ->
-                  (\s ->
+program 1.1.0 ((\go ->
+                  (\eta -> go eta)
+                    ((\c ->
+                        c
+                          1
+                          (c
+                             2
+                             (c
+                                3
+                                (c
+                                   4
+                                   (c
+                                      5
+                                      (c
+                                         6
+                                         (c
+                                            7
+                                            (c
+                                               8
+                                               (c
+                                                  9
+                                                  (c 10 (constr 0 [])))))))))))
+                       (\ds ds -> constr 1 [ds, ds])))
+                 ((\s ->
                      s s)
                     (\s
                       x ->
@@ -15,22 +37,4 @@ program 1.1.0 ((\eta ->
                                                                    (constr 0 [  ])) [ (delay (go
                                                                                                 xs))
                                                                                     , (delay (constr 1 [  ])) ]))) ]))
-                         (s s))
-                    eta)
-                 ((\c ->
-                     c
-                       1
-                       (c
-                          2
-                          (c
-                             3
-                             (c
-                                4
-                                (c
-                                   5
-                                   (c
-                                      6
-                                      (c
-                                         7
-                                         (c 8 (c 9 (c 10 (constr 0 [])))))))))))
-                    (\ds ds -> constr 1 [ds, ds])))
+                         (s s))))