Add a constant folding pass to UPLC

plutus-core/changelog.d/20260504_235314_ana.pantilie95_constant_folding_uplc_2.md

@@ -0,0 +1,3 @@
+### Added
+
+- The UPLC optimizer now also simplifies terms by evaluating builtins.

plutus-core/plutus-core.cabal

@@ -235,6 +235,7 @@ library
     UntypedPlutusCore.Transform.Certify.Hints
     UntypedPlutusCore.Transform.Certify.Trace
     UntypedPlutusCore.Transform.Cse
+    UntypedPlutusCore.Transform.EvaluateBuiltins
     UntypedPlutusCore.Transform.FloatDelay
     UntypedPlutusCore.Transform.ForceCaseDelay
     UntypedPlutusCore.Transform.ForceDelay
@@ -282,6 +283,7 @@ library
     PlutusCore.Pretty.Readable
     PlutusCore.Pretty.Utils
     Universe.Core
+    UntypedPlutusCore.Analysis.Builtins
     UntypedPlutusCore.Analysis.Definitions
     UntypedPlutusCore.Analysis.Usages
     UntypedPlutusCore.Core.Instance
@@ -485,6 +487,7 @@ library untyped-plutus-core-testlib
     Generators.Spec
     Scoping.Spec
     Transform.CaseOfCase.Spec
+    Transform.EvaluateBuiltins.Spec
     Transform.Inline.Spec
     Transform.Simplify.Lib
     Transform.Simplify.Spec

plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Analysis/Builtins.hs

@@ -0,0 +1,78 @@
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeApplications #-}
+
+module UntypedPlutusCore.Analysis.Builtins
+  ( BuiltinsInfo (..)
+  , biSemanticsVariant
+  , biUnserializableConstants
+  , builtinArityInfo
+  , constantIsSerializable
+  , termIsSerializable
+  , defaultUniUnserializableConstants
+  ) where
+
+import Control.Lens hiding (parts)
+import Data.Kind
+import Data.Proxy
+import PlutusCore.Arity
+import PlutusCore.Builtin
+import PlutusCore.Builtin qualified as PLC
+import PlutusCore.Default
+import PlutusPrelude (Default (..))
+import UntypedPlutusCore.Core (Term)
+import UntypedPlutusCore.Core.Plated (termSubtermsDeep, _Constant)
+
+-- | All non-static information about builtins that the compiler might want.
+data BuiltinsInfo (uni :: Type -> Type) fun = BuiltinsInfo
+  { _biSemanticsVariant :: PLC.BuiltinSemanticsVariant fun
+  , -- See Note [Unserializable constants]
+    _biUnserializableConstants :: Some (ValueOf uni) -> Bool
+  }
+
+makeLenses ''BuiltinsInfo
+
+instance Default (BuiltinsInfo DefaultUni DefaultFun) where
+  def =
+    BuiltinsInfo
+      { _biSemanticsVariant = def
+      , _biUnserializableConstants = defaultUniUnserializableConstants
+      }
+
+-- | Get the arity of a builtin function from the 'PLC.BuiltinInfo'.
+builtinArityInfo
+  :: forall uni fun
+   . ToBuiltinMeaning uni fun
+  => BuiltinsInfo uni fun
+  -> fun
+  -> Arity
+builtinArityInfo binfo = builtinArity (Proxy @uni) (binfo ^. biSemanticsVariant)
+
+constantIsSerializable
+  :: forall uni fun
+   . BuiltinsInfo uni fun
+  -> Some (ValueOf uni)
+  -> Bool
+constantIsSerializable bi v = not $ _biUnserializableConstants bi v
+
+termIsSerializable :: BuiltinsInfo uni fun -> Term name uni fun a -> Bool
+termIsSerializable binfo =
+  allOf
+    (termSubtermsDeep . _Constant)
+    (constantIsSerializable binfo . snd)
+
+-- See Note [Unserializable constants]
+defaultUniUnserializableConstants :: Some (ValueOf DefaultUni) -> Bool
+defaultUniUnserializableConstants = \case
+  Some (ValueOf DefaultUniBLS12_381_G1_Element _) -> True
+  Some (ValueOf DefaultUniBLS12_381_G2_Element _) -> True
+  Some (ValueOf DefaultUniBLS12_381_MlResult _) -> True
+  _ -> False
+
+{- See Note [Unserializable constants] in PlutusIR.Analysis.Builtins.
+-}

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

@@ -1,7 +1,9 @@
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE RankNTypes #-}
 
 module UntypedPlutusCore.Contexts where
 
+import PlutusCore.Arity
 import UntypedPlutusCore.Core (Term (..))
 import UntypedPlutusCore.Core.Instance.Eq ()
 
@@ -23,7 +25,7 @@ splitAppCtx :: Term nam uni fun a -> (Term nam uni fun a, AppCtx nam uni fun a)
 splitAppCtx = go AppCtxEnd
   where
     go appCtx = \case
-      Apply ann function argument -> go (AppCtxTerm ann argument appCtx) function
+      Apply ann function arg -> go (AppCtxTerm ann arg appCtx) function
       Force ann forcedTerm -> go (AppCtxType ann appCtx) forcedTerm
       term -> (term, appCtx)
 
@@ -36,3 +38,22 @@ fillAppCtx term = \case
   AppCtxEnd -> term
   AppCtxTerm ann arg ctx -> fillAppCtx (Apply ann term arg) ctx
   AppCtxType ann ctx -> fillAppCtx (Force ann term) ctx
+
+data Saturation = Oversaturated | Undersaturated | Saturated
+
+-- | Do the given arguments saturate the given arity?
+saturates :: AppCtx name uni fun a -> Arity -> Maybe Saturation
+-- Exactly right
+saturates AppCtxEnd [] = Just Saturated
+-- Parameters left - undersaturated
+saturates AppCtxEnd _ = Just Undersaturated
+-- Match a term parameter to a term arg
+saturates (AppCtxTerm _ _ ctx) (TermParam : arities) = saturates ctx arities
+-- Match a type parameter to a type arg
+saturates (AppCtxType _ ctx) (TypeParam : arities) = saturates ctx arities
+-- Param/arg mismatch
+saturates (AppCtxTerm {}) (TypeParam : _) = Nothing
+saturates (AppCtxType {}) (TermParam : _) = Nothing
+-- Arguments left - undersaturated
+saturates (AppCtxTerm {}) [] = Just Oversaturated
+saturates (AppCtxType {}) [] = Just Oversaturated

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

@@ -10,8 +10,10 @@ module UntypedPlutusCore.Core.Plated
   , termConstantsDeep
   , termSubtermsDeep
   , termUniquesDeep
+  , _Constant
   ) where
 
+import PlutusCore qualified as PLC
 import PlutusCore.Core (HasUniques)
 import PlutusCore.Name.Unique
 import UntypedPlutusCore.Core.Type
@@ -78,3 +80,7 @@ termSubtermsDeep = cosmosOf termSubterms
 -- | Get all the transitive child 'Unique's of the given 'Term'.
 termUniquesDeep :: HasUniques (Term name uni fun ann) => Fold (Term name uni fun ann) Unique
 termUniquesDeep = termSubtermsDeep . termUniques
+
+-- | View a term as a constant.
+_Constant :: Prism' (Term name uni fun a) (a, PLC.Some (PLC.ValueOf uni))
+_Constant = prism' (uncurry Constant) (\case Constant a v -> Just (a, v); _ -> Nothing)

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

@@ -32,6 +32,7 @@ import PlutusPrelude
 import Data.Vector
 import PlutusCore.Builtin qualified as TPLC
 import PlutusCore.Core qualified as TPLC
+import PlutusCore.Evaluation.Machine.ExMemoryUsage
 import PlutusCore.MkPlc
 import PlutusCore.Name.Unique qualified as TPLC
 import Universe
@@ -110,6 +111,11 @@ deriving anyclass instance
   (NFData name, NFData fun, NFData ann, Everywhere uni NFData, Closed uni)
   => NFData (Term name uni fun ann)
 
+-- See Note [ExMemoryUsage instances for non-constants].
+instance ExMemoryUsage (Term name uni fun ann) where
+  memoryUsage =
+    Prelude.error "Internal error: 'memoryUsage' for UPLC 'Term' is not supposed to be forced"
+
 -- | A 'Program' is simply a 'Term' coupled with a 'Version' of the core language.
 data Program name uni fun ann = Program
   { _progAnn :: ann
@@ -165,25 +171,25 @@ makeLenses ''UVarDecl
 -- | Return the outermost annotation of a 'Term'.
 instance HasAnn (Term name uni fun) where
   getAnn (Constant ann _) = ann
-  getAnn (Builtin ann _)  = ann
-  getAnn (Var ann _)      = ann
+  getAnn (Builtin ann _) = ann
+  getAnn (Var ann _) = ann
   getAnn (LamAbs ann _ _) = ann
-  getAnn (Apply ann _ _)  = ann
-  getAnn (Delay ann _)    = ann
-  getAnn (Force ann _)    = ann
-  getAnn (Error ann)      = ann
+  getAnn (Apply ann _ _) = ann
+  getAnn (Delay ann _) = ann
+  getAnn (Force ann _) = ann
+  getAnn (Error ann) = ann
   getAnn (Constr ann _ _) = ann
-  getAnn (Case ann _ _)   = ann
+  getAnn (Case ann _ _) = ann
   modifyAnn f = \case
-    Constant ann c      -> Constant (f ann) c
-    Builtin ann b       -> Builtin (f ann) b
-    Var ann v           -> Var (f ann) v
-    LamAbs ann x body   -> LamAbs (f ann) x body
-    Apply ann fun arg   -> Apply (f ann) fun arg
-    Delay ann body      -> Delay (f ann) body
-    Force ann body      -> Force (f ann) body
-    Error ann           -> Error (f ann)
-    Constr ann i args   -> Constr (f ann) i args
+    Constant ann c -> Constant (f ann) c
+    Builtin ann b -> Builtin (f ann) b
+    Var ann v -> Var (f ann) v
+    LamAbs ann x body -> LamAbs (f ann) x body
+    Apply ann fun arg -> Apply (f ann) fun arg
+    Delay ann body -> Delay (f ann) body
+    Force ann body -> Force (f ann) body
+    Error ann -> Error (f ann)
+    Constr ann i args -> Constr (f ann) i args
     Case ann scrut alts -> Case (f ann) scrut alts
 
 bindFunM

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

@@ -13,24 +13,29 @@ module UntypedPlutusCore.Optimize
   , module UntypedPlutusCore.Transform.Optimizer
   ) where
 
+import PlutusCore.Builtin (CostingPart)
 import PlutusCore.Compiler.Types
 import PlutusCore.Default qualified as PLC
 import PlutusCore.Default.Builtins
 import PlutusCore.Name.Unique
+import UntypedPlutusCore.Analysis.Builtins (BuiltinsInfo, biSemanticsVariant)
 import UntypedPlutusCore.Core.Type
 import UntypedPlutusCore.Optimize.Opts as Opts
 import UntypedPlutusCore.Transform.ApplyToCase (applyToCase)
 import UntypedPlutusCore.Transform.CaseOfCase
 import UntypedPlutusCore.Transform.CaseReduce
 import UntypedPlutusCore.Transform.Cse
+import UntypedPlutusCore.Transform.EvaluateBuiltins (evaluateBuiltinsPass)
 import UntypedPlutusCore.Transform.FloatDelay (floatDelay)
 import UntypedPlutusCore.Transform.ForceCaseDelay (forceCaseDelay)
 import UntypedPlutusCore.Transform.ForceDelay (forceDelay)
 import UntypedPlutusCore.Transform.Inline (InlineHints (..), inline)
 import UntypedPlutusCore.Transform.LetFloatOut (letFloatOut)
 import UntypedPlutusCore.Transform.Optimizer
 
+import Control.Lens ((&), (.~))
 import Control.Monad
+import Data.Default.Class (def)
 import Data.Either (isRight)
 import Data.List as List (foldl')
 import Data.Typeable
@@ -109,6 +114,7 @@ termOptimizer opts builtinSemanticsVariant =
         >=> runStage CaseOfCaseStage
         >=> runStage CaseReduceStage
         >=> runStage InlineStage
+        >=> runStage ConstantFoldingStage
 
     certifiedOnly = _ooCertifiedOptsOnly opts
 
@@ -144,6 +150,14 @@ termOptimizer opts builtinSemanticsVariant =
             if _ooApplyToCase opts then applyToCase else pure
           LetFloatOutStage ->
             letFloatOut
+          ConstantFoldingStage ->
+            case (eqT @uni @PLC.DefaultUni, eqT @fun @DefaultFun) of
+              (Just Refl, Just Refl) ->
+                evaluateBuiltinsPass
+                  (_ooPreserveLogging opts)
+                  ((def :: BuiltinsInfo PLC.DefaultUni DefaultFun) & biSemanticsVariant .~ builtinSemanticsVariant)
+                  (def :: CostingPart PLC.DefaultUni DefaultFun)
+              _ -> pure
 
     caseOfCase'
       :: Term name uni fun a

plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/Certify/Trace.hs

@@ -43,6 +43,7 @@ corresponding constructor to this type. Please also open an issue
 at https://github.com/IntersectMBO/plutus/issues. -}
 data UncertifiedOptStage
   = CaseOfCase
+  | ConstantFolding
   deriving stock (Show, Generic)
   deriving anyclass (NFData)
 
@@ -75,6 +76,9 @@ pattern CaseOfCaseStage = Left CaseOfCase
 pattern LetFloatOutStage :: OptStage
 pattern LetFloatOutStage = Right LetFloatOut
 
+pattern ConstantFoldingStage :: OptStage
+pattern ConstantFoldingStage = Left ConstantFolding
+
 {-# COMPLETE
   FloatDelayStage
   , ForceDelayStage
@@ -85,6 +89,7 @@ pattern LetFloatOutStage = Right LetFloatOut
   , ApplyToCaseStage
   , CaseOfCaseStage
   , LetFloatOutStage
+  , ConstantFoldingStage
   #-}
 
 data Optimization name uni fun a

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

@@ -0,0 +1,71 @@
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE ViewPatterns #-}
+
+{-| This module mirrors 'PlutusIR.Transform.EvaluateBuiltins', adapted for UPLC's
+untyped 'Term' and 'AppCtx'. See that module for detailed commentary. -}
+module UntypedPlutusCore.Transform.EvaluateBuiltins
+  ( evaluateBuiltinsPass
+  ) where
+
+import PlutusCore.Builtin
+import UntypedPlutusCore.Analysis.Builtins
+import UntypedPlutusCore.Contexts
+import UntypedPlutusCore.Core
+import UntypedPlutusCore.Transform.Certify.Trace
+import UntypedPlutusCore.Transform.Optimizer
+
+import Control.Lens (transformOf, (^.))
+import Data.Functor (void)
+
+evaluateBuiltinsPass
+  :: (Monad m, ToBuiltinMeaning uni fun, Typeable name)
+  => Bool
+  -- ^ Whether to be conservative and try to retain logging behaviour.
+  -> BuiltinsInfo uni fun
+  -> CostingPart uni fun
+  -> Term name uni fun a
+  -> OptimizerT name uni fun a m (Term name uni fun a)
+evaluateBuiltinsPass preserveLogging binfo costModel term = do
+  result <- evaluateBuiltins preserveLogging binfo costModel term
+  recordOptimization term ConstantFoldingStage result
+  return result
+
+evaluateBuiltins
+  :: forall m name uni fun a
+   . ( Monad m
+     , ToBuiltinMeaning uni fun
+     , Typeable name
+     )
+  => Bool
+  -- ^ Whether to be conservative and try to retain logging behaviour.
+  -> BuiltinsInfo uni fun
+  -> CostingPart uni fun
+  -> Term name uni fun a
+  -> OptimizerT name uni fun a m (Term name uni fun a)
+evaluateBuiltins preserveLogging binfo costModel =
+  pure . transformOf termSubterms processTerm
+  where
+    eval
+      :: BuiltinRuntime (Term name uni fun ())
+      -> AppCtx name uni fun a
+      -> Maybe (Term name uni fun ())
+    eval (BuiltinCostedResult _ getFXs) AppCtxEnd =
+      case getFXs of
+        BuiltinSuccess y -> Just y
+        BuiltinSuccessWithLogs _ y -> if preserveLogging then Nothing else Just y
+        BuiltinFailure {} -> Nothing
+    eval (BuiltinExpectArgument toRuntime) (AppCtxTerm _ arg ctx) =
+      eval (toRuntime $ void arg) ctx
+    eval (BuiltinExpectForce runtime) (AppCtxType _ ctx) =
+      eval runtime ctx
+    eval _ _ = Nothing
+
+    processTerm :: Term name uni fun a -> Term name uni fun a
+    -- See Note [Context splitting in a recursive pass]
+    processTerm t@(splitAppCtx -> (Builtin x bn, argCtx)) =
+      let runtime = toBuiltinRuntime costModel (toBuiltinMeaning (binfo ^. biSemanticsVariant) bn)
+       in case eval runtime argCtx of
+            -- See Note [Unserializable constants] in PlutusIR.Analysis.Builtins
+            Just t' | termIsSerializable binfo t' -> x <$ t'
+            _ -> t
+    processTerm t = t

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

@@ -12,6 +12,7 @@ module UntypedPlutusCore.Transform.Optimizer
   , pattern Trace.ApplyToCaseStage
   , pattern Trace.CaseOfCaseStage
   , pattern Trace.LetFloatOutStage
+  , pattern Trace.ConstantFoldingStage
   , Trace.OptimizerTrace (..)
   , Trace.Optimization (..)
   , runOptimizerT