From 768652deb26d1e9339af8ecd4361a3c54051ef01 Mon Sep 17 00:00:00 2001
From: fanghr <chfanghr@gmail.com>
Date: Mon, 16 May 2022 18:11:14 +0800
Subject: [PATCH] docs for utils; refactor tests; rewrite `pnubSortBy`

---
 agora-test/Spec/Utils.hs        |  66 +++++++++--------
 agora/Agora/Proposal/Scripts.hs |   7 +-
 agora/Agora/Utils.hs            | 124 +++++++++++++++++++-------------
 3 files changed, 111 insertions(+), 86 deletions(-)

diff --git a/agora-test/Spec/Utils.hs b/agora-test/Spec/Utils.hs
index 596a6a3..45718de 100644
--- a/agora-test/Spec/Utils.hs
+++ b/agora-test/Spec/Utils.hs
@@ -7,58 +7,66 @@ Tests for utility functions in 'Agora.Utils'.
 -}
 module Spec.Utils (tests) where
 
-import Agora.Utils (phalve, pmerge, pmsortOrd, tcmatch)
+import Agora.Utils (phalve, pmergeBy, pmsort)
 import Data.List (sort)
 import Test.Tasty (TestTree)
 import Test.Tasty.QuickCheck (testProperty)
 
 tests :: [TestTree]
 tests =
-  [ testProperty "Merge sort sorts a list properly" prop_msort_sorted
-  , testProperty "Two sorted lists are merged into one sorted list" prop_pmerge_sorted
-  , testProperty "Split a list in half as expected" prop_halve_properly
+  [ testProperty "Merge sort sorts a list properly" prop_msortSorted
+  , testProperty "Two sorted lists are merged into one sorted list" prop_pmergeSorted
+  , testProperty "Split a list in half as expected" prop_halveProperly
   ]
 
 --------------------------------------------------------------------------------
 
-prop_msort_sorted :: [Integer] -> Bool
-prop_msort_sorted arr = sort arr == sorted
+prop_msortSorted :: [Integer] -> Bool
+prop_msortSorted arr = sorted == expected
   where
-    parr :: Term _ (PBuiltinList PInteger)
-    parr = pconstant arr
+    -- Expected sorted list, using 'Data.List.sort'.
+    expected :: [Integer]
+    expected = sort arr
+
+    --
 
     psorted :: Term _ (PBuiltinList PInteger)
-    psorted = pmsortOrd # parr
+    psorted = pmsort # pconstant arr
 
     sorted :: [Integer]
     sorted = plift psorted
 
-prop_pmerge_sorted :: ([Integer], [Integer]) -> Bool
-prop_pmerge_sorted (a, b) = merge sa sb == merged
+prop_pmergeSorted :: [Integer] -> [Integer] -> Bool
+prop_pmergeSorted a b = merged == expected
   where
+    -- Sorted list a and b
     sa = sort a
     sb = sort b
 
+    -- Merge two lists which are assumed to be ordered.
+    merge :: [Integer] -> [Integer] -> [Integer]
     merge xs [] = xs
     merge [] ys = ys
     merge sx@(x : xs) sy@(y : ys)
       | x <= y = x : merge xs sy
       | otherwise = y : merge sx ys
 
-    psa :: Term _ (PBuiltinList PInteger)
-    psa = pconstant @(PBuiltinList PInteger) sa
-    psb :: Term _ (PBuiltinList PInteger)
-    psb = pconstant @(PBuiltinList PInteger) sb
+    expected :: [Integer]
+    expected = merge sa sb
+
+    --
 
     pmerged :: Term _ (PBuiltinList PInteger)
-    pmerged = pmerge # plam (#<) # psa # psb
+    pmerged = pmergeBy # plam (#<) # pconstant sa # pconstant sb
 
     merged :: [Integer]
     merged = plift pmerged
 
-prop_halve_properly :: [Integer] -> Bool
-prop_halve_properly arr = halve arr == halved
+prop_halveProperly :: [Integer] -> Bool
+prop_halveProperly arr = halved == expected
   where
+    -- Halve a list.
+    halve :: [Integer] -> ([Integer], [Integer])
     halve xs = go xs xs
       where
         go xs [] = ([], xs)
@@ -69,26 +77,16 @@ prop_halve_properly arr = halve arr == halved
            in (x : first, last)
         go [] _ = ([], [])
 
-    parr :: Term _ (PBuiltinList PInteger)
-    parr = pconstant arr
+    expected :: ([Integer], [Integer])
+    expected = halve arr
 
-    ppairFst :: Term _ (PPair a b :--> a)
-    ppairFst = phoistAcyclic $
-      plam $ \p -> unTermCont $ do
-        PPair x _ <- tcmatch p
-        return x
-
-    ppairSnd :: Term _ (PPair a b :--> b)
-    ppairSnd = phoistAcyclic $
-      plam $ \p -> unTermCont $ do
-        PPair _ y <- tcmatch p
-        return y
+    --
 
     phalved :: Term _ (PPair (PBuiltinList PInteger) (PBuiltinList PInteger))
-    phalved = phalve # parr
+    phalved = phalve # pconstant arr
 
     halved :: ([Integer], [Integer])
     halved =
-      let f = plift $ ppairFst # phalved
-          s = plift $ ppairSnd # phalved
+      let f = plift $ pmatch phalved $ \(PPair x _) -> x
+          s = plift $ pmatch phalved $ \(PPair _ x) -> x
        in (f, s)
diff --git a/agora/Agora/Proposal/Scripts.hs b/agora/Agora/Proposal/Scripts.hs
index 57ba4bc..ce2dcbe 100644
--- a/agora/Agora/Proposal/Scripts.hs
+++ b/agora/Agora/Proposal/Scripts.hs
@@ -21,7 +21,7 @@ import Agora.Utils (
   anyOutput,
   findTxOutByTxOutRef,
   getMintingPolicySymbol,
-  pisUniq,
+  pisUniqBy,
   psymbolValueOf,
   ptokenSpent,
   ptxSignedBy,
@@ -171,7 +171,10 @@ proposalValidator proposal =
           newSigs <- tclet $ pfield @"newCosigners" # r
 
           tcassert "Cosigners are unique" $
-            pisUniq # phoistAcyclic (plam $ \(pfromData -> x) (pfromData -> y) -> x #< y) # newSigs
+            pisUniqBy
+              # phoistAcyclic (plam (#==))
+              # phoistAcyclic (plam $ \(pfromData -> x) (pfromData -> y) -> x #< y)
+              # newSigs
 
           tcassert "Signed by all new cosigners" $
             pall # signedBy # newSigs
diff --git a/agora/Agora/Utils.hs b/agora/Agora/Utils.hs
index 7ce3bf3..558bc13 100644
--- a/agora/Agora/Utils.hs
+++ b/agora/Agora/Utils.hs
@@ -32,13 +32,14 @@ module Agora.Utils (
   pisJust,
   ptokenSpent,
   pkeysEqual,
-  pnub,
+  pnubSortBy,
   pisUniq,
-  pisUniqOrd,
+  pisUniqBy,
   pisDJust,
   pisUTXOSpent,
+  pmsortBy,
   pmsort,
-  pmsortOrd,
+  pnubSort,
 
   -- * Functions which should (probably) not be upstreamed
   anyOutput,
@@ -56,7 +57,7 @@ module Agora.Utils (
   mustBePJust,
   mustBePDJust,
   validatorHashToAddress,
-  pmerge,
+  pmergeBy,
   phalve,
 ) where
 
@@ -371,46 +372,67 @@ pkeysEqual = phoistAcyclic $
         (plength # pks #== plength # qks)
         ( unTermCont $ do
             let comp = phoistAcyclic $ plam $ \(pfromData -> x) (pfromData -> y) -> x #< y
-                spks = pmsort # comp # pks
-                sqks = pmsort # comp # qks
+                spks = pmsortBy # comp # pks
+                sqks = pmsortBy # comp # qks
 
             pure $ plistEquals # spks # sqks
         )
         (pcon PFalse)
 
--- | / O(nlogn) /. Clear out duplicates in a list. The order is not preserved.
-pnub :: forall list a (s :: S). (PEq a, PIsListLike list a) => Term s ((a :--> a :--> PBool) :--> list a :--> list a)
-pnub = phoistAcyclic $
-  plam $ \comp xs -> unTermCont $ do
-    sorted <- tclet $ pmsort # comp # xs
-    pure $ pnubOrd # comp # sorted
-  where
-    pnubOrd = phoistAcyclic $ pfix #$ plam pnubOrd'
-    pnubOrd' self comp xs =
-      pif (pnull # xs) pnil $
-        unTermCont $ do
-          xh <- tclet $ phead # xs
-          xt <- tclet $ ptail # xs
+-- | / O(nlogn) /. Sort and remove dupicate elements in a list.
+pnubSortBy ::
+  forall list a (s :: S).
+  (PIsListLike list a) =>
+  Term s ((a :--> a :--> PBool) :--> (a :--> a :--> PBool) :--> list a :--> list a)
+pnubSortBy = phoistAcyclic $
+  plam $ \eq comp l -> pif (pnull # l) l $
+    unTermCont $ do
+      sl <- tclet $ pmsortBy # comp # l
 
-          pure $
-            pif (pnull # xt) xs $
-              unTermCont $ do
-                xh' <- tclet $ phead # xt
-                pure $
-                  pif
-                    (xh #== xh')
-                    (self # comp # xt)
-                    (pcons # xh #$ self # comp # xt)
+      let x = phead # sl
+          xs = ptail # sl
+
+      return $ pgo # eq # x # xs
+  where
+    pgo = phoistAcyclic pfix #$ plam pgo'
+    pgo' self eq seen l =
+      pif (pnull # l) (psingleton # seen) $
+        unTermCont $ do
+          x <- tclet $ phead # l
+          xs <- tclet $ ptail # l
+
+          return $
+            pif
+              (eq # x # seen)
+              (self # eq # seen # xs)
+              (pcons # seen #$ self # eq # x # xs)
+
+-- | Special version of 'pnubSortBy', which requires elements have 'POrd'.
+pnubSort ::
+  forall list a (s :: S).
+  (PIsListLike list a, POrd a) =>
+  Term s (list a :--> list a)
+pnubSort = phoistAcyclic $ pnubSortBy # eq # comp
+  where
+    eq = phoistAcyclic $ plam (#==)
+    comp = phoistAcyclic $ plam (#<)
 
 -- | / O(nlogn) /. Check if a list contains no duplicates.
-pisUniq :: forall list a (s :: S). (PEq a, PIsListLike list a) => Term s ((a :--> a :--> PBool) :--> list a :--> PBool)
-pisUniq = phoistAcyclic $
-  plam $ \comp xs ->
-    let nubbed = pnub # comp # xs in plength # xs #== plength # nubbed
+pisUniqBy ::
+  forall list a (s :: S).
+  (PIsListLike list a) =>
+  Term s ((a :--> a :--> PBool) :--> (a :--> a :--> PBool) :--> list a :--> PBool)
+pisUniqBy = phoistAcyclic $
+  plam $ \eq comp xs ->
+    let nubbed = pnubSortBy # eq # comp # xs
+     in plength # xs #== plength # nubbed
 
--- | List elements should have 'POrd' instance.
-pisUniqOrd :: forall list a (s :: S). (POrd a, PIsListLike list a) => Term s (list a :--> PBool)
-pisUniqOrd = phoistAcyclic $ pisUniq # plam (#<)
+-- | A special case of 'pisUniqBy' which requires elements have 'POrd' instance.
+pisUniq :: forall list a (s :: S). (POrd a, PIsListLike list a) => Term s (list a :--> PBool)
+pisUniq = phoistAcyclic $ pisUniqBy # eq # comp
+  where
+    eq = phoistAcyclic $ plam (#==)
+    comp = phoistAcyclic $ plam (#<)
 
 -- | Yield True if a given PMaybeData is of form @'PDJust' _@.
 pisDJust :: Term s (PMaybeData a :--> PBool)
@@ -423,19 +445,17 @@ pisDJust = phoistAcyclic $
           _ -> pconstant False
       )
 
-{- | Determines if a given UTXO is spent.
- TODO: no need to pass the whole TxInfo here.
--}
+-- | Determines if a given UTXO is spent.
 pisUTXOSpent :: Term s (PTxOutRef :--> PBuiltinList (PAsData PTxInInfo) :--> PBool)
 pisUTXOSpent = phoistAcyclic $
   plam $ \oref inputs -> P.do
     pisJust #$ pfindTxInByTxOutRef # oref # inputs
 
--- | Merge two ordered lists together.
-pmerge :: (PIsListLike l a) => Term s ((a :--> a :--> PBool) :--> l a :--> l a :--> l a)
-pmerge = phoistAcyclic $ pfix #$ plam pmerge'
+-- | / O(n) /. Merge two lists which are assumed to be ordered, given a custom comparator.
+pmergeBy :: (PIsListLike l a) => Term s ((a :--> a :--> PBool) :--> l a :--> l a :--> l a)
+pmergeBy = phoistAcyclic $ pfix #$ plam pmergeBy'
   where
-    pmerge' self comp a b =
+    pmergeBy' self comp a b =
       pif (pnull # a) b $
         pif (pnull # b) a $
           unTermCont $ do
@@ -450,20 +470,24 @@ pmerge = phoistAcyclic $ pfix #$ plam pmerge'
                 (pcons # ah #$ self # comp # at # b)
                 (pcons # bh #$ self # comp # a # bt)
 
--- | / O(nlogn) /. Merge sort, bottom-up version.
-pmsort :: (PIsListLike l a) => Term s ((a :--> a :--> PBool) :--> l a :--> l a)
-pmsort = phoistAcyclic $ pfix #$ plam pmsort'
+{- | / O(nlogn) /. Merge sort, bottom-up version, given a custom comparator.
+
+   Elements are arranged from lowest to highest,
+   keeping duplicates in the order they appeared in the input.
+-}
+pmsortBy :: (PIsListLike l a) => Term s ((a :--> a :--> PBool) :--> l a :--> l a)
+pmsortBy = phoistAcyclic $ pfix #$ plam pmsortBy'
   where
-    pmsort' self comp xs = pif (pnull # xs) pnil $
+    pmsortBy' self comp xs = pif (pnull # xs) pnil $
       pif (pnull #$ ptail # xs) xs $
         pmatch (phalve # xs) $ \(PPair fh sh) ->
           let sfh = self # comp # fh
               ssh = self # comp # sh
-           in pmerge # comp # sfh # ssh
+           in pmergeBy # comp # sfh # ssh
 
--- | Required list elements have 'POrd' instance.
-pmsortOrd :: (POrd a, PIsListLike l a) => Term s (l a :--> l a)
-pmsortOrd = phoistAcyclic $ pmsort # comp
+-- | A special case of 'pmsortBy' which requires elements have 'POrd' instance.
+pmsort :: (POrd a, PIsListLike l a) => Term s (l a :--> l a)
+pmsort = phoistAcyclic $ pmsortBy # comp
   where
     comp = phoistAcyclic $ plam (#<)