Add reward withdrawals.

lib/core/src/Cardano/Wallet/Primitive/Migration/Selection.hs

@@ -28,6 +28,7 @@ module Cardano.Wallet.Primitive.Migration.Selection
     , finalize
     -- * Extending a selection
     , addEntry
+    , addRewardWithdrawal
 
     ----------------------------------------------------------------------------
     -- Internal interface
@@ -101,12 +102,16 @@ data SelectionParameters s = SelectionParameters
       -- ^ The constant fee for a selection input.
     , feeForOutput :: TokenBundle -> Coin
       -- ^ The variable fee for a selection output.
+    , feeForRewardWithdrawal :: Coin
+      -- ^ The constant fee for a reward withdrawal.
     , sizeOfEmptySelection :: s
       -- ^ The constant size of an empty selection.
     , sizeOfInput :: s
       -- ^ The constant size of a selection input.
     , sizeOfOutput :: TokenBundle -> s
       -- ^ The variable size of a selection output.
+    , sizeOfRewardWithdrawal :: s
+      -- ^ The constant size of a reward withdrawal.
     , maximumSizeOfOutput :: SelectionOutputSizeAssessor
       -- ^ The maximum size of a selection output.
     , maximumSizeOfSelection :: s
@@ -179,6 +184,8 @@ data Selection i s = Selection
       -- ^ The excess over the minimum permissible fee for this selection.
     , size :: !s
       -- ^ The current size of this selection.
+    , rewardWithdrawal :: !Coin
+      -- ^ The current reward withdrawal amount, if any.
     }
     deriving (Eq, Generic, Show)
 
@@ -474,13 +481,14 @@ checkSizeWithinLimit params selection
 -- | Calculates the current fee for a selection.
 --
 currentFee :: Selection i s -> Either NegativeCoin Coin
-currentFee Selection {inputs, outputs}
+currentFee Selection {inputs, outputs, rewardWithdrawal}
     | adaBalanceIn >= adaBalanceOut =
         Right adaDifference
     | otherwise =
         Left (NegativeCoin adaDifference)
   where
-    adaBalanceIn  = F.foldMap (TokenBundle.getCoin . snd) inputs
+    adaBalanceIn = F.foldMap (TokenBundle.getCoin . snd) inputs
+        <> rewardWithdrawal
     adaBalanceOut = F.foldMap (TokenBundle.getCoin) outputs
     adaDifference = Coin.distance adaBalanceIn adaBalanceOut
 
@@ -502,6 +510,9 @@ minimumFee params selection = mconcat
     [ feeForEmptySelection
     , F.foldMap (const feeForInput) (inputs selection)
     , F.foldMap feeForOutput (outputs selection)
+    , if (rewardWithdrawal selection > Coin 0)
+        then feeForRewardWithdrawal params
+        else Coin 0
     ]
   where
     SelectionParameters
@@ -560,6 +571,7 @@ initialize params entries = do
         , outputs = []
         , feeExcess = Coin 0
         , size = mempty
+        , rewardWithdrawal = Coin 0
         }
 
 --------------------------------------------------------------------------------
@@ -613,6 +625,16 @@ addEntryWithFirstSuccessfulStrategy
 addEntryWithFirstSuccessfulStrategy strategies params selection input =
     eithersToEither $ strategies <&> (\s -> s params selection input)
 
+addRewardWithdrawal
+    :: Selection i s
+    -> Coin
+    -> Selection i s
+addRewardWithdrawal selection withdrawal = selection
+    -- TODO: check that the invariant is not violated.
+    { rewardWithdrawal = rewardWithdrawal selection <> withdrawal
+    , feeExcess = feeExcess selection <> withdrawal
+    }
+
 --------------------------------------------------------------------------------
 -- Adding coins to a selection
 --------------------------------------------------------------------------------

lib/core/test/unit/Cardano/Wallet/Primitive/Migration/SelectionSpec.hs

@@ -75,6 +75,7 @@ import Test.QuickCheck
     , vector
     )
 
+import qualified Cardano.Wallet.Primitive.Types.Coin as Coin
 import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
 import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
 import qualified Data.ByteString as BS
@@ -212,9 +213,43 @@ genMockSelection mockParams =
         ]
 
 genMockSelectionSmall :: MockSelectionParameters -> Gen MockSelection
-genMockSelectionSmall _mockParams = do
-    -- lots of possibilities here
-    undefined
+genMockSelectionSmall mockParams = oneof
+    [ genSingleInputCoinNoOutput
+    , genSingleInputCoinSingleOutput
+    --, genSingleInputBundleSingleOutput
+    --, genMultipleInputBundlesSingleOutput
+    ]
+  where
+    genSingleInputCoinNoOutput :: Gen MockSelection
+    genSingleInputCoinNoOutput = do
+        coin <- genCoin
+        inputId <- genMockInputId
+        pure Selection
+            { inputs = [(inputId, TokenBundle.fromCoin coin)]
+            , outputs = []
+            , feeExcess = coin `Coin.distance` fee
+            , size = sizeOfInput params <> sizeOfEmptySelection params
+            , rewardWithdrawal = Coin 0
+            }
+      where
+        genCoin :: Gen Coin
+        genCoin = oneof
+            [ pure fee
+            , genCoinRange (fee <> Coin 1) (stimes (1000 :: Int) fee)
+            ]
+        fee :: Coin
+        fee = feeForEmptySelection params <> feeForInput params
+
+    genSingleInputCoinSingleOutput :: Gen MockSelection
+    genSingleInputCoinSingleOutput = undefined
+
+    --genSingleInputBundleSingleOutput :: Gen MockSelection
+    --genSingleInputBundleSingleOutput = undefined
+
+    --genMultipleInputBundlesSingleOutput :: Gen MockSelection
+    --genMultipleInputBundlesSingleOutput = undefined
+
+    params = unMockSelectionParameters mockParams
 
 genMockSelectionHalfFull :: MockSelectionParameters -> Gen MockSelection
 genMockSelectionHalfFull mockParams =
@@ -285,6 +320,8 @@ joinMockSelections mockParams s1 s2
         , size
             = size s1 <> size s2
             & flip mockSizeSubtractSafe (sizeOfEmptySelection params)
+        , rewardWithdrawal
+            = rewardWithdrawal s1 <> rewardWithdrawal s2
         }
     params = unMockSelectionParameters mockParams
 
@@ -314,12 +351,16 @@ data MockSelectionParameters = MockSelectionParameters
         :: MockFeeForInput
     , mockFeeForOutput
         :: MockFeeForOutput
+    , mockFeeForRewardWithdrawal
+        :: MockFeeForRewardWithdrawal
     , mockSizeOfEmptySelection
         :: MockSizeOfEmptySelection
     , mockSizeOfInput
         :: MockSizeOfInput
     , mockSizeOfOutput
         :: MockSizeOfOutput
+    , mockSizeOfRewardWithdrawal
+        :: MockSizeOfRewardWithdrawal
     , mockMaximumSizeOfOutput
         :: MockMaximumSizeOfOutput
     , mockMaximumSizeOfSelection
@@ -341,6 +382,9 @@ unMockSelectionParameters m = SelectionParameters
     , feeForOutput =
         unMockFeeForOutput
             $ view #mockFeeForOutput m
+    , feeForRewardWithdrawal =
+        unMockFeeForRewardWithdrawal
+            $ view #mockFeeForRewardWithdrawal m
     , sizeOfEmptySelection =
         unMockSizeOfEmptySelection
             $ view #mockSizeOfEmptySelection m
@@ -350,6 +394,9 @@ unMockSelectionParameters m = SelectionParameters
     , sizeOfOutput =
         unMockSizeOfOutput
             $ view #mockSizeOfOutput m
+    , sizeOfRewardWithdrawal =
+        unMockSizeOfRewardWithdrawal
+            $ view #mockSizeOfRewardWithdrawal m
     , maximumSizeOfOutput =
         unMockMaximumSizeOfOutput
             $ view #mockMaximumSizeOfOutput m
@@ -366,9 +413,11 @@ genMockSelectionParameters = MockSelectionParameters
     <$> genMockFeeForEmptySelection
     <*> genMockFeeForInput
     <*> genMockFeeForOutput
+    <*> genMockFeeForRewardWithdrawal
     <*> genMockSizeOfEmptySelection
     <*> genMockSizeOfInput
     <*> genMockSizeOfOutput
+    <*> genMockSizeOfRewardWithdrawal
     <*> genMockMaximumSizeOfOutput
     <*> genMockMaximumSizeOfSelection
     <*> genMockMinimumAdaQuantityForOutput
@@ -427,6 +476,21 @@ genMockFeeForOutput = MockFeeForOutput
     <$> genCoinRange (Coin 0) (Coin 10)
     <*> genCoinRange (Coin 0) (Coin 10)
 
+--------------------------------------------------------------------------------
+-- Mock fees for reward withdrawal
+--------------------------------------------------------------------------------
+
+newtype MockFeeForRewardWithdrawal = MockFeeForRewardWithdrawal
+    { unMockFeeForRewardWithdrawal :: Coin }
+    deriving (Eq, Show)
+
+genMockFeeForRewardWithdrawal :: Gen MockFeeForRewardWithdrawal
+genMockFeeForRewardWithdrawal = MockFeeForRewardWithdrawal
+    <$> genCoinRange (Coin 0) (Coin 10)
+
+instance Arbitrary MockFeeForRewardWithdrawal where
+    arbitrary = genMockFeeForRewardWithdrawal
+
 --------------------------------------------------------------------------------
 -- Mock sizes
 --------------------------------------------------------------------------------
@@ -500,6 +564,18 @@ genMockSizeOfOutput = MockSizeOfOutput
     <$> genMockSizeRange 0 10
     <*> genMockSizeRange 0 10
 
+--------------------------------------------------------------------------------
+-- Mock sizes of reward withdrawals
+--------------------------------------------------------------------------------
+
+newtype MockSizeOfRewardWithdrawal = MockSizeOfRewardWithdrawal
+    { unMockSizeOfRewardWithdrawal :: MockSize }
+    deriving (Eq, Generic, Ord, Show)
+
+genMockSizeOfRewardWithdrawal :: Gen MockSizeOfRewardWithdrawal
+genMockSizeOfRewardWithdrawal =
+    MockSizeOfRewardWithdrawal <$> genMockSizeRange 0 10
+
 --------------------------------------------------------------------------------
 -- Mock maximum sizes of outputs
 --------------------------------------------------------------------------------
@@ -599,12 +675,16 @@ nullSelectionParameters = SelectionParameters
         Coin 0
     , feeForOutput =
         const (Coin 0)
+    , feeForRewardWithdrawal =
+        Coin 0
     , sizeOfEmptySelection =
         MockSize 0
     , sizeOfInput =
         MockSize 0
     , sizeOfOutput =
         const (MockSize 0)
+    , sizeOfRewardWithdrawal =
+        MockSize 0
     , maximumSizeOfOutput =
         SelectionOutputSizeAssessor (const SelectionOutputSizeWithinLimit)
     , maximumSizeOfSelection =