feat: support RequiredMint for both adv. endpts

Anastasia-Labs · Sep 4, 2024 · 3aca0d9 · 3aca0d9
1 parent fe30aef
commit 3aca0d9
Show file tree

Hide file tree

Showing 4 changed files with 105 additions and 86 deletions.
diff --git a/src/SingleValidator.hs b/src/SingleValidator.hs
@@ -1,11 +1,9 @@
 module SingleValidator where
 
-import PlutusLedgerApi.V2 (Address, CurrencySymbol, TokenName)
+import PlutusLedgerApi.V2 (Address)
 import PlutusTx qualified
 
 import Plutarch.Api.V1.Address (PCredential (..))
-import Plutarch.Api.V1.Value (PCurrencySymbol, PTokenName, passertPositive, pforgetPositive)
-import Plutarch.Api.V1.Value qualified as Value
 import Plutarch.Api.V2 (PAddress, PMaybeData (..), PScriptContext, PScriptHash, PScriptPurpose (..), PTxInInfo, PTxOut, PTxOutRef, PValidator)
 import Plutarch.DataRepr
 import Plutarch.Lift (PConstantDecl, PUnsafeLiftDecl (..))
@@ -16,7 +14,7 @@ import Plutarch.Unsafe (punsafeCoerce)
 import "liqwid-plutarch-extra" Plutarch.Extra.ScriptContext ()
 
 import Constants (negativeRouterFeeForSimpleRoutes, routerFeeForSimpleRoutes)
-import Utils (PCustomValidator, pand'List, pconvertChecked, pconvertUnsafe, pmintIsSameAsSingleton, presolveDatum, psignedByOwner, pvalueHasChangedByLovelaces)
+import Utils (PCustomValidator, PRequiredMint (..), RequiredMint, pand'List, papplyRequiredMintToInputValue, pconvertChecked, pconvertUnsafe, presolveDatum, psignedByOwner, pvalueHasChangedByLovelaces)
 
 pcountInputsAtScript :: Term s (PScriptHash :--> PBuiltinList PTxInInfo :--> PInteger)
 pcountInputsAtScript =
@@ -33,46 +31,12 @@ pcountInputsAtScript =
             n
      in go # 0
 
-data RequiredMint
-  = Singleton CurrencySymbol TokenName Integer
-  | None
-
-PlutusTx.makeLift ''RequiredMint
-PlutusTx.makeIsDataIndexed
-  ''RequiredMint
-  [ ('Singleton, 0)
-  , ('None, 1)
-  ]
-
-data PRequiredMint (s :: S)
-  = PSingleton
-      ( Term
-          s
-          ( PDataRecord
-              '[ "policy" ':= PCurrencySymbol
-               , "name" ':= PTokenName
-               , "quantity" ':= PInteger
-               ]
-          )
-      )
-  | PNone (Term s (PDataRecord '[]))
-  deriving stock (Generic)
-  deriving anyclass (PlutusType, PIsData)
-
-instance DerivePlutusType PRequiredMint where
-  type DPTStrat _ = PlutusTypeData
-
-instance PTryFrom PData PRequiredMint
-
-instance PUnsafeLiftDecl PRequiredMint where type PLifted PRequiredMint = RequiredMint
-deriving via (DerivePConstantViaData RequiredMint PRequiredMint) instance PConstantDecl RequiredMint
-
 data SmartHandleDatum
   = Simple Address -- <-- owner
-  | Advanced (Maybe Address) Integer Integer RequiredMint PlutusTx.BuiltinData
+  | Advanced (Maybe Address) Integer Integer RequiredMint RequiredMint PlutusTx.BuiltinData
 
---           ^-------------^ ^-----^ ^-----^              ^------------------^
---               mOwner    routerFee reclaimRouterFee          extraInfo
+--           ^-------------^ ^-----^ ^-----^                           ^------------------^
+--               mOwner    routerFee reclaimRouterFee                       extraInfo
 
 PlutusTx.makeLift ''SmartHandleDatum
 PlutusTx.makeIsDataIndexed
@@ -97,7 +61,8 @@ data PSmartHandleDatum (s :: S)
               '[ "mOwner" ':= PMaybeData PAddress
                , "routerFee" ':= PInteger
                , "reclaimRouterFee" ':= PInteger
-               , "requiredMint" ':= PRequiredMint
+               , "routeRequiredMint" ':= PRequiredMint
+               , "reclaimRequiredMint" ':= PRequiredMint
                , "extraInfo" ':= PData
                ]
           )
@@ -260,20 +225,11 @@ prouter = phoistAcyclic $ plam $ \validateFn routeAddress dat ownIndex routerInd
                 , ptraceIfFalse "Incorrect Route Output Value" (pvalueHasChangedByLovelaces # ownInputF.value # routerOutputF.value # negativeRouterFeeForSimpleRoutes)
                 ]
             PAdvanced dat' -> P.do
-              datF <- pletFields @'["mOwner", "routerFee", "reclaimRouterFee", "requiredMint", "extraInfo"] dat'
-              let inputIncludingMint = pmatch datF.requiredMint $ \case
-                    PSingleton rm -> P.do
-                      rmF <- pletFields @'["policy", "name", "quantity"] rm
-                      let requiredMintValue = Value.psingleton # rmF.policy # rmF.name # rmF.quantity
-                          inputAppendedWithMint = requiredMintValue <> pforgetPositive ownInputF.value
-                      pif
-                        ( ptraceIfFalse
-                            "Tx mint doesn't match the reclaim mint"
-                            (pmintIsSameAsSingleton rmF.policy rmF.name rmF.quantity infoF.mint)
-                        )
-                        (passertPositive # inputAppendedWithMint)
-                        perror
-                    PNone _ ->
+              datF <- pletFields @'["mOwner", "routerFee", "reclaimRouterFee", "routeRequiredMint", "reclaimRequiredMint", "extraInfo"] dat'
+              let inputIncludingMint =
+                    papplyRequiredMintToInputValue
+                      infoF.mint
+                      (pif forRoute datF.routeRequiredMint datF.reclaimRequiredMint)
                       ownInputF.value
               routerFee <- plet $ pif forRoute datF.routerFee datF.reclaimRouterFee
               pand'List

diff --git a/src/StakingValidator.hs b/src/StakingValidator.hs
@@ -28,7 +28,7 @@ import "liqwid-plutarch-extra" Plutarch.Extra.TermCont
 import BatchValidator (PSmartRedeemer (..))
 import Constants (negativeRouterFeeForSimpleRoutes, routerFeeForSimpleRoutes)
 import Plutarch.Builtin (PIsData (pdataImpl), ppairDataBuiltin)
-import SingleValidator (PRequiredMint (..), PSmartHandleDatum (..))
+import SingleValidator (PSmartHandleDatum (..))
 import Utils
 
 pcountScriptInputs :: Term s (PBuiltinList PTxInInfo :--> PInteger)
@@ -125,38 +125,38 @@ psmartHandleSuccessor validateFn datums ctx routeAddress smartInputRouteFlagPair
             )
             (psignedByOwner # ctx # owner)
         PAdvanced dat' -> P.do
-          datF <- pletFields @'["mOwner", "routerFee", "reclaimRouterFee", "requiredMint", "extraInfo"] dat'
+          let txInfo = pfield @"txInfo" # ctx
+              mint = pfield @"mint" # txInfo
+          datF <- pletFields @'["mOwner", "routerFee", "reclaimRouterFee", "routeRequiredMint", "reclaimRequiredMint", "extraInfo"] dat'
           pif
             forRoute
-            ( pand'List
-                [ validateFn # datF.mOwner # datF.routerFee # smartInputF.value # datF.extraInfo # routeOutputDatum # forRoute # ctx
-                , ptraceIfFalse "Incorrect Route Output Value" (pvalueHasChangedByLovelaces # smartInputF.value # routeOutputF.value # (pnegate # datF.routerFee))
-                , ptraceIfFalse "Incorrect Route Address" (routeOutputF.address #== routeAddress)
-                ]
+            ( let
+                inputIncludingMint =
+                  papplyRequiredMintToInputValue
+                    mint
+                    datF.routeRequiredMint
+                    smartInputF.value
+               in
+                pand'List
+                  [ validateFn # datF.mOwner # datF.routerFee # smartInputF.value # datF.extraInfo # routeOutputDatum # forRoute # ctx
+                  , ptraceIfFalse "Incorrect Route Output Value" (pvalueHasChangedByLovelaces # inputIncludingMint # routeOutputF.value # (pnegate # datF.routerFee))
+                  , ptraceIfFalse "Incorrect Route Address" (routeOutputF.address #== routeAddress)
+                  ]
             )
             ( pmatch (pfield @"mOwner" # dat') $ \case
-                PDJust ((pfield @"_0" #) -> owner) -> P.do
-                  let txInfo = pfield @"txInfo" # ctx
-                      mint = pfield @"mint" # txInfo
-                      inputIncludingMint = pmatch datF.requiredMint $ \case
-                        PSingleton rm -> P.do
-                          rmF <- pletFields @'["policy", "name", "quantity"] rm
-                          let requiredMintValue = Value.psingleton # rmF.policy # rmF.name # rmF.quantity
-                              inputAppendedWithMint = requiredMintValue <> pforgetPositive smartInputF.value
-                          pif
-                            ( ptraceIfFalse
-                                "Tx mint doesn't match the reclaim mint"
-                                (pmintIsSameAsSingleton rmF.policy rmF.name rmF.quantity mint)
-                            )
-                            (passertPositive # inputAppendedWithMint)
-                            perror
-                        PNone _ ->
-                          smartInputF.value
-                  pand'List
-                    [ validateFn # datF.mOwner # datF.reclaimRouterFee # smartInputF.value # datF.extraInfo # routeOutputDatum # forRoute # ctx
-                    , ptraceIfFalse "Incorrect Route Output Value" (pvalueHasChangedByLovelaces # inputIncludingMint # routeOutputF.value # (pnegate # datF.reclaimRouterFee))
-                    , ptraceIfFalse "Incorrect Route Address" (routeOutputF.address #== owner)
-                    ]
+                PDJust ((pfield @"_0" #) -> owner) ->
+                  let
+                    inputIncludingMint =
+                      papplyRequiredMintToInputValue
+                        mint
+                        datF.reclaimRequiredMint
+                        smartInputF.value
+                   in
+                    pand'List
+                      [ validateFn # datF.mOwner # datF.reclaimRouterFee # smartInputF.value # datF.extraInfo # routeOutputDatum # forRoute # ctx
+                      , ptraceIfFalse "Incorrect Reclaim Output Value" (pvalueHasChangedByLovelaces # inputIncludingMint # routeOutputF.value # (pnegate # datF.reclaimRouterFee))
+                      , ptraceIfFalse "Incorrect Reclaim Address" (routeOutputF.address #== owner)
+                      ]
                 PDNothing _ ->
                   perror
             )

diff --git a/src/Utils.hs b/src/Utils.hs
@@ -2,13 +2,19 @@
 
 module Utils where
 
+import PlutusLedgerApi.V2 (CurrencySymbol, TokenName)
+import PlutusTx qualified
+
 import Plutarch.Api.V1.Address (PCredential (..))
 import Plutarch.Api.V1.AssocMap (plookup)
 import Plutarch.Api.V1.Value (padaSymbol, padaToken, pforgetPositive, psingleton)
+import Plutarch.Api.V1.Value qualified as Value
 import Plutarch.Api.V2
 import Plutarch.Bool
 import Plutarch.DataRepr
+import Plutarch.Lift (PConstantDecl, PUnsafeLiftDecl (..))
 import Plutarch.Maybe (pfromJust)
+import Plutarch.Monadic qualified as P
 import Plutarch.Prelude hiding (psingleton)
 import Plutarch.Unsafe (punsafeCoerce)
 import "liqwid-plutarch-extra" Plutarch.Extra.List (plookupAssoc)
@@ -34,6 +40,40 @@ instance DerivePlutusType PAssetClass where
 
 instance PTryFrom PData PAssetClass
 
+data RequiredMint
+  = Singleton CurrencySymbol TokenName Integer
+  | None
+
+PlutusTx.makeLift ''RequiredMint
+PlutusTx.makeIsDataIndexed
+  ''RequiredMint
+  [ ('Singleton, 0)
+  , ('None, 1)
+  ]
+
+data PRequiredMint (s :: S)
+  = PSingleton
+      ( Term
+          s
+          ( PDataRecord
+              '[ "policy" ':= PCurrencySymbol
+               , "name" ':= PTokenName
+               , "quantity" ':= PInteger
+               ]
+          )
+      )
+  | PNone (Term s (PDataRecord '[]))
+  deriving stock (Generic)
+  deriving anyclass (PlutusType, PIsData)
+
+instance DerivePlutusType PRequiredMint where
+  type DPTStrat _ = PlutusTypeData
+
+instance PTryFrom PData PRequiredMint
+
+instance PUnsafeLiftDecl PRequiredMint where type PLifted PRequiredMint = RequiredMint
+deriving via (DerivePConstantViaData RequiredMint PRequiredMint) instance PConstantDecl RequiredMint
+
 pexpectJust :: Term s r -> Term s (PMaybe a) -> TermCont @r s (Term s a)
 pexpectJust escape ma = tcont $ \f -> pmatch ma $ \case
   PJust v -> f v
@@ -232,3 +272,24 @@ pmintIsSameAsSingleton policy name qty mintVal =
           , name #== pfromData mintTN
           , qty #== pfromData mintQty
           ]
+
+papplyRequiredMintToInputValue ::
+  Term s (PValue 'Sorted 'NoGuarantees) ->
+  Term s PRequiredMint ->
+  Term s (PValue 'Sorted 'Positive) ->
+  Term s (PValue 'Sorted 'Positive)
+papplyRequiredMintToInputValue mint requiredMint inputValue =
+  pmatch requiredMint $ \case
+    PSingleton rm -> P.do
+      rmF <- pletFields @'["policy", "name", "quantity"] rm
+      let requiredMintValue = Value.psingleton # rmF.policy # rmF.name # rmF.quantity
+          inputAppendedWithMint = requiredMintValue <> pforgetPositive inputValue
+      pif
+        ( ptraceIfFalse
+            "Tx mint doesn't match the reclaim mint"
+            (pmintIsSameAsSingleton rmF.policy rmF.name rmF.quantity mint)
+        )
+        (Value.passertPositive # inputAppendedWithMint)
+        perror
+    PNone _ ->
+      inputValue
diff --git a/test/SmartHandlesSpec.hs b/test/SmartHandlesSpec.hs
@@ -19,9 +19,10 @@ import Test.Tasty.QuickCheck (Gen, Property, chooseInt, chooseInteger, forAll, s
 import BatchValidator (SmartRedeemer (..), smartHandleRouteValidatorW)
 import Compilation
 import Debug.Trace (trace)
-import SingleValidator (RequiredMint (..), SmartHandleDatum (..))
+import SingleValidator (SmartHandleDatum (..))
 import Specialized.Minswap
 import StakingValidator (RouterRedeemer (..), puniqueOrdered)
+import Utils (RequiredMint (..))
 
 tests :: TestTree
 tests = testGroup "Smart handles" [uniqueOrderedTests, stakingValidatorTests]
@@ -110,6 +111,7 @@ scriptInput =
             1_000_000
             0
             None
+            None
             ( toBuiltinData $
                 MinswapRequestInfo
                   (fst $ unAssetClass minAssetClass)