From 139565421f7ff9cc0377928ab35a17024d260411 Mon Sep 17 00:00:00 2001
From: Skylar-Ray O'Quinn <soraskymizu@gmail.com>
Date: Wed, 8 Dec 2021 05:03:18 +0000
Subject: [PATCH] Write arbitrary instance for data that uses sized

By using sized we can have an arbitrary instance that actually
terminates and is configurable via standard QuickCheck
methods (e.g. generate, resize). This will allow testing with Datums and
will prevent users unintentionally stalling their tests when generating
Datums, BuiltinData, and Data. Fixes #172.
---
 plutus-pab/src/Plutus/PAB/Arbitrary.hs | 52 ++++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 3 deletions(-)

diff --git a/plutus-pab/src/Plutus/PAB/Arbitrary.hs b/plutus-pab/src/Plutus/PAB/Arbitrary.hs
index 71aa268834..5a060a71d8 100644
--- a/plutus-pab/src/Plutus/PAB/Arbitrary.hs
+++ b/plutus-pab/src/Plutus/PAB/Arbitrary.hs
@@ -8,6 +8,7 @@
 -- across to the test suite.
 module Plutus.PAB.Arbitrary where
 
+import Control.Monad (replicateM)
 import Data.Aeson (Value)
 import Data.Aeson qualified as Aeson
 import Data.ByteString (ByteString)
@@ -29,7 +30,7 @@ import Plutus.Contract.StateMachine (ThreadToken)
 import PlutusTx qualified
 import PlutusTx.AssocMap qualified as AssocMap
 import PlutusTx.Prelude qualified as PlutusTx
-import Test.QuickCheck (Gen, oneof)
+import Test.QuickCheck (Gen, Positive (..), oneof, sized)
 import Test.QuickCheck.Arbitrary.Generic (Arbitrary, arbitrary, genericArbitrary, genericShrink, shrink)
 import Test.QuickCheck.Instances ()
 import Wallet (WalletAPIError)
@@ -156,8 +157,53 @@ instance Arbitrary ThreadToken where
     shrink = genericShrink
 
 instance Arbitrary PlutusTx.Data where
-    arbitrary = genericArbitrary
-    shrink = genericShrink
+    arbitrary = sized arbitraryData
+      where
+        arbitraryData :: Int -> Gen PlutusTx.Data
+        arbitraryData n =
+          oneof [ arbitraryConstr n
+                , arbitraryMap n
+                , arbitraryList n
+                , arbitraryI
+                , arbitraryB
+                ]
+
+        arbitraryConstr 0 = do
+          (Positive ix) <- arbitrary
+          pure $ PlutusTx.Constr ix []
+        arbitraryConstr n = do
+          (n', m) <- segmentRange n
+          (Positive ix) <- arbitrary
+          args <- replicateM m (arbitraryData n')
+          pure $ PlutusTx.Constr ix args
+
+        arbitraryMap 0 = pure $ PlutusTx.Map []
+        arbitraryMap n = do
+          (n', m) <- segmentRange n
+          kvs <- replicateM m (arbitraryPair n')
+          pure $ PlutusTx.Map kvs
+
+        arbitraryPair n = do
+          (n', _) <- segmentRange n
+          (,) <$> arbitraryData n' <*> arbitraryData n'
+
+        arbitraryList 0 = pure $ PlutusTx.List []
+        arbitraryList n = do
+          (n', m) <- segmentRange n
+          es <- replicateM m (arbitraryData n')
+          pure $ PlutusTx.List es
+
+        arbitraryI =
+          PlutusTx.I <$> arbitrary
+
+        arbitraryB =
+          PlutusTx.B <$> arbitrary
+
+        -- Used to break the sized generator up more or less evenly
+        segmentRange n = do
+          (Positive m) <- arbitrary
+          let n' = n `div` (m + 1) -- Prevent division by 0
+          pure (n', m)
 
 instance Arbitrary PlutusTx.BuiltinData where
     arbitrary = PlutusTx.dataToBuiltinData <$> arbitrary