From 253004fbb28a6b82de3561a62acdae28a4d258ef Mon Sep 17 00:00:00 2001
From: Ziyang Liu <unsafeFixIO@gmail.com>
Date: Thu, 9 May 2024 12:13:38 -0700
Subject: [PATCH] Avoid evaluating `nilCase` strictly in `matchList`, and
 rename the original `matchList` to `matchList'` (#5901)

---
 .../PlutusBenchmark/Lists/Lookup/Compiled.hs  |  6 ++--
 .../src/PlutusBenchmark/Lists/Sum/Compiled.hs |  4 +--
 .../9.6/builtinListIndexing.budget.golden     |  2 ++
 .../9.6/builtinListIndexing.eval.golden       |  1 +
 .../Budget/9.6/builtinListIndexing.pir.golden | 32 ++++++++++++++++++
 .../9.6/builtinListIndexing.uplc.golden       | 20 +++++++++++
 .../Budget/9.6/listIndexing.budget.golden     |  2 ++
 .../test/Budget/9.6/listIndexing.eval.golden  |  1 +
 .../test/Budget/9.6/listIndexing.pir.golden   | 31 +++++++++++++++++
 .../test/Budget/9.6/listIndexing.uplc.golden  | 19 +++++++++++
 plutus-tx-plugin/test/Budget/Spec.hs          | 33 +++++++++++++++++++
 ...0240509_114920_unsafeFixIO_index_budget.md |  5 +++
 plutus-tx/src/PlutusTx/Builtins.hs            | 14 +++++---
 13 files changed, 161 insertions(+), 9 deletions(-)
 create mode 100644 plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.budget.golden
 create mode 100644 plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.eval.golden
 create mode 100644 plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.pir.golden
 create mode 100644 plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.uplc.golden
 create mode 100644 plutus-tx-plugin/test/Budget/9.6/listIndexing.budget.golden
 create mode 100644 plutus-tx-plugin/test/Budget/9.6/listIndexing.eval.golden
 create mode 100644 plutus-tx-plugin/test/Budget/9.6/listIndexing.pir.golden
 create mode 100644 plutus-tx-plugin/test/Budget/9.6/listIndexing.uplc.golden
 create mode 100644 plutus-tx/changelog.d/20240509_114920_unsafeFixIO_index_budget.md

diff --git a/plutus-benchmark/lists/src/PlutusBenchmark/Lists/Lookup/Compiled.hs b/plutus-benchmark/lists/src/PlutusBenchmark/Lists/Lookup/Compiled.hs
index d8950458afe..e573a0df697 100644
--- a/plutus-benchmark/lists/src/PlutusBenchmark/Lists/Lookup/Compiled.hs
+++ b/plutus-benchmark/lists/src/PlutusBenchmark/Lists/Lookup/Compiled.hs
@@ -52,14 +52,14 @@ matchWithBuiltinLists :: Workload BI.BuiltinList -> Integer
 matchWithBuiltinLists (lixs, rixs, ls, rs) = go lixs rixs 0
   where
     go ltodo rtodo acc =
-      B.matchList
+      B.matchList'
         ltodo
         acc
-        (\lix lrest -> B.matchList rtodo acc
+        (\lix lrest -> B.matchList' rtodo acc
           (\rix rrest -> go lrest rrest
             ((ls !! lix) `B.addInteger` (rs !! rix) `B.addInteger` acc)))
     l !! ix =
-      B.matchList
+      B.matchList'
         l
         (\() -> P.traceError "empty list")
         (\h t -> \() -> if ix P.== 0 then h else t !! (ix `B.subtractInteger` 1))
diff --git a/plutus-benchmark/lists/src/PlutusBenchmark/Lists/Sum/Compiled.hs b/plutus-benchmark/lists/src/PlutusBenchmark/Lists/Sum/Compiled.hs
index 37908caa248..a07047d36aa 100644
--- a/plutus-benchmark/lists/src/PlutusBenchmark/Lists/Sum/Compiled.hs
+++ b/plutus-benchmark/lists/src/PlutusBenchmark/Lists/Sum/Compiled.hs
@@ -53,7 +53,7 @@ mkSumRightScottTerm l = compiledCodeToTerm $ mkSumRightScottCode l
 
 {-# INLINABLE foldLeftBuiltin #-}
 foldLeftBuiltin :: (b -> a -> b) -> b -> BI.BuiltinList a -> b
-foldLeftBuiltin f z l = B.matchList l z (\x xs -> (foldLeftBuiltin f (f z x) xs))
+foldLeftBuiltin f z l = B.matchList' l z (\x xs -> (foldLeftBuiltin f (f z x) xs))
 
 {-# INLINABLE sumLeftBuiltin #-}
 sumLeftBuiltin :: BI.BuiltinList Integer -> Integer
@@ -61,7 +61,7 @@ sumLeftBuiltin l = foldLeftBuiltin B.addInteger 0 l
 
 {-# INLINABLE foldRightBuiltin #-}
 foldRightBuiltin :: (a -> b -> b) -> b -> BI.BuiltinList a -> b
-foldRightBuiltin f z l = B.matchList l z (\x xs -> f x $! (foldRightBuiltin f z xs))
+foldRightBuiltin f z l = B.matchList' l z (\x xs -> f x $! (foldRightBuiltin f z xs))
 
 {-# INLINABLE sumRightBuiltin #-}
 sumRightBuiltin :: BI.BuiltinList Integer -> Integer
diff --git a/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.budget.golden b/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.budget.golden
new file mode 100644
index 00000000000..0ad82f19258
--- /dev/null
+++ b/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.budget.golden
@@ -0,0 +1,2 @@
+({cpu: 11720376
+| mem: 32730})
\ No newline at end of file
diff --git a/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.eval.golden b/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.eval.golden
new file mode 100644
index 00000000000..f5d13a24393
--- /dev/null
+++ b/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.eval.golden
@@ -0,0 +1 @@
+(con data (I 6))
\ No newline at end of file
diff --git a/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.pir.golden b/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.pir.golden
new file mode 100644
index 00000000000..f1593ac7d9f
--- /dev/null
+++ b/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.pir.golden
@@ -0,0 +1,32 @@
+let
+  data Unit | Unit_match where
+    Unit : Unit
+in
+letrec
+  !go : list data -> integer -> data
+    = \(xs : list data) (i : integer) ->
+        chooseList
+          {data}
+          {Unit -> Unit -> data}
+          xs
+          (\(ds : Unit) -> error {Unit -> data})
+          (\(ds : Unit) (ds : Unit) ->
+             let
+               !hd : data = headList {data} xs
+               !tl : list data = tailList {data} xs
+             in
+             ifThenElse
+               {all dead. data}
+               (equalsInteger 0 i)
+               (/\dead -> hd)
+               (/\dead -> go tl (subtractInteger i 1))
+               {all dead. dead})
+          Unit
+          Unit
+in
+let
+  data Bool | Bool_match where
+    True : Bool
+    False : Bool
+in
+\(d : data) -> let !xs : list data = unListData d in go xs 5
\ No newline at end of file
diff --git a/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.uplc.golden b/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.uplc.golden
new file mode 100644
index 00000000000..21752c6f5d3
--- /dev/null
+++ b/plutus-tx-plugin/test/Budget/9.6/builtinListIndexing.uplc.golden
@@ -0,0 +1,20 @@
+program
+  1.1.0
+  ((\go d -> go (unListData d) 5)
+     ((\s -> s s)
+        (\s xs i ->
+           force (force chooseList)
+             xs
+             (\ds -> error)
+             (\ds ds ->
+                (\hd ->
+                   (\tl ->
+                      force
+                        (force ifThenElse
+                           (equalsInteger 0 i)
+                           (delay hd)
+                           (delay (s s tl (subtractInteger i 1)))))
+                     (force tailList xs))
+                  (force headList xs))
+             (constr 0 [])
+             (constr 0 []))))
\ No newline at end of file
diff --git a/plutus-tx-plugin/test/Budget/9.6/listIndexing.budget.golden b/plutus-tx-plugin/test/Budget/9.6/listIndexing.budget.golden
new file mode 100644
index 00000000000..d40f47bf8e1
--- /dev/null
+++ b/plutus-tx-plugin/test/Budget/9.6/listIndexing.budget.golden
@@ -0,0 +1,2 @@
+({cpu: 10267419
+| mem: 32722})
\ No newline at end of file
diff --git a/plutus-tx-plugin/test/Budget/9.6/listIndexing.eval.golden b/plutus-tx-plugin/test/Budget/9.6/listIndexing.eval.golden
new file mode 100644
index 00000000000..f5d13a24393
--- /dev/null
+++ b/plutus-tx-plugin/test/Budget/9.6/listIndexing.eval.golden
@@ -0,0 +1 @@
+(con data (I 6))
\ No newline at end of file
diff --git a/plutus-tx-plugin/test/Budget/9.6/listIndexing.pir.golden b/plutus-tx-plugin/test/Budget/9.6/listIndexing.pir.golden
new file mode 100644
index 00000000000..fa5d613e7f7
--- /dev/null
+++ b/plutus-tx-plugin/test/Budget/9.6/listIndexing.pir.golden
@@ -0,0 +1,31 @@
+letrec
+  data (List :: * -> *) a | List_match where
+    Nil : List a
+    Cons : a -> List a -> List a
+in
+letrec
+  !go : integer -> List data -> data
+    = \(ds : integer) (ds : List data) ->
+        List_match
+          {data}
+          ds
+          {all dead. data}
+          (/\dead -> error {data})
+          (\(x : data) (xs : List data) ->
+             /\dead ->
+               ifThenElse
+                 {all dead. data}
+                 (equalsInteger 0 ds)
+                 (/\dead -> x)
+                 (/\dead -> go (subtractInteger ds 1) xs)
+                 {all dead. dead})
+          {all dead. dead}
+in
+let
+  data Bool | Bool_match where
+    True : Bool
+    False : Bool
+  data Unit | Unit_match where
+    Unit : Unit
+in
+\(xs : List data) -> go 5 xs
\ No newline at end of file
diff --git a/plutus-tx-plugin/test/Budget/9.6/listIndexing.uplc.golden b/plutus-tx-plugin/test/Budget/9.6/listIndexing.uplc.golden
new file mode 100644
index 00000000000..4044ad577c1
--- /dev/null
+++ b/plutus-tx-plugin/test/Budget/9.6/listIndexing.uplc.golden
@@ -0,0 +1,19 @@
+program
+  1.1.0
+  ((\go xs -> go 5 xs)
+     ((\s -> s s)
+        (\s ds ds ->
+           force
+             (case
+                ds
+                [ (delay error)
+                , (\x xs ->
+                     delay
+                       (force
+                          (force ifThenElse
+                             (equalsInteger 0 ds)
+                             (delay x)
+                             (delay
+                                ((\x -> s s x)
+                                   (subtractInteger ds 1)
+                                   xs))))) ]))))
\ No newline at end of file
diff --git a/plutus-tx-plugin/test/Budget/Spec.hs b/plutus-tx-plugin/test/Budget/Spec.hs
index 3c779f5177d..b9200ed45c8 100644
--- a/plutus-tx-plugin/test/Budget/Spec.hs
+++ b/plutus-tx-plugin/test/Budget/Spec.hs
@@ -20,6 +20,7 @@ import Budget.WithGHCOptimisations qualified as WithGHCOptTest
 import Budget.WithoutGHCOptimisations qualified as WithoutGHCOptTest
 import PlutusTx.AsData qualified as AsData
 import PlutusTx.Builtins qualified as PlutusTx
+import PlutusTx.Builtins.Internal qualified as BI
 import PlutusTx.Code
 import PlutusTx.IsData qualified as IsData
 import PlutusTx.Lift (liftCodeDef, makeLift)
@@ -201,6 +202,24 @@ tests = testNestedGhc "Budget" [
   , goldenPirReadable "null" compiledNull
   , goldenEvalCekCatch "null" [compiledNull]
 
+  , goldenUPlcReadable "listIndexing" compiledListIndexing
+  , goldenPirReadable "listIndexing" compiledListIndexing
+  , goldenEvalCekCatch
+      "listIndexing"
+      [compiledListIndexing `unsafeApplyCode` liftCodeDef listIndexingInput]
+  , goldenBudget
+      "listIndexing"
+      (compiledListIndexing `unsafeApplyCode` liftCodeDef listIndexingInput)
+
+  , goldenUPlcReadable "builtinListIndexing" compiledBuiltinListIndexing
+  , goldenPirReadable "builtinListIndexing" compiledBuiltinListIndexing
+  , goldenEvalCekCatch
+      "builtinListIndexing"
+      [compiledBuiltinListIndexing `unsafeApplyCode` liftCodeDef builtinListIndexingInput]
+  , goldenBudget
+      "builtinListIndexing"
+      (compiledBuiltinListIndexing `unsafeApplyCode` liftCodeDef builtinListIndexingInput)
+
   , goldenBudget "toFromData" compiledToFromData
   , goldenUPlcReadable "toFromData" compiledToFromData
   , goldenPirReadable "toFromData" compiledToFromData
@@ -461,6 +480,20 @@ compiledNull = $$(compile [||
       let ls = [1,2,3,4,5,6,7,8,9,10] :: [Integer]
        in PlutusTx.null ls ||])
 
+compiledListIndexing :: CompiledCode ([PlutusTx.BuiltinData] -> PlutusTx.BuiltinData)
+compiledListIndexing = $$(compile [||
+      \xs -> xs List.!! 5 ||])
+
+compiledBuiltinListIndexing :: CompiledCode (PlutusTx.BuiltinData -> PlutusTx.BuiltinData)
+compiledBuiltinListIndexing = $$(compile [||
+      \d -> BI.unsafeDataAsList d `List.indexBuiltinList` 5 ||])
+
+listIndexingInput :: [PlutusTx.BuiltinData]
+listIndexingInput = IsData.toBuiltinData <$> [1 :: Integer .. 10]
+
+builtinListIndexingInput :: PlutusTx.BuiltinData
+builtinListIndexingInput = IsData.toBuiltinData listIndexingInput
+
 compiledToFromData :: CompiledCode (Either Integer (Maybe (Bool, Integer, Bool)))
 compiledToFromData = $$(compile [||
       let
diff --git a/plutus-tx/changelog.d/20240509_114920_unsafeFixIO_index_budget.md b/plutus-tx/changelog.d/20240509_114920_unsafeFixIO_index_budget.md
new file mode 100644
index 00000000000..14265d10c5d
--- /dev/null
+++ b/plutus-tx/changelog.d/20240509_114920_unsafeFixIO_index_budget.md
@@ -0,0 +1,5 @@
+### Changed
+
+- Renamed `PlutusTx.Builtins.matchList` to `matchList'`. The new `matchList` takes
+  an argument of type `() -> r` for the `nil` case, ensuring that the nil case
+  isn't evaluated if the list is non-empty.
diff --git a/plutus-tx/src/PlutusTx/Builtins.hs b/plutus-tx/src/PlutusTx/Builtins.hs
index 775d24c9bfc..3f7c92e98e8 100644
--- a/plutus-tx/src/PlutusTx/Builtins.hs
+++ b/plutus-tx/src/PlutusTx/Builtins.hs
@@ -69,6 +69,7 @@ module PlutusTx.Builtins (
                          , pairToPair
                          -- * Lists
                          , matchList
+                         , matchList'
                          , headMaybe
                          , BI.head
                          , BI.tail
@@ -384,17 +385,22 @@ encodeUtf8 :: BuiltinString -> BuiltinByteString
 encodeUtf8 = BI.encodeUtf8
 
 {-# INLINABLE matchList #-}
-matchList :: forall a r . BI.BuiltinList a -> r -> (a -> BI.BuiltinList a -> r) -> r
-matchList l nilCase consCase = BI.chooseList l (const nilCase) (\_ -> consCase (BI.head l) (BI.tail l)) ()
+matchList :: forall a r . BI.BuiltinList a -> (() -> r) -> (a -> BI.BuiltinList a -> r) -> r
+matchList l nilCase consCase = BI.chooseList l nilCase (\_ -> consCase (BI.head l) (BI.tail l)) ()
+
+{-# INLINABLE matchList' #-}
+-- | Like `matchList` but evaluates @nilCase@ strictly.
+matchList' :: forall a r . BI.BuiltinList a -> r -> (a -> BI.BuiltinList a -> r) -> r
+matchList' l nilCase consCase = BI.chooseList l (const nilCase) (\_ -> consCase (BI.head l) (BI.tail l)) ()
 
 {-# INLINE headMaybe #-}
 headMaybe :: BI.BuiltinList a -> Maybe a
-headMaybe l = matchList l Nothing (\h _ -> Just h)
+headMaybe l = matchList' l Nothing (\h _ -> Just h)
 
 {-# INLINE uncons #-}
 -- | Uncons a builtin list, failing if the list is empty, useful in patterns.
 uncons :: BI.BuiltinList a -> Maybe (a, BI.BuiltinList a)
-uncons l = matchList l Nothing (\h t -> Just (h, t))
+uncons l = matchList' l Nothing (\h t -> Just (h, t))
 
 {-# INLINE unsafeUncons #-}
 -- | Uncons a builtin list, failing if the list is empty, useful in patterns.