Make forcing builtin functions pure and work-free

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

@@ -4,6 +4,7 @@ module UntypedPlutusCore.Core (
   module Export,
   splitParams,
   splitApplication,
+  splitForces,
 ) where
 
 import UntypedPlutusCore.Core.Instance as Export
@@ -25,3 +26,11 @@ splitApplication = go []
     go acc = \case
       Apply ann fun arg -> go ((ann, arg) : acc) fun
       t                 -> (t, acc)
+
+-- | Extract the term inside a number of nested `Force` nodes.
+--
+-- This also returns the annotations on the `Force` noes, in the order of outer to inner.
+splitForces :: Term name uni fun a -> (Term name uni fun a, [a])
+splitForces = \case
+  Force ann body -> second (ann:) (splitForces body)
+  t              -> (t, [])

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

@@ -17,7 +17,7 @@ module UntypedPlutusCore.Purity
     ) where
 
 import PlutusCore.Pretty
-import UntypedPlutusCore.Core.Type
+import UntypedPlutusCore.Core
 
 import Data.DList qualified as DList
 import Prettyprinter
@@ -108,15 +108,23 @@ termEvaluationOrder = goTerm
         t@(Force _ dterm)        ->
           -- first delayed term
           goTerm dterm
-          -- then the whole term, which will mean forcing, so work
-          <> evalThis (EvalTerm Pure MaybeWork t)
+          <> evalThis (EvalTerm Pure workFreedom t)
           <> dest
           where
+            workFreedom = case dtermInner of
+              -- Forcing a builtin is workfree
+              Builtin{} -> WorkFree
+              -- Forcing a delayed term may not be workfree
+              _         -> MaybeWork
             dest = case dterm of
               -- known delayed term
               (Delay _ body) -> goTerm body
-              -- unknown delayed term
-              _              -> evalThis Unknown
+              _ -> case dtermInner of
+                -- known builtin. It is already accounted for in `goTerm dterm`, hence `mempty`.
+                Builtin{} -> mempty
+                -- unknown delayed term
+                _         -> evalThis Unknown
+            dtermInner = fst (splitForces dterm)
 
         t@(Constr _ _ ts)     ->
           -- first the arguments, in left-to-right order

plutus-core/untyped-plutus-core/test/Transform/Simplify.hs

@@ -192,6 +192,24 @@ inlineImpure4 =
   mkInlinePurityTest $
     Force () . Force () . Force () . Delay () . Delay () . Var () <$> freshName "a"
 
+{- | @(\x -> error x) (force ifThenElse)@.
+@x@ should be inlined since @force ifThenElse@ is pure.
+-}
+inline1 :: Term Name PLC.DefaultUni PLC.DefaultFun ()
+inline1 = runQuote $ do
+  x <- freshName "x"
+  let xRhs = Force () (Builtin () PLC.IfThenElse)
+  pure $ Apply () (LamAbs () x (Apply () (Error ()) (Var () x))) xRhs
+
+{- | @(\x -> error x) (force (force fstPair))@.
+@x@ should be inlined since @force (force fstPair)@ is pure.
+-}
+inline2 :: Term Name PLC.DefaultUni PLC.DefaultFun ()
+inline2 = runQuote $ do
+  x <- freshName "x"
+  let xRhs = Force () (Force () (Builtin () PLC.FstPair))
+  pure $ Apply () (LamAbs () x (Apply () (Error ()) (Var () x))) xRhs
+
 {- | @(\a -> f (a 0 1) (a 2)) (\x y -> g x y)@
 
 The first occurrence of `a` should be inlined because doing so does not increase
@@ -264,5 +282,7 @@ test_simplify =
     , goldenVsSimplified "inlineImpure2" inlineImpure2
     , goldenVsSimplified "inlineImpure3" inlineImpure3
     , goldenVsSimplified "inlineImpure4" inlineImpure4
+    , goldenVsSimplified "inline1" inline1
+    , goldenVsSimplified "inline2" inline2
     , goldenVsSimplified "multiApp" multiApp
     ]

plutus-core/untyped-plutus-core/test/Transform/inline1.uplc.golden

@@ -0,0 +1 @@
+(error (force ifThenElse))

plutus-core/untyped-plutus-core/test/Transform/inline2.uplc.golden

@@ -0,0 +1 @@
+(error (force (force fstPair)))