Fix CaseOfCase UPLC transform.

IntersectMBO · May 7, 2024 · 9aa9b46 · 9aa9b46
1 parent ea8c8c3
commit 9aa9b46
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 7 deletions.
diff --git a/plutus-core/untyped-plutus-core/src/UntypedPlutusCore/Transform/CaseOfCase.hs b/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
diff --git a/plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/1.uplc.golden b/plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/1.uplc.golden
@@ -1 +1 @@
-(force ifThenElse b_0 1 2)
+(force (force ifThenElse b_0 (delay 1) (delay 2)))
diff --git a/plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/3.uplc.golden b/plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/3.uplc.golden
@@ -1 +1 @@
-(force ifThenElse b_0 (f_3 x_1 xs_2) 2)
+(force (force ifThenElse b_0 (delay (f_3 x_1 xs_2)) (delay 2)))
diff --git a/plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/Test.hs b/plutus-core/untyped-plutus-core/test/Transform/CaseOfCase/Test.hs
@@ -73,11 +73,29 @@ caseOfCase3 = runQuote do
       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 :: TestTree
 caseOfCaseWithError =
   testCase "Transformation doesn't evaluate error eagerly" do
     let
-      originalTerm =
+       originalTerm =
         Case
           ()
           ( mkIterApp