{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_HADDOCK prune #-}

-- |
-- Copyright: © 2018-2024 Intersect MBO
-- License: Apache-2.0
--
-- This module contains an algorithm for migrating all funds from one wallet
-- to another.
--
-- See 'selectCoins'.
--
module Cardano.CoinSelection.Algorithm.Migration
    (
      -- * Coin Selection for Migration
      selectCoins
    , BatchSize (..)
    , idealBatchSize
    ) where

import Prelude

import Cardano.CoinSelection
    ( CoinMap
    , CoinMapEntry (..)
    , CoinSelection (..)
    , CoinSelectionLimit (..)
    , coinMapFromList
    , coinMapToList
    , coinMapValue
    , sumChange
    , sumInputs
    )
import Cardano.CoinSelection.Fee
    ( DustThreshold (..)
    , Fee (..)
    , FeeBalancingPolicy (..)
    , FeeEstimator (..)
    , FeeOptions (..)
    , isDust
    )
import Control.Monad.Trans.State ( State, evalState, get, put )
import Data.List.NonEmpty ( NonEmpty ((:|)) )
import Data.Maybe ( fromMaybe )
import Data.Word ( Word16 )
import GHC.Generics ( Generic )
import Internal.Coin ( Coin, coinFromIntegral, coinToIntegral )

import qualified Internal.Coin as C

--------------------------------------------------------------------------------
-- Coin Selection for Migration
--------------------------------------------------------------------------------

-- | Creates a __series__ of coin selections that, when published as
--   transactions, will have the effect of migrating all funds from one
--   wallet to another.
--
-- Since UTxO-based blockchains typically impose limits on the sizes of
-- individual transactions, and since individual UTxO sets can contain
-- /arbitrarily/ many entries, migrating all funds from one wallet to another
-- may require the creation of /several/ transactions.
--
-- This function therefore /partitions/ the given set of inputs into multiple
-- /batches/ of up to __/b/__ inputs, where __/b/__ is specified by the given
-- 'BatchSize' parameter. (See 'idealBatchSize' for an automatic way to
-- calculate a suitable batch size.)
--
-- For each batch of inputs, this function creates a separate 'CoinSelection'
-- with the given 'inputs' /and/ a generated 'change' set, where the 'change'
-- set represents the value to be transferred to the target wallet, carefully
-- adjusted to deduct a fee in accordance with the given 'FeeOptions'
-- parameter. The set of 'outputs' for each coin selection is /purposefully/
-- left empty, as /all/ value is captured in the 'change' set.
--
-- @since 1.0.0
selectCoins
    :: forall i o . (Ord i, Ord o)
    => FeeOptions i o
        -- ^ The fee options.
    -> BatchSize
        -- ^ The maximum number of inputs to include in each selection.
    -> CoinMap i
        -- ^ The UTxO set to migrate.
    -> [CoinSelection i o]
selectCoins :: forall i o.
(Ord i, Ord o) =>
FeeOptions i o -> BatchSize -> CoinMap i -> [CoinSelection i o]
selectCoins FeeOptions i o
options (BatchSize Word16
batchSize) CoinMap i
utxo =
    State [CoinMapEntry i] [CoinSelection i o]
-> [CoinMapEntry i] -> [CoinSelection i o]
forall s a. State s a -> s -> a
evalState State [CoinMapEntry i] [CoinSelection i o]
migrate (CoinMap i -> [CoinMapEntry i]
forall a. CoinMap a -> [CoinMapEntry a]
coinMapToList CoinMap i
utxo)
  where
    FeeOptions {DustThreshold
dustThreshold :: DustThreshold
dustThreshold :: forall i o. FeeOptions i o -> DustThreshold
dustThreshold, FeeEstimator i o
feeEstimator :: FeeEstimator i o
feeEstimator :: forall i o. FeeOptions i o -> FeeEstimator i o
feeEstimator, FeeBalancingPolicy
feeBalancingPolicy :: FeeBalancingPolicy
feeBalancingPolicy :: forall i o. FeeOptions i o -> FeeBalancingPolicy
feeBalancingPolicy} = FeeOptions i o
options

    migrate :: State [CoinMapEntry i] [CoinSelection i o]
    migrate :: State [CoinMapEntry i] [CoinSelection i o]
migrate = do
        [CoinMapEntry i]
batch <- State [CoinMapEntry i] [CoinMapEntry i]
forall a. State [a] [a]
getNextBatch
        if [CoinMapEntry i] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [CoinMapEntry i]
batch then
            [CoinSelection i o] -> State [CoinMapEntry i] [CoinSelection i o]
forall a. a -> StateT [CoinMapEntry i] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
        else case CoinSelection i o -> Maybe (CoinSelection i o)
adjustForFee ([CoinMapEntry i] -> CoinSelection i o
mkCoinSelection [CoinMapEntry i]
batch) of
            Maybe (CoinSelection i o)
Nothing -> [CoinSelection i o] -> State [CoinMapEntry i] [CoinSelection i o]
forall a. a -> StateT [CoinMapEntry i] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
            Just CoinSelection i o
coinSel -> do
                [CoinSelection i o]
rest <- State [CoinMapEntry i] [CoinSelection i o]
migrate
                [CoinSelection i o] -> State [CoinMapEntry i] [CoinSelection i o]
forall a. a -> StateT [CoinMapEntry i] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CoinSelection i o
coinSelCoinSelection i o -> [CoinSelection i o] -> [CoinSelection i o]
forall a. a -> [a] -> [a]
:[CoinSelection i o]
rest)

    -- Construct a provisional 'CoinSelection' from the given selected inputs.
    -- Note that the selection may look a bit weird at first sight as it has
    -- no outputs (we are paying everything to ourselves!).
    mkCoinSelection :: [CoinMapEntry i] -> CoinSelection i o
    mkCoinSelection :: [CoinMapEntry i] -> CoinSelection i o
mkCoinSelection [CoinMapEntry i]
inputEntries = CoinSelection {CoinMap i
inputs :: CoinMap i
inputs :: CoinMap i
inputs, CoinMap o
outputs :: CoinMap o
outputs :: CoinMap o
outputs, [Coin]
change :: [Coin]
change :: [Coin]
change}
      where
        inputs :: CoinMap i
inputs = [CoinMapEntry i] -> CoinMap i
forall a. Ord a => [CoinMapEntry a] -> CoinMap a
coinMapFromList [CoinMapEntry i]
inputEntries
        outputs :: CoinMap o
outputs = CoinMap o
forall a. Monoid a => a
mempty
        change :: [Coin]
change
            | [Coin] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Coin]
nonDustInputCoins Bool -> Bool -> Bool
&& Coin
totalInputValue Coin -> Coin -> Bool
forall a. Ord a => a -> a -> Bool
>= Coin
smallestNonDustCoin =
                [Coin
smallestNonDustCoin]
            | Bool
otherwise =
                [Coin]
nonDustInputCoins
        nonDustInputCoins :: [Coin]
nonDustInputCoins = (Coin -> Bool) -> [Coin] -> [Coin]
forall a. (a -> Bool) -> [a] -> [a]
filter
            (Bool -> Bool
not (Bool -> Bool) -> (Coin -> Bool) -> Coin -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DustThreshold -> Coin -> Bool
isDust DustThreshold
dustThreshold)
            (CoinMapEntry i -> Coin
forall a. CoinMapEntry a -> Coin
entryValue (CoinMapEntry i -> Coin) -> [CoinMapEntry i] -> [Coin]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CoinMapEntry i]
inputEntries)
        smallestNonDustCoin :: Coin
smallestNonDustCoin = Coin -> Coin
C.succ (Coin -> Coin) -> Coin -> Coin
forall a b. (a -> b) -> a -> b
$ DustThreshold -> Coin
unDustThreshold DustThreshold
dustThreshold
        totalInputValue :: Coin
totalInputValue = CoinMap i -> Coin
forall a. CoinMap a -> Coin
coinMapValue CoinMap i
inputs

    -- | Attempt to balance the coin selection by reducing or increasing the
    -- change values based on the computed fees.
    adjustForFee :: CoinSelection i o -> Maybe (CoinSelection i o)
    adjustForFee :: CoinSelection i o -> Maybe (CoinSelection i o)
adjustForFee !CoinSelection i o
coinSel = case CoinSelection i o -> [Coin]
forall i o. CoinSelection i o -> [Coin]
change CoinSelection i o
coinSel of
        -- If there's no change, nothing to adjust
        [] -> Maybe (CoinSelection i o)
forall a. Maybe a
Nothing

        -- No difference between required and computed, we're done
        (Coin
_ : [Coin]
_) | Integer
diff Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 -> CoinSelection i o -> Maybe (CoinSelection i o)
forall a. a -> Maybe a
Just CoinSelection i o
coinSel

        -- Otherwise, we have 2 cases:
        --
        -- 1/ diff < 0
        -- We aren't giving enough as fee, so we need to reduce one output.
        --
        -- 2/ diff > 0
        -- We have some surplus so we add it to an arbitrary output
        --
        -- If both cases we can simply modify one output by adding `diff`, the
        -- sign of `diff` making for the right modification.
        -- We then recursively call ourselves for this might reduce the number
        -- of outputs and change the fee.
        (Coin
c : [Coin]
cs) -> do
            let coinSel' :: CoinSelection i o
coinSel' = CoinSelection i o
coinSel
                    { change = modifyFirst (c :| cs) (applyDiff diff) }
            let costOfSurplus :: Integer
costOfSurplus
                    = Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral
                    (Natural -> Integer) -> Natural -> Integer
forall a b. (a -> b) -> a -> b
$ Coin -> Natural
C.coinToNatural
                    (Coin -> Natural) -> Coin -> Natural
forall a b. (a -> b) -> a -> b
$ Coin -> Coin -> Coin
C.distance
                        (Fee -> Coin
unFee (Fee -> Coin) -> Fee -> Coin
forall a b. (a -> b) -> a -> b
$ FeeEstimator i o -> CoinSelection i o -> Fee
forall i o. FeeEstimator i o -> CoinSelection i o -> Fee
estimateFee FeeEstimator i o
feeEstimator CoinSelection i o
coinSel')
                        (Fee -> Coin
unFee (Fee -> Coin) -> Fee -> Coin
forall a b. (a -> b) -> a -> b
$ FeeEstimator i o -> CoinSelection i o -> Fee
forall i o. FeeEstimator i o -> CoinSelection i o -> Fee
estimateFee FeeEstimator i o
feeEstimator CoinSelection i o
coinSel )
            if
                -- Adding the change costs less than not having it, so it's
                -- worth trying.
                | Integer
costOfSurplus Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
actualFee ->
                    CoinSelection i o -> Maybe (CoinSelection i o)
adjustForFee CoinSelection i o
coinSel'

                -- Adding the change costs more than not having it, If we don't
                -- require strict balancing, we can leave the selection as-is.
                | FeeBalancingPolicy
feeBalancingPolicy FeeBalancingPolicy -> FeeBalancingPolicy -> Bool
forall a. Eq a => a -> a -> Bool
== FeeBalancingPolicy
RequireMinimalFee ->
                    CoinSelection i o -> Maybe (CoinSelection i o)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CoinSelection i o
coinSel

                -- Adding the change costs more than not having it. So,
                -- depending on our balancing policy, we may stop the balancing
                -- right here, or, if we must balance the selection discard the
                -- whole selection: it can't be balanced with this algorithm.
                --
                -- Note that this last extreme case is reached when using an
                -- unstable fee policy (where values of outputs can influence
                -- the policy) AND, require transactions to be 100% balanced.
                -- This is a silly thing to do.
                | Bool
otherwise ->
                    Maybe (CoinSelection i o)
forall a. Maybe a
Nothing
      where
        applyDiff :: Integer -> Coin -> Coin
        applyDiff :: Integer -> Coin -> Coin
applyDiff Integer
i Coin
c
            = Coin -> Maybe Coin -> Coin
forall a. a -> Maybe a -> a
fromMaybe Coin
C.zero
            (Maybe Coin -> Coin) -> Maybe Coin -> Coin
forall a b. (a -> b) -> a -> b
$ Integer -> Maybe Coin
forall i. Integral i => i -> Maybe Coin
coinFromIntegral (Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Coin -> Integer
forall i. Integral i => Coin -> i
coinToIntegral Coin
c)

        diff :: Integer
        diff :: Integer
diff = Integer
actualFee Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
requiredFee
          where
            requiredFee :: Integer
requiredFee
                = Coin -> Integer
forall i. Integral i => Coin -> i
coinToIntegral (Coin -> Integer) -> Coin -> Integer
forall a b. (a -> b) -> a -> b
$ Fee -> Coin
unFee
                (Fee -> Coin) -> Fee -> Coin
forall a b. (a -> b) -> a -> b
$ FeeEstimator i o -> CoinSelection i o -> Fee
forall i o. FeeEstimator i o -> CoinSelection i o -> Fee
estimateFee FeeEstimator i o
feeEstimator CoinSelection i o
coinSel

        actualFee :: Integer
        actualFee :: Integer
actualFee
            = Coin -> Integer
forall i. Integral i => Coin -> i
coinToIntegral (CoinSelection i o -> Coin
forall i o. CoinSelection i o -> Coin
sumInputs CoinSelection i o
coinSel)
            Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Coin -> Integer
forall i. Integral i => Coin -> i
coinToIntegral (CoinSelection i o -> Coin
forall i o. CoinSelection i o -> Coin
sumChange CoinSelection i o
coinSel)

    -- | Apply the given function to the first coin of the list. If the
    -- operation makes the 'Coin' smaller than the dust threshold, the coin is
    -- discarded.
    modifyFirst :: NonEmpty Coin -> (Coin -> Coin) -> [Coin]
    modifyFirst :: NonEmpty Coin -> (Coin -> Coin) -> [Coin]
modifyFirst (Coin
c :| [Coin]
cs) Coin -> Coin
op
        | Coin
c' Coin -> Coin -> Bool
forall a. Ord a => a -> a -> Bool
<= Coin
threshold = [Coin]
cs
        | Bool
otherwise = Coin
c' Coin -> [Coin] -> [Coin]
forall a. a -> [a] -> [a]
: [Coin]
cs
      where
        c' :: Coin
c' = Coin -> Coin
op Coin
c
        threshold :: Coin
threshold = DustThreshold -> Coin
unDustThreshold DustThreshold
dustThreshold

    getNextBatch :: State [a] [a]
    getNextBatch :: forall a. State [a] [a]
getNextBatch = do
        [a]
xs <- State [a] [a]
forall (m :: * -> *) s. Monad m => StateT s m s
get
        let ([a]
batch, [a]
rest) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt (Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
batchSize) [a]
xs
        [a] -> StateT [a] Identity ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put [a]
rest
        [a] -> State [a] [a]
forall a. a -> StateT [a] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
batch

-- | An upper limit for the number of 'inputs' to include in each coin selection
--   generated by 'selectCoins'.
--
-- @since 1.0.0
newtype BatchSize = BatchSize Word16
    deriving (BatchSize -> BatchSize -> Bool
(BatchSize -> BatchSize -> Bool)
-> (BatchSize -> BatchSize -> Bool) -> Eq BatchSize
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BatchSize -> BatchSize -> Bool
== :: BatchSize -> BatchSize -> Bool
$c/= :: BatchSize -> BatchSize -> Bool
/= :: BatchSize -> BatchSize -> Bool
Eq, (forall x. BatchSize -> Rep BatchSize x)
-> (forall x. Rep BatchSize x -> BatchSize) -> Generic BatchSize
forall x. Rep BatchSize x -> BatchSize
forall x. BatchSize -> Rep BatchSize x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. BatchSize -> Rep BatchSize x
from :: forall x. BatchSize -> Rep BatchSize x
$cto :: forall x. Rep BatchSize x -> BatchSize
to :: forall x. Rep BatchSize x -> BatchSize
Generic, Eq BatchSize
Eq BatchSize =>
(BatchSize -> BatchSize -> Ordering)
-> (BatchSize -> BatchSize -> Bool)
-> (BatchSize -> BatchSize -> Bool)
-> (BatchSize -> BatchSize -> Bool)
-> (BatchSize -> BatchSize -> Bool)
-> (BatchSize -> BatchSize -> BatchSize)
-> (BatchSize -> BatchSize -> BatchSize)
-> Ord BatchSize
BatchSize -> BatchSize -> Bool
BatchSize -> BatchSize -> Ordering
BatchSize -> BatchSize -> BatchSize
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: BatchSize -> BatchSize -> Ordering
compare :: BatchSize -> BatchSize -> Ordering
$c< :: BatchSize -> BatchSize -> Bool
< :: BatchSize -> BatchSize -> Bool
$c<= :: BatchSize -> BatchSize -> Bool
<= :: BatchSize -> BatchSize -> Bool
$c> :: BatchSize -> BatchSize -> Bool
> :: BatchSize -> BatchSize -> Bool
$c>= :: BatchSize -> BatchSize -> Bool
>= :: BatchSize -> BatchSize -> Bool
$cmax :: BatchSize -> BatchSize -> BatchSize
max :: BatchSize -> BatchSize -> BatchSize
$cmin :: BatchSize -> BatchSize -> BatchSize
min :: BatchSize -> BatchSize -> BatchSize
Ord, Int -> BatchSize -> ShowS
[BatchSize] -> ShowS
BatchSize -> String
(Int -> BatchSize -> ShowS)
-> (BatchSize -> String)
-> ([BatchSize] -> ShowS)
-> Show BatchSize
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BatchSize -> ShowS
showsPrec :: Int -> BatchSize -> ShowS
$cshow :: BatchSize -> String
show :: BatchSize -> String
$cshowList :: [BatchSize] -> ShowS
showList :: [BatchSize] -> ShowS
Show)

-- | Calculate an ideal batch size based on the given coin selection limit.
--
-- @since 1.0.0
idealBatchSize :: CoinSelectionLimit -> BatchSize
idealBatchSize :: CoinSelectionLimit -> BatchSize
idealBatchSize CoinSelectionLimit
coinselOpts = Word16 -> BatchSize
BatchSize (Word16 -> BatchSize) -> Word16 -> BatchSize
forall a b. (a -> b) -> a -> b
$ Word16 -> Word16
fixPoint Word16
1
  where
    fixPoint :: Word16 -> Word16
    fixPoint :: Word16 -> Word16
fixPoint !Word16
n
        | Word16 -> Word16
maxN Word16
n Word16 -> Word16 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word16
n = Word16
n
        | Word16
n Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
forall a. Bounded a => a
maxBound = Word16
n
        | Bool
otherwise = Word16 -> Word16
fixPoint (Word16
n Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word16
1)
      where
        maxN :: Word16 -> Word16
        maxN :: Word16 -> Word16
maxN = CoinSelectionLimit -> Word16 -> Word16
calculateLimit CoinSelectionLimit
coinselOpts