From 99af5e26fac8307213bfdf7b27f5c27b45402a94 Mon Sep 17 00:00:00 2001
From: Emily Martins <emi@haskell.fyi>
Date: Fri, 18 Feb 2022 17:21:05 +0100
Subject: [PATCH] add Utils module, fix stake policy missing ST

---
 agora.cabal                 |   1 +
 src/Agora/AuthorityToken.hs |   2 +-
 src/Agora/SafeMoney.hs      |  58 +-------------
 src/Agora/Stake.hs          |  27 +------
 src/Agora/Utils.hs          | 151 ++++++++++++++++++++++++++++++++++++
 5 files changed, 160 insertions(+), 79 deletions(-)
 create mode 100644 src/Agora/Utils.hs

diff --git a/agora.cabal b/agora.cabal
index c874b3f..5f8ded4 100644
--- a/agora.cabal
+++ b/agora.cabal
@@ -121,6 +121,7 @@ library
     Agora.Voting
 
   other-modules:
+    Agora.Utils
   hs-source-dirs:  src
 
 library pprelude
diff --git a/src/Agora/AuthorityToken.hs b/src/Agora/AuthorityToken.hs
index 0734987..e32906e 100644
--- a/src/Agora/AuthorityToken.hs
+++ b/src/Agora/AuthorityToken.hs
@@ -20,7 +20,7 @@ import Prelude
 
 --------------------------------------------------------------------------------
 
-import Agora.SafeMoney
+import Agora.Utils (passetClassValueOf, passetClassValueOf')
 
 --------------------------------------------------------------------------------
 
diff --git a/src/Agora/SafeMoney.hs b/src/Agora/SafeMoney.hs
index 1f8faa4..f7abcaf 100644
--- a/src/Agora/SafeMoney.hs
+++ b/src/Agora/SafeMoney.hs
@@ -37,6 +37,10 @@ import Plutarch.Prelude
 
 --------------------------------------------------------------------------------
 
+import Agora.Utils
+
+--------------------------------------------------------------------------------
+
 -- | Type-level unique identifier for an AssetClass
 type MoneyClass =
   ( -- AssetClass
@@ -86,60 +90,6 @@ type ADA = '("", "", 6)
 
 --------------------------------------------------------------------------------
 
--- TODO: upstream something like this
-pfind' ::
-  PIsListLike list a =>
-  (Term s a -> Term s PBool) ->
-  Term s (list a :--> PMaybe a)
-pfind' p =
-  precList
-    (\self x xs -> pif (p x) (pcon (PJust x)) (self # xs))
-    (const $ pcon PNothing)
-
--- TODO: upstream something like this
-plookup ::
-  (PEq a, PIsListLike list (PBuiltinPair a b)) =>
-  Term s (a :--> list (PBuiltinPair a b) :--> PMaybe b)
-plookup =
-  phoistAcyclic $
-    plam $ \k xs ->
-      pmatch (pfind' (\p -> pfstBuiltin # p #== k) # xs) $ \case
-        PNothing -> pcon PNothing
-        PJust p -> pcon (PJust (psndBuiltin # p))
-
--- This is quite silly.
-plookupTuple ::
-  (PEq a, PIsListLike list (PAsData (PTuple a b)), PIsData a, PIsData b) =>
-  Term s (a :--> list (PAsData (PTuple a b)) :--> PMaybe b)
-plookupTuple =
-  phoistAcyclic $
-    plam $ \k xs ->
-      pmatch (pfind' (\p -> (pfield @"_0" # pfromData p) #== k) # xs) $ \case
-        PNothing -> pcon PNothing
-        PJust p -> pcon (PJust (pfield @"_1" # pfromData p))
-
-matchMaybe :: Term s r -> Term s (PMaybe a) -> TermCont @r s (Term s a)
-matchMaybe r f = TermCont $ \k ->
-  pmatch f $ \case
-    PJust v -> k v
-    PNothing -> r
-
-passetClassValueOf ::
-  Term s (PCurrencySymbol :--> PTokenName :--> PValue :--> PInteger)
-passetClassValueOf =
-  phoistAcyclic $
-    plam $ \sym token value'' -> unTermCont $ do
-      PValue value' <- tcont $ pmatch value''
-      PMap value <- tcont $ pmatch value'
-      m' <- matchMaybe 0 (plookup # pdata sym # value)
-      PMap m <- tcont (pmatch (pfromData m'))
-      v <- matchMaybe 0 (plookup # pdata token # m)
-      pure (pfromData v)
-
-passetClassValueOf' :: AssetClass -> Term s (PValue :--> PInteger)
-passetClassValueOf' (AssetClass (sym, token)) =
-  passetClassValueOf # pconstant sym # pconstant token
-
 -- | Downcast a 'PValue' to a 'Discrete' unit
 valueDiscrete ::
   forall (moneyClass :: MoneyClass) (ac :: Symbol) (n :: Symbol) (scale :: Nat) s.
diff --git a/src/Agora/Stake.hs b/src/Agora/Stake.hs
index 34d7230..1e1915c 100644
--- a/src/Agora/Stake.hs
+++ b/src/Agora/Stake.hs
@@ -31,6 +31,7 @@ import Plutarch.Monadic qualified as P
 --------------------------------------------------------------------------------
 
 import Agora.SafeMoney
+import Agora.Utils
 
 --------------------------------------------------------------------------------
 
@@ -66,28 +67,6 @@ newtype StakeDatum (gt :: MoneyClass) (s :: S) = StakeDatum
     (PlutusType, PIsData, PDataFields)
     via (PIsDataReprInstances (StakeDatum gt))
 
--- | Assert a particular bool, trace on falsehood. Use in monadic context
-passert :: Term s PString -> Term s PBool -> Term s k -> Term s k
-passert errorMessage check k = pif check k (ptraceError errorMessage)
-
--- | Find a datum with the given hash.
-pfindDatum :: Term s (PDatumHash :--> PTxInfo :--> PMaybe PDatum)
-pfindDatum = phoistAcyclic $
-  plam $ \datumHash txInfo'' -> P.do
-    PTxInfo txInfo' <- pmatch txInfo''
-    plookupTuple # datumHash #$ pfield @"data" # txInfo'
-
--- | Find a datum with the given hash. NOTE: this is unsafe in the sense that, if the data layout is wrong, this is UB.
-pfindDatum' :: PIsData a => Term s (PDatumHash :--> PTxInfo :--> PMaybe (PAsData a))
-pfindDatum' = phoistAcyclic $ plam $ \dh x -> punsafeCoerce $ pfindDatum # dh # x
-
--- | Check if a PubKeyHash signs this transaction
-ptxSignedBy :: Term s (PTxInfo :--> PAsData PPubKeyHash :--> PBool)
-ptxSignedBy = phoistAcyclic $
-  plam $ \txInfo' pkh -> P.do
-    txInfo <- pletFields @'["signatories"] txInfo'
-    pelem @PBuiltinList # pkh # txInfo.signatories
-
 -- | Check if any output matches the predicate
 anyOutput ::
   forall (datum :: PType) s.
@@ -125,14 +104,14 @@ stakePolicy _stake =
 
     PMinting ownSymbol <- pmatch $ pfromData ctx.purpose
     -- TODO: add this to 'valueCorrect'
-    let _stValue = psingletonValue # (pfield @"_0" # ownSymbol) # pconstant "ST" # 1
+    let stValue = psingletonValue # (pfield @"_0" # ownSymbol) # pconstant "ST" # 1
 
     passert "A UTXO must exist with the correct output" $
       anyOutput @(StakeDatum gt) # pfromData ctx.txInfo
         # ( plam $ \value stakeDatum' -> P.do
               stakeDatum <- pletFields @'["owner", "stakedAmount"] stakeDatum'
               let ownerSignsTransaction = ptxSignedBy # ctx.txInfo # stakeDatum.owner
-              let valueCorrect = pdata value #== pdata (discreteValue # stakeDatum.stakedAmount)
+              let valueCorrect = pdata value #== pdata (paddValue # (discreteValue # stakeDatum.stakedAmount) # stValue)
               ownerSignsTransaction #&& valueCorrect
           )
 
diff --git a/src/Agora/Utils.hs b/src/Agora/Utils.hs
new file mode 100644
index 0000000..dde004e
--- /dev/null
+++ b/src/Agora/Utils.hs
@@ -0,0 +1,151 @@
+-- | Plutarch utility functions that should be upstreamed or don't belong anywhere else
+module Agora.Utils (module Agora.Utils) where
+
+--------------------------------------------------------------------------------
+
+import Plutus.V1.Ledger.Value (AssetClass (..))
+
+--------------------------------------------------------------------------------
+
+import Plutarch.Api.V1 (
+  PCurrencySymbol,
+  PDatum,
+  PDatumHash,
+  PMap (PMap),
+  PPubKeyHash,
+  PTokenName,
+  PTuple,
+  PTxInfo (PTxInfo),
+  PValue (PValue),
+ )
+import Plutarch.Builtin (ppairDataBuiltin)
+import Plutarch.Internal (punsafeCoerce)
+import Plutarch.Monadic qualified as P
+
+--------------------------------------------------------------------------------
+
+-- | Assert a particular bool, trace on falsehood. Use in monadic context
+passert :: Term s PString -> Term s PBool -> Term s k -> Term s k
+passert errorMessage check k = pif check k (ptraceError errorMessage)
+
+-- | Find a datum with the given hash.
+pfindDatum :: Term s (PDatumHash :--> PTxInfo :--> PMaybe PDatum)
+pfindDatum = phoistAcyclic $
+  plam $ \datumHash txInfo'' -> P.do
+    PTxInfo txInfo' <- pmatch txInfo''
+    plookupTuple # datumHash #$ pfield @"data" # txInfo'
+
+-- | Find a datum with the given hash. NOTE: this is unsafe in the sense that, if the data layout is wrong, this is UB.
+pfindDatum' :: PIsData a => Term s (PDatumHash :--> PTxInfo :--> PMaybe (PAsData a))
+pfindDatum' = phoistAcyclic $ plam $ \dh x -> punsafeCoerce $ pfindDatum # dh # x
+
+-- | Check if a PubKeyHash signs this transaction
+ptxSignedBy :: Term s (PTxInfo :--> PAsData PPubKeyHash :--> PBool)
+ptxSignedBy = phoistAcyclic $
+  plam $ \txInfo' pkh -> P.do
+    txInfo <- pletFields @'["signatories"] txInfo'
+    pelem @PBuiltinList # pkh # txInfo.signatories
+
+-- | Get the first element that matches a predicate or return Nothing
+pfind' ::
+  PIsListLike list a =>
+  (Term s a -> Term s PBool) ->
+  Term s (list a :--> PMaybe a)
+pfind' p =
+  precList
+    (\self x xs -> pif (p x) (pcon (PJust x)) (self # xs))
+    (const $ pcon PNothing)
+
+-- | Find the value for a given key in an assoclist
+plookup ::
+  (PEq a, PIsListLike list (PBuiltinPair a b)) =>
+  Term s (a :--> list (PBuiltinPair a b) :--> PMaybe b)
+plookup =
+  phoistAcyclic $
+    plam $ \k xs ->
+      pmatch (pfind' (\p -> pfstBuiltin # p #== k) # xs) $ \case
+        PNothing -> pcon PNothing
+        PJust p -> pcon (PJust (psndBuiltin # p))
+
+-- | Find the value for a given key in an assoclist which uses 'PTuple's
+plookupTuple ::
+  (PEq a, PIsListLike list (PAsData (PTuple a b)), PIsData a, PIsData b) =>
+  Term s (a :--> list (PAsData (PTuple a b)) :--> PMaybe b)
+plookupTuple =
+  phoistAcyclic $
+    plam $ \k xs ->
+      pmatch (pfind' (\p -> (pfield @"_0" # pfromData p) #== k) # xs) $ \case
+        PNothing -> pcon PNothing
+        PJust p -> pcon (PJust (pfield @"_1" # pfromData p))
+
+-- | Extract a Maybe by providing a default value in case of Just
+pfromMaybe :: forall a s. Term s (a :--> PMaybe a :--> a)
+pfromMaybe = phoistAcyclic $
+  plam $ \e a ->
+    pmatch a $ \case
+      PJust a' -> a'
+      PNothing -> e
+
+-- | Escape with a particular value on expecting 'Just'. For use in monadic context
+pexpectJust :: forall r a s. Term s r -> Term s (PMaybe a) -> (Term s a -> Term s r) -> Term s r
+pexpectJust escape ma f =
+  pmatch ma $ \case
+    PJust v -> f v
+    PNothing -> escape
+
+passetClassValueOf ::
+  Term s (PCurrencySymbol :--> PTokenName :--> PValue :--> PInteger)
+passetClassValueOf =
+  phoistAcyclic $
+    plam $ \sym token value'' -> P.do
+      PValue value' <- pmatch value''
+      PMap value <- pmatch value'
+      m' <- pexpectJust 0 (plookup # pdata sym # value)
+      PMap m <- pmatch (pfromData m')
+      v <- pexpectJust 0 (plookup # pdata token # m)
+      pfromData v
+
+passetClassValueOf' :: AssetClass -> Term s (PValue :--> PInteger)
+passetClassValueOf' (AssetClass (sym, token)) =
+  passetClassValueOf # pconstant sym # pconstant token
+
+-- | Union two maps using a merge function on collisions
+pmapUnionWith :: forall k v s. PIsData v => Term s ((v :--> v :--> v) :--> PMap k v :--> PMap k v :--> PMap k v)
+pmapUnionWith = phoistAcyclic $
+  -- TODO: this function is kinda suspect. I feel like a lot of optimizations could be done here
+  plam $ \f xs' ys' -> P.do
+    PMap xs <- pmatch xs'
+    PMap ys <- pmatch ys'
+    let ls =
+          pmap
+            # ( plam $ \p -> P.do
+                  pf <- plet $ pfstBuiltin # p
+                  ps <- plet $ psndBuiltin # p
+                  pmatch (plookup # pf # ys) $ \case
+                    PJust v -> P.do
+                      -- Data conversions here are silly, aren't they?
+                      ppairDataBuiltin # pf # (pdata (f # pfromData ps # pfromData v))
+                    PNothing -> p
+              )
+            # xs
+        rs =
+          pfilter
+            # ( plam $ \p ->
+                  pnot # (pany # (plam $ \p' -> pfstBuiltin # p' #== pfstBuiltin # p) # xs)
+              )
+            # ys
+    pcon (PMap $ pconcat # ls # rs)
+
+-- | Add two 'PValue's together
+paddValue :: forall s. Term s (PValue :--> PValue :--> PValue)
+paddValue = phoistAcyclic $
+  plam $ \a' b' -> P.do
+    PValue a <- pmatch a'
+    PValue b <- pmatch b'
+    pcon
+      ( PValue $
+          pmapUnionWith
+            # (plam $ \a' b' -> pmapUnionWith # (plam (+)) # a' # b')
+            # a
+            # b
+      )