CaseOfCase UPLC transformation tests. (#5960)

plutus-benchmark/nofib/test/9.6/knights10-4x4.budget.golden

@@ -1,2 +1,2 @@
-({cpu: 2339861234
-| mem: 7524812})
+({cpu: 2340321234
+| mem: 7526812})

plutus-benchmark/nofib/test/9.6/knights10-4x4.size.golden

@@ -1 +1 @@
-1995
+1998

plutus-core/plutus-core.cabal

@@ -202,6 +202,7 @@ library
     UntypedPlutusCore.Parser
     UntypedPlutusCore.Purity
     UntypedPlutusCore.Rename
+    UntypedPlutusCore.Transform.CaseOfCase
 
   other-modules:
     Data.Aeson.Flatten
@@ -261,7 +262,6 @@ library
     UntypedPlutusCore.Simplify.Opts
     UntypedPlutusCore.Size
     UntypedPlutusCore.Subst
-    UntypedPlutusCore.Transform.CaseOfCase
     UntypedPlutusCore.Transform.CaseReduce
     UntypedPlutusCore.Transform.Cse
     UntypedPlutusCore.Transform.FloatDelay
@@ -427,7 +427,9 @@ test-suite untyped-plutus-core-test
     Evaluation.Regressions
     Flat.Spec
     Generators
+    Transform.CaseOfCase.Test
     Transform.Simplify
+    Transform.Simplify.Lib
 
   build-depends:
     , base                                            >=4.9   && <5

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

@@ -27,11 +27,15 @@ import Control.Lens
 caseOfCase :: (fun ~ PLC.DefaultFun) => Term name uni fun a -> Term name uni fun a
 caseOfCase = transformOf termSubterms $ \case
   Case ann scrut alts
-    | ( ite@(Force _ (Builtin _ PLC.IfThenElse))
+    | ( ite@(Force a (Builtin _ PLC.IfThenElse))
         , [cond, (trueAnn, true@Constr{}), (falseAnn, false@Constr{})]
         ) <-
         splitApplication scrut ->
-        mkIterApp
-          ite
-          [cond, (trueAnn, Case ann true alts), (falseAnn, Case ann false alts)]
+        Force a $
+          mkIterApp
+            ite
+            [ cond
+            , (trueAnn, Delay trueAnn (Case ann true alts))
+            , (falseAnn, Delay falseAnn (Case ann false alts))
+            ]
   other -> other

plutus-core/untyped-plutus-core/test/Spec.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE DataKinds         #-}
 {-# LANGUAGE OverloadedStrings #-}
+
 module Main where
 
 import GHC.IO.Encoding (setLocaleEncoding, utf8)
@@ -10,29 +11,33 @@ import Evaluation.Builtins (test_builtins)
 import Evaluation.Debug (test_debug)
 import Evaluation.FreeVars (test_freevars)
 import Evaluation.Golden (test_golden)
-import Evaluation.Machines
+import Evaluation.Machines (test_budget, test_machines, test_tallying)
 import Evaluation.Regressions (schnorrVerifyRegressions)
 import Flat.Spec (test_flat)
 import Generators (test_parsing)
+import Transform.CaseOfCase.Test (test_caseOfCase)
 import Transform.Simplify (test_simplify)
 
 import Test.Tasty
 
 main :: IO ()
 main = do
   setLocaleEncoding utf8
-  defaultMain $ testGroup "Untyped Plutus Core"
-    [ test_machines
-    , test_builtins
-    , test_budget
-    , test_golden
-    , test_tallying
-    , test_simplify
-    , test_debruijn
-    , test_freevars
-    , test_parsing
-    , test_debug
-    , test_flat
-    , schnorrVerifyRegressions
-    , evalOrder
-    ]
+  defaultMain $
+    testGroup
+      "Untyped Plutus Core"
+      [ test_machines
+      , test_builtins
+      , test_budget
+      , test_caseOfCase
+      , test_golden
+      , test_tallying
+      , test_simplify
+      , test_debruijn
+      , test_freevars
+      , test_parsing
+      , test_debug
+      , test_flat
+      , schnorrVerifyRegressions
+      , evalOrder
+      ]

plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/1.uplc.golden

@@ -0,0 +1,9 @@
+(force
+  [
+    [
+      [ (force (builtin ifThenElse)) b_0 ]
+      (delay (case (constr 0) (con integer 1) (con integer 2)))
+    ]
+    (delay (case (constr 1) (con integer 1) (con integer 2)))
+  ]
+)

plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/2.uplc.golden

@@ -0,0 +1,5 @@
+(case
+  [ [ [ (force (builtin ifThenElse)) b_0 ] t_1 ] (constr 1) ]
+  (con integer 1)
+  (con integer 2)
+)

plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/3.uplc.golden

@@ -0,0 +1,9 @@
+(force
+  [
+    [
+      [ (force (builtin ifThenElse)) b_0 ]
+      (delay (case (constr 0 x_1 xs_2) f_3 (con integer 2)))
+    ]
+    (delay (case (constr 1) f_3 (con integer 2)))
+  ]
+)

plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/Test.hs

@@ -0,0 +1,138 @@
+{-# LANGUAGE BlockArguments    #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TypeApplications  #-}
+
+module Transform.CaseOfCase.Test where
+
+import Data.ByteString.Lazy qualified as BSL
+import Data.Text.Encoding (encodeUtf8)
+import Data.Vector qualified as V
+import PlutusCore qualified as PLC
+import PlutusCore.Evaluation.Machine.BuiltinCostModel (BuiltinCostModel)
+import PlutusCore.Evaluation.Machine.ExBudgetingDefaults (defaultBuiltinCostModel,
+                                                          defaultCekMachineCosts)
+import PlutusCore.Evaluation.Machine.MachineParameters (CostModel (..), MachineParameters,
+                                                        mkMachineParameters)
+import PlutusCore.MkPlc (mkConstant, mkIterApp)
+import PlutusCore.Pretty
+import PlutusCore.Quote (freshName, runQuote)
+import PlutusPrelude (Default (def))
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.Golden (goldenVsString)
+import Test.Tasty.HUnit (testCase, (@?=))
+import UntypedPlutusCore (DefaultFun, DefaultUni, Name, Term (..))
+import UntypedPlutusCore.Core qualified as UPLC
+import UntypedPlutusCore.Evaluation.Machine.Cek (CekMachineCosts, CekValue, EvaluationResult (..),
+                                                 noEmitter, unsafeEvaluateCek)
+import UntypedPlutusCore.Transform.CaseOfCase (caseOfCase)
+
+test_caseOfCase :: TestTree
+test_caseOfCase =
+  testGroup
+    "CaseOfCase"
+    [ goldenVsSimplified "1" caseOfCase1
+    , goldenVsSimplified "2" caseOfCase2
+    , goldenVsSimplified "3" caseOfCase3
+    , goldenVsSimplified "withError" caseOfCaseWithError
+    , testCaseOfCaseWithError
+    ]
+
+caseOfCase1 :: Term Name PLC.DefaultUni PLC.DefaultFun ()
+caseOfCase1 = runQuote do
+  b <- freshName "b"
+  let ite = Force () (Builtin () PLC.IfThenElse)
+      true = Constr () 0 []
+      false = Constr () 1 []
+      alts = V.fromList [mkConstant @Integer () 1, mkConstant @Integer () 2]
+  pure $ Case () (mkIterApp ite [((), Var () b), ((), true), ((), false)]) alts
+
+{- | This should not simplify, because one of the branches of `ifThenElse` is not a `Constr`.
+Unless both branches are known constructors, the case-of-case transformation
+may increase the program size.
+-}
+caseOfCase2 :: Term Name PLC.DefaultUni PLC.DefaultFun ()
+caseOfCase2 = runQuote do
+  b <- freshName "b"
+  t <- freshName "t"
+  let ite = Force () (Builtin () PLC.IfThenElse)
+      true = Var () t
+      false = Constr () 1 []
+      alts = V.fromList [mkConstant @Integer () 1, mkConstant @Integer () 2]
+  pure $ Case () (mkIterApp ite [((), Var () b), ((), true), ((), false)]) alts
+
+{- | Similar to `caseOfCase1`, but the type of the @true@ and @false@ branches is
+@[Integer]@ rather than Bool (note that @Constr 0@ has two parameters, @x@ and @xs@).
+-}
+caseOfCase3 :: Term Name PLC.DefaultUni PLC.DefaultFun ()
+caseOfCase3 = runQuote do
+  b <- freshName "b"
+  x <- freshName "x"
+  xs <- freshName "xs"
+  f <- freshName "f"
+  let ite = Force () (Builtin () PLC.IfThenElse)
+      true = Constr () 0 [Var () x, Var () xs]
+      false = Constr () 1 []
+      altTrue = Var () f
+      altFalse = mkConstant @Integer () 2
+      alts = V.fromList [altTrue, altFalse]
+  pure $ Case () (mkIterApp ite [((), Var () b), ((), true), ((), false)]) alts
+
+{- |
+
+@
+  case (force ifThenElse) True True False of
+    True -> ()
+    False -> _|_
+@
+
+Evaluates to `()` because the first case alternative is selected.
+(The _|_ is not evaluated because case alternatives are evaluated lazily).
+
+After the `CaseOfCase` transformation the program should evaluate to `()` as well.
+
+@
+  force ((force ifThenElse) True (delay ()) (delay _|_))
+@
+-}
+caseOfCaseWithError :: Term Name DefaultUni DefaultFun ()
+caseOfCaseWithError =
+  Case
+    ()
+    ( mkIterApp
+        (Force () (Builtin () PLC.IfThenElse))
+        [ ((), mkConstant @Bool () True)
+        , ((), Constr () 0 []) -- True
+        , ((), Constr () 1 []) -- False
+        ]
+    )
+    (V.fromList [mkConstant @() () (), Error ()])
+
+testCaseOfCaseWithError :: TestTree
+testCaseOfCaseWithError =
+  testCase "Transformation doesn't evaluate error eagerly" do
+    let simplifiedTerm = caseOfCase caseOfCaseWithError
+    evaluateUplc simplifiedTerm @?= evaluateUplc caseOfCaseWithError
+
+----------------------------------------------------------------------------------------------------
+-- Helper functions --------------------------------------------------------------------------------
+
+evaluateUplc
+  :: UPLC.Term Name DefaultUni DefaultFun ()
+  -> EvaluationResult (UPLC.Term Name DefaultUni DefaultFun ())
+evaluateUplc = fst <$> unsafeEvaluateCek noEmitter machineParameters
+ where
+  costModel :: CostModel CekMachineCosts BuiltinCostModel =
+    CostModel defaultCekMachineCosts defaultBuiltinCostModel
+  machineParameters
+    :: MachineParameters CekMachineCosts DefaultFun (CekValue DefaultUni DefaultFun ()) =
+      mkMachineParameters def costModel
+
+goldenVsSimplified :: String -> Term Name PLC.DefaultUni PLC.DefaultFun () -> TestTree
+goldenVsSimplified name =
+  goldenVsString name ("untyped-plutus-core/test/Transform/CaseOfCase/" ++ name ++ ".uplc.golden")
+    . pure
+    . BSL.fromStrict
+    . encodeUtf8
+    . render
+    . prettyClassicDebug
+    . caseOfCase

plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/withError.uplc.golden

@@ -0,0 +1,9 @@
+(force
+  [
+    [
+      [ (force (builtin ifThenElse)) (con bool True) ]
+      (delay (case (constr 0) (con unit ()) (error)))
+    ]
+    (delay (case (constr 1) (con unit ()) (error)))
+  ]
+)