Multi-Asset Migration Algorithm

lib/core/cardano-wallet-core.cabal

@@ -175,6 +175,9 @@ library
       Cardano.Wallet.Primitive.AddressDiscovery.SharedState
       Cardano.Wallet.Primitive.SyncProgress
       Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobin
+      Cardano.Wallet.Primitive.Migration
+      Cardano.Wallet.Primitive.Migration.Planning
+      Cardano.Wallet.Primitive.Migration.Selection
       Cardano.Wallet.Primitive.Model
       Cardano.Wallet.Primitive.Types
       Cardano.Wallet.Primitive.Types.Address
@@ -356,6 +359,8 @@ test-suite unit
       Cardano.Wallet.Primitive.AddressDiscovery.SharedStateSpec
       Cardano.Wallet.Primitive.AddressDiscoverySpec
       Cardano.Wallet.Primitive.CoinSelection.MA.RoundRobinSpec
+      Cardano.Wallet.Primitive.Migration.PlanningSpec
+      Cardano.Wallet.Primitive.Migration.SelectionSpec
       Cardano.Wallet.Primitive.ModelSpec
       Cardano.Wallet.Primitive.Slotting.Legacy
       Cardano.Wallet.Primitive.SlottingSpec

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

@@ -0,0 +1,74 @@
+{-# LANGUAGE DuplicateRecordFields #-}
+{-# LANGUAGE OverloadedLabels #-}
+
+-- |
+-- Copyright: © 2021 IOHK
+-- License: Apache-2.0
+--
+-- This module provides a public API for planning wallet migrations.
+--
+-- Use 'createPlan' to create a migration plan.
+--
+module Cardano.Wallet.Primitive.Migration
+    (
+    -- * Creating a migration plan
+      createPlan
+    , MigrationPlan (..)
+    , RewardWithdrawal (..)
+    , Selection (..)
+    , TxSize (..)
+
+    ) where
+
+import Prelude
+
+import Cardano.Wallet.Primitive.Migration.Selection
+    ( RewardWithdrawal (..), Selection (..), TxSize (..) )
+import Cardano.Wallet.Primitive.Types.Coin
+    ( Coin )
+import Cardano.Wallet.Primitive.Types.Tx
+    ( TxConstraints (..), TxIn, TxOut )
+import Cardano.Wallet.Primitive.Types.UTxO
+    ( UTxO )
+import Data.Generics.Internal.VL.Lens
+    ( view )
+import Data.Generics.Labels
+    ()
+
+import qualified Cardano.Wallet.Primitive.Migration.Planning as Planning
+
+-- | Represents a plan for migrating a 'UTxO' set.
+--
+-- See 'createPlan' to create a migration plan.
+--
+data MigrationPlan s = MigrationPlan
+    { selections :: ![Selection (TxIn, TxOut) s]
+      -- ^ A list of generated selections: each selection is the basis for a
+      -- single transaction.
+    , unselected :: !UTxO
+      -- ^ The portion of the UTxO that was not selected.
+    , totalFee :: !Coin
+      -- ^ The total fee payable: equal to the sum of the fees of the
+      -- individual selections.
+    }
+    deriving (Eq, Show)
+
+-- | Creates a migration plan for the given UTxO set and reward withdrawal
+--   amount.
+--
+-- See 'MigrationPlan'.
+--
+createPlan
+    :: TxSize s
+    => TxConstraints s
+    -> UTxO
+    -> RewardWithdrawal
+    -> MigrationPlan s
+createPlan constraints utxo reward = MigrationPlan
+    { selections = view #selections plan
+    , unselected = Planning.uncategorizeUTxO (view #unselected plan)
+    , totalFee = view #totalFee plan
+    }
+  where
+    categorizedUTxO = Planning.categorizeUTxO constraints utxo
+    plan = Planning.createPlan constraints categorizedUTxO reward

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

@@ -0,0 +1,271 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DuplicateRecordFields #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE OverloadedLabels #-}
+{-# LANGUAGE OverloadedLists #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+-- |
+-- Copyright: © 2021 IOHK
+-- License: Apache-2.0
+--
+-- This module contains an algorithm for planning migrations at a high level.
+--
+-- It determines how to partition the UTxO set into entries of different types,
+-- and in which order to add entries to selections, in order to maximize the
+-- number of entries that can be successfully migrated.
+--
+-- Use 'createPlan' to create a migration plan.
+--
+module Cardano.Wallet.Primitive.Migration.Planning
+    (
+    -- * Migration planning
+      createPlan
+    , MigrationPlan (..)
+
+    -- * UTxO entry categorization
+    , CategorizedUTxO (..)
+    , UTxOEntryCategory (..)
+    , categorizeUTxO
+    , categorizeUTxOEntries
+    , categorizeUTxOEntry
+    , uncategorizeUTxO
+    , uncategorizeUTxOEntries
+
+    ) where
+
+import Prelude
+
+import Cardano.Wallet.Primitive.Migration.Selection
+    ( RewardWithdrawal (..), Selection (..), SelectionError (..), TxSize (..) )
+import Cardano.Wallet.Primitive.Types.Coin
+    ( Coin (..) )
+import Cardano.Wallet.Primitive.Types.TokenBundle
+    ( TokenBundle (..) )
+import Cardano.Wallet.Primitive.Types.Tx
+    ( TxConstraints (..), TxIn, TxOut )
+import Cardano.Wallet.Primitive.Types.UTxO
+    ( UTxO (..) )
+import Data.Either
+    ( isRight )
+import Data.Functor
+    ( (<&>) )
+import Data.Generics.Internal.VL.Lens
+    ( view )
+import Data.Generics.Labels
+    ()
+import GHC.Generics
+    ( Generic )
+
+import qualified Cardano.Wallet.Primitive.Migration.Selection as Selection
+import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
+import qualified Data.Foldable as F
+import qualified Data.List as L
+import qualified Data.Map.Strict as Map
+
+--------------------------------------------------------------------------------
+-- Migration planning
+--------------------------------------------------------------------------------
+
+data MigrationPlan i s = MigrationPlan
+    { selections :: ![Selection i s]
+    , unselected :: !(CategorizedUTxO i)
+    , totalFee :: !Coin
+    }
+    deriving (Eq, Generic, Show)
+
+createPlan
+    :: TxSize s
+    => TxConstraints s
+    -> CategorizedUTxO i
+    -> RewardWithdrawal
+    -> MigrationPlan i s
+createPlan constraints =
+    run []
+  where
+    run !selections !utxo !reward =
+        case createSelection constraints utxo reward of
+            Just (utxo', selection) ->
+                run (selection : selections) utxo' (RewardWithdrawal $ Coin 0)
+            Nothing -> MigrationPlan
+                { selections
+                , unselected = utxo
+                , totalFee = F.foldMap (view #fee) selections
+                }
+
+createSelection
+    :: TxSize s
+    => TxConstraints s
+    -> CategorizedUTxO i
+    -> RewardWithdrawal
+    -> Maybe (CategorizedUTxO i, Selection i s)
+createSelection constraints utxo rewardWithdrawal =
+    initializeSelection constraints utxo rewardWithdrawal
+    <&> extendSelection constraints
+
+initializeSelection
+    :: forall i s. TxSize s
+    => TxConstraints s
+    -> CategorizedUTxO i
+    -> RewardWithdrawal
+    -> Maybe (CategorizedUTxO i, Selection i s)
+initializeSelection constraints utxoAtStart reward =
+    initializeWith =<< utxoAtStart `select` Supporter
+  where
+    initializeWith (entry, utxo) =
+        case Selection.create constraints reward [entry] of
+            Right selection -> Just (utxo, selection)
+            Left _ -> Nothing
+
+extendSelection
+    :: TxSize s
+    => TxConstraints s
+    -> (CategorizedUTxO i, Selection i s)
+    -> (CategorizedUTxO i, Selection i s)
+extendSelection constraints = extendWithFreerider
+  where
+    extendWithFreerider (!utxo, !selection) =
+        case extendWith Freerider constraints (utxo, selection) of
+            Right (utxo', selection') ->
+                extendWithFreerider (utxo', selection')
+            Left ExtendSelectionAdaInsufficient ->
+                extendWithSupporter (utxo, selection)
+            Left ExtendSelectionEntriesExhausted ->
+                extendWithSupporter (utxo, selection)
+            Left ExtendSelectionFull ->
+                (utxo, selection)
+
+    extendWithSupporter (!utxo, !selection) =
+        case extendWith Supporter constraints (utxo, selection) of
+            Right (utxo', selection') ->
+                extendWithFreerider (utxo', selection')
+            Left ExtendSelectionAdaInsufficient ->
+                (utxo, selection)
+            Left ExtendSelectionEntriesExhausted ->
+                (utxo, selection)
+            Left ExtendSelectionFull ->
+                (utxo, selection)
+
+data ExtendSelectionError
+    = ExtendSelectionAdaInsufficient
+    | ExtendSelectionEntriesExhausted
+    | ExtendSelectionFull
+
+extendWith
+    :: TxSize s
+    => UTxOEntryCategory
+    -> TxConstraints s
+    -> (CategorizedUTxO i, Selection i s)
+    -> Either ExtendSelectionError (CategorizedUTxO i, Selection i s)
+extendWith category constraints (utxo, selection) =
+    case utxo `select` category of
+        Just (entry, utxo') ->
+            case Selection.extend constraints selection entry of
+                Right selection' ->
+                    Right (utxo', selection')
+                Left SelectionAdaInsufficient ->
+                    Left ExtendSelectionAdaInsufficient
+                Left SelectionFull {} ->
+                    Left ExtendSelectionFull
+        Nothing ->
+            Left ExtendSelectionEntriesExhausted
+
+select
+    :: CategorizedUTxO i
+    -> UTxOEntryCategory
+    -> Maybe ((i, TokenBundle), CategorizedUTxO i)
+select utxo = \case
+    Supporter -> selectSupporter
+    Freerider -> selectFreerider
+    Ignorable -> selectIgnorable
+  where
+    selectSupporter = case supporters utxo of
+        entry : remaining -> Just (entry, utxo {supporters = remaining})
+        [] -> Nothing
+    selectFreerider = case freeriders utxo of
+        entry : remaining -> Just (entry, utxo {freeriders = remaining})
+        [] ->  Nothing
+    selectIgnorable =
+        -- We never select an entry that should be ignored:
+        Nothing
+
+--------------------------------------------------------------------------------
+-- Categorization of UTxO entries
+--------------------------------------------------------------------------------
+
+data UTxOEntryCategory
+    = Supporter
+    -- ^ A coin or bundle that is capable of paying for its own marginal fee
+    -- and the base transaction fee.
+    | Freerider
+    -- ^ A coin or bundle that is not capable of paying for itself.
+    | Ignorable
+    -- ^ A coin that should not be added to a selection, because its value is
+    -- lower than the marginal fee for an input.
+    deriving (Eq, Show)
+
+data CategorizedUTxO i = CategorizedUTxO
+    { supporters :: ![(i, TokenBundle)]
+    , freeriders :: ![(i, TokenBundle)]
+    , ignorables :: ![(i, TokenBundle)]
+    }
+    deriving (Eq, Show)
+
+categorizeUTxO
+    :: TxSize s
+    => TxConstraints s
+    -> UTxO
+    -> CategorizedUTxO (TxIn, TxOut)
+categorizeUTxO constraints (UTxO u) = categorizeUTxOEntries constraints $
+    (\(i, o) -> ((i, o), view #tokens o)) <$> Map.toList u
+
+categorizeUTxOEntries
+    :: forall i s. TxSize s
+    => TxConstraints s
+    -> [(i, TokenBundle)]
+    -> CategorizedUTxO i
+categorizeUTxOEntries constraints uncategorizedEntries = CategorizedUTxO
+    { supporters = entriesMatching Supporter
+    , freeriders = entriesMatching Freerider
+    , ignorables = entriesMatching Ignorable
+    }
+  where
+    categorizedEntries :: [(i, (TokenBundle, UTxOEntryCategory))]
+    categorizedEntries = uncategorizedEntries
+        <&> (\(i, b) -> (i, (b, categorizeUTxOEntry constraints b)))
+
+    entriesMatching :: UTxOEntryCategory -> [(i, TokenBundle)]
+    entriesMatching category =
+        fmap fst <$> L.filter ((== category) . snd . snd) categorizedEntries
+
+categorizeUTxOEntry
+    :: TxSize s
+    => TxConstraints s
+    -> TokenBundle
+    -> UTxOEntryCategory
+categorizeUTxOEntry constraints b
+    | Just c <- TokenBundle.toCoin b, coinIsIgnorable c =
+        Ignorable
+    | bundleIsSupporter =
+        Supporter
+    | otherwise =
+        Freerider
+  where
+    bundleIsSupporter :: Bool
+    bundleIsSupporter = isRight $
+        Selection.create constraints (RewardWithdrawal $ Coin 0) [((), b)]
+
+    coinIsIgnorable :: Coin -> Bool
+    coinIsIgnorable c = c <= txInputCost constraints
+
+uncategorizeUTxO :: CategorizedUTxO (TxIn, TxOut) -> UTxO
+uncategorizeUTxO = UTxO . Map.fromList . fmap fst . uncategorizeUTxOEntries
+
+uncategorizeUTxOEntries :: CategorizedUTxO i -> [(i, TokenBundle)]
+uncategorizeUTxOEntries utxo = mconcat
+    [ supporters utxo
+    , freeriders utxo
+    , ignorables utxo
+    ]