Add function runRoundRobin.

lib/core/src/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobin.hs

@@ -1,4 +1,6 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE OverloadedLabels #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
@@ -19,6 +21,10 @@ module Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
     , groupByKey
     , ungroupByKey
 
+    -- * Round-robin processing
+    , runRoundRobin
+    , runRoundRobinM
+
     ) where
 
 import Prelude
@@ -37,6 +43,8 @@ import Cardano.Wallet.Primitive.Types.TokenQuantity
     ( TokenQuantity (..) )
 import Data.Function
     ( (&) )
+import Data.Functor.Identity
+    ( Identity (..) )
 import Data.Generics.Internal.VL.Lens
     ( view )
 import Data.Generics.Labels
@@ -57,6 +65,7 @@ import Numeric.Natural
 import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
 import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
 import qualified Data.Foldable as F
+import qualified Data.List as L
 import qualified Data.List.NonEmpty as NE
 import qualified Data.Map.Strict as Map
 import qualified Data.Set as Set
@@ -283,3 +292,21 @@ groupByKey = F.foldl' acc mempty
 
 ungroupByKey :: forall k v. Map k (NonEmpty v) -> [(k, v)]
 ungroupByKey m = [(k, v) | (k, vs) <- Map.toList m, v <- NE.toList vs]
+
+--------------------------------------------------------------------------------
+-- Round-robin processing
+--------------------------------------------------------------------------------
+
+runRoundRobin :: s -> [(s -> Maybe s)] -> s
+runRoundRobin state processors =
+    runIdentity $ runRoundRobinM state $ fmap Identity <$> processors
+
+runRoundRobinM :: Monad m => s -> [(s -> m (Maybe s))] -> m s
+runRoundRobinM state processors = go state processors []
+  where
+    go !s []        [] = pure s
+    go !s []       !qs = go s (L.reverse qs) []
+    go !s (p : ps) !qs = p s >>=
+        \case
+            Nothing -> go s  ps      qs
+            Just s' -> go s' ps (p : qs)

lib/core/test/unit/Cardano/Wallet/Primitive/CoinSelection/MA/RoundRobinSpec.hs

@@ -1,4 +1,6 @@
+{-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
@@ -11,7 +13,12 @@ import Prelude
 import Algebra.PartialOrd
     ( PartialOrd (..) )
 import Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
-    ( groupByKey, makeChange, makeChangeForSurplusAssets, ungroupByKey )
+    ( groupByKey
+    , makeChange
+    , makeChangeForSurplusAssets
+    , runRoundRobin
+    , ungroupByKey
+    )
 import Cardano.Wallet.Primitive.Types.Coin
     ( Coin (..) )
 import Cardano.Wallet.Primitive.Types.TokenBundle
@@ -22,16 +29,28 @@ import Cardano.Wallet.Primitive.Types.TokenMap
     ( AssetId (..) )
 import Cardano.Wallet.Primitive.Types.TokenMap.Gen
     ( genAssetIdSmallRange, shrinkAssetIdSmallRange )
+import Cardano.Wallet.Primitive.Types.TokenPolicy
+    ( TokenName )
+import Cardano.Wallet.Primitive.Types.TokenPolicy.Gen
+    ( genTokenNameMediumRange )
 import Cardano.Wallet.Primitive.Types.TokenQuantity
     ( TokenQuantity )
 import Cardano.Wallet.Primitive.Types.TokenQuantity.Gen
     ( genTokenQuantitySmallPositive, shrinkTokenQuantitySmallPositive )
 import Control.Monad
     ( replicateM )
+import Data.Function
+    ( (&) )
 import Data.List.NonEmpty
     ( NonEmpty (..) )
 import Data.Map.Strict
     ( Map )
+import Data.Set
+    ( Set )
+import Data.Tuple
+    ( swap )
+import Data.Word
+    ( Word8 )
 import Safe
     ( tailMay )
 import Test.Hspec
@@ -45,6 +64,7 @@ import Test.QuickCheck
     , choose
     , genericShrink
     , property
+    , shrinkList
     , suchThat
     , (===)
     )
@@ -55,6 +75,7 @@ import qualified Data.Foldable as F
 import qualified Data.List as L
 import qualified Data.List.NonEmpty as NE
 import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
 
 spec :: Spec
 spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
@@ -86,6 +107,19 @@ spec = describe "Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec" $
         it "prop_ungroupByKey_groupByKey" $
             property $ prop_ungroupByKey_groupByKey @Int @Int
 
+    parallel $ describe "Round-robin processing" $ do
+
+        it "prop_runRoundRobin_identity" $
+            property $ prop_runRoundRobin_identity @Int
+        it "prop_runRoundRobin_iterationCount" $
+            property $ prop_runRoundRobin_iterationCount @TokenName @Word8
+        it "prop_runRoundRobin_iterationOrder" $
+            property $ prop_runRoundRobin_iterationOrder @TokenName @Word8
+        it "prop_runRoundRobin_generationCount" $
+            property $ prop_runRoundRobin_generationCount @TokenName @Word8
+        it "prop_runRoundRobin_generationOrder" $
+            property $ prop_runRoundRobin_generationOrder @TokenName @Word8
+
 --------------------------------------------------------------------------------
 -- Making change
 --------------------------------------------------------------------------------
@@ -183,6 +217,129 @@ prop_ungroupByKey_groupByKey
 prop_ungroupByKey_groupByKey kvs =
     fmap NE.sort kvs === fmap NE.sort (groupByKey $ ungroupByKey kvs)
 
+--------------------------------------------------------------------------------
+-- Round-robin processing
+--------------------------------------------------------------------------------
+
+data MockRoundRobinState k n = MockRoundRobinState
+    { processorLifetimes :: Map k n
+    , accumulatedEntries :: [(k, n)]
+    } deriving (Eq, Show)
+
+genMockRoundRobinState
+    :: forall k n. Ord k => Gen k -> Gen n -> Gen (MockRoundRobinState k n)
+genMockRoundRobinState genKey genLifetime = do
+    processorCount <- choose (0, 16)
+    MockRoundRobinState
+        <$> genProcessorLifetimes processorCount
+        <*> pure []
+  where
+    genProcessorLifetimes :: Int -> Gen (Map k n)
+    genProcessorLifetimes processorCount =
+        Map.fromList <$> replicateM processorCount genProcessorLifetime
+
+    genProcessorLifetime :: Gen (k, n)
+    genProcessorLifetime = (,)
+        <$> genKey
+        <*> genLifetime
+
+shrinkMockRoundRobinState
+    :: Ord k
+    => (n -> [n])
+    -> MockRoundRobinState k n
+    -> [MockRoundRobinState k n]
+shrinkMockRoundRobinState shrinkLifetime s =
+    [ s { processorLifetimes = processorLifetimes' }
+    | processorLifetimes' <- shrinkProcessorLifetimes $ processorLifetimes s
+    ]
+  where
+    shrinkProcessorLifetimes
+        = fmap Map.fromList
+        . shrinkList shrinkProcessorLifetime
+        . Map.toList
+    shrinkProcessorLifetime (k, n) = (k, ) <$> shrinkLifetime n
+
+runMockRoundRobin
+    :: forall k n. (Ord k, Integral n)
+    => MockRoundRobinState k n
+    -> MockRoundRobinState k n
+runMockRoundRobin initialState = runRoundRobin initialState processors
+  where
+    processors :: [MockRoundRobinState k n -> Maybe (MockRoundRobinState k n)]
+    processors = mkProcessor <$> Map.toList (processorLifetimes initialState)
+
+    mkProcessor
+        :: (k, n) -> MockRoundRobinState k n -> Maybe (MockRoundRobinState k n)
+    mkProcessor (k, n) s
+        | remainingLifetime <= 0 =
+            Nothing
+        | otherwise = Just $ MockRoundRobinState
+            { processorLifetimes = Map.adjust pred k (processorLifetimes s)
+            , accumulatedEntries = entry : accumulatedEntries s
+            }
+      where
+        entry :: (k, n)
+        entry = (k, n - remainingLifetime)
+
+        remainingLifetime :: n
+        remainingLifetime = Map.findWithDefault 0 k (processorLifetimes s)
+
+prop_runRoundRobin_identity
+    :: forall state. (Eq state, Show state) => state -> [()] -> Property
+prop_runRoundRobin_identity state processors =
+    runRoundRobin state (const Nothing <$ processors) === state
+
+prop_runRoundRobin_iterationCount
+    :: forall k n. (Ord k, Integral n)
+    => MockRoundRobinState k n
+    -> Property
+prop_runRoundRobin_iterationCount initialState = (===)
+    (toInteger $ length $ accumulatedEntries finalState)
+    (F.sum $ toInteger <$> processorLifetimes initialState)
+  where
+    finalState = runMockRoundRobin initialState
+
+prop_runRoundRobin_iterationOrder
+    :: forall k n. (Ord k, Show k, Integral n, Show n)
+    => MockRoundRobinState k n
+    -> Property
+prop_runRoundRobin_iterationOrder initialState =
+    sortDescending entries === entries
+  where
+    finalState = runMockRoundRobin initialState
+    entries = swap <$> accumulatedEntries finalState
+    sortDescending = L.sortBy (flip compare)
+
+prop_runRoundRobin_generationCount
+    :: forall k n. (Ord k, Show k, Integral n, Show n)
+    => MockRoundRobinState k n
+    -> Property
+prop_runRoundRobin_generationCount initialState =
+    Map.filter (> 0) (processorLifetimes initialState)
+        === generationCounts
+  where
+    finalState = runMockRoundRobin initialState
+    generationCounts :: Map k n
+    generationCounts = accumulatedEntries finalState
+        & groupByKey
+        & fmap (fromIntegral . NE.length)
+
+prop_runRoundRobin_generationOrder
+    :: forall k n. (Ord k, Integral n)
+    => MockRoundRobinState k n
+    -> Property
+prop_runRoundRobin_generationOrder initialState = property $
+    all (uncurry Set.isSubsetOf)
+        $ consecutivePairs
+        $ snd <$> Map.toDescList generations
+  where
+    finalState = runMockRoundRobin initialState
+    generations :: Map n (Set k)
+    generations = accumulatedEntries finalState
+        & fmap swap
+        & groupByKey
+        & fmap (Set.fromList . F.toList)
+
 --------------------------------------------------------------------------------
 -- Utility functions
 --------------------------------------------------------------------------------
@@ -210,6 +367,10 @@ instance Arbitrary AssetId where
 instance Arbitrary MakeChangeData where
     arbitrary = genMakeChangeData
 
+instance Arbitrary (MockRoundRobinState TokenName Word8) where
+    arbitrary = genMockRoundRobinState genTokenNameMediumRange arbitrary
+    shrink = shrinkMockRoundRobinState shrink
+
 instance Arbitrary TokenBundle where
     arbitrary = genTokenBundleSmallRange
     shrink = shrinkTokenBundleSmallRange