refactor readRewardAccount and use it also for shared wallets

lib/wallet/api/http/Cardano/Wallet/Api/Http/Shelley/Server.hs

@@ -2584,7 +2584,7 @@ constructTransaction api argGenChange knownPools poolStatus apiWalletId body = d
 
         apiDecoded <- decodeTransaction @_ @n api apiWalletId balancedTx
 
-        (_, _, rewardPath) <- handler $ W.readRewardAccount @n db
+        (_, _, rewardPath) <- handler $ W.readRewardAccount @s db
 
         let deposits = case txDelegationAction transactionCtx2 of
                 Just (JoinRegisteringKey _poolId) -> [W.stakeKeyDeposit pp]
@@ -2954,26 +2954,26 @@ constructSharedTransaction
                 let refunds = case optionalDelegationAction of
                         Just Quit -> [W.stakeKeyDeposit pp]
                         _ -> []
-                rewardAccountM <- handler $ W.readSharedRewardAccount @n db
-                delCerts <- case optionalDelegationAction of
+                delCertsWithPath <- case optionalDelegationAction of
                     Nothing -> pure Nothing
                     Just action -> do
-                        --at this moment we are sure reward account is present
-                        --if not ErrConstructTxDelegationInvalid would be thrown already
-                        pure $ Just (action, snd $ fromJust rewardAccountM)
+                        (_, _, path) <-
+                            handler $ W.readRewardAccount @((SharedState n SharedKey)) db
+                        pure $ Just (action, path)
 
                 pure $ ApiConstructTransaction
                     { transaction = balancedTx
                     , coinSelection =
                         mkApiCoinSelection deposits refunds
-                        delCerts md (unsignedTx outs apiDecoded rewardAccountM)
+                        delCertsWithPath md
+                        (unsignedTx outs apiDecoded (snd <$> delCertsWithPath))
                     , fee = apiDecoded ^. #fee
                     }
   where
     ti :: TimeInterpreter (ExceptT PastHorizonException IO)
     ti = timeInterpreter (api ^. networkLayer)
 
-    unsignedTx initialOuts decodedTx rewardAccountM = UnsignedTx
+    unsignedTx initialOuts decodedTx pathM = UnsignedTx
         { unsignedCollateral =
             mapMaybe toUnsignedTxInp (decodedTx ^. #collateral)
         , unsignedInputs =
@@ -2984,9 +2984,9 @@ constructSharedTransaction
         , unsignedChange =
             drop (length initialOuts)
                 $ map toUnsignedTxChange (decodedTx ^. #outputs)
-        , unsignedWithdrawals = case rewardAccountM of
+        , unsignedWithdrawals = case pathM of
                 Nothing -> []
-                Just (_, path) ->
+                Just path ->
                     mapMaybe (toUsignedTxWdrl path) (decodedTx ^. #withdrawals)
         }
 
@@ -3734,7 +3734,7 @@ listStakeKeys lookupStakeRef ctx@ApiLayer{..} (ApiT wid) =
         (wal, (_, delegation) ,pending) <- W.readWallet @_ @s wrk
         let utxo = availableUTxO @s pending wal
         let takeFst (a,_,_) = a
-        ourAccount <- takeFst <$> liftIO (W.readRewardAccount @n db)
+        ourAccount <- takeFst <$> liftIO (W.readRewardAccount @s db)
         ourApiDelegation <- liftIO $ toApiWalletDelegation delegation
             (unsafeExtendSafeZone (timeInterpreter $ ctx ^. networkLayer))
         let ourKeys = [(ourAccount, 0, ourApiDelegation)]

lib/wallet/src/Cardano/Wallet.hs

@@ -84,7 +84,6 @@ module Cardano.Wallet
     , fetchRewardBalance
     , manageRewardBalance
     , manageSharedRewardBalance
-    , readSharedRewardAccount
     , rollbackBlocks
     , checkWalletIntegrity
     , mkExternalWithdrawal
@@ -347,12 +346,16 @@ import Cardano.Wallet.DB.WalletState
     , getSlot
     )
 import Cardano.Wallet.Flavor
-    ( CredFromOf
+    ( AllFlavors
+    , CredFromOf
     , Excluding
+    , FlavorOf
+    , Including
     , KeyFlavorS (..)
     , KeyOf
     , WalletFlavor (..)
     , WalletFlavorS (..)
+    , WalletFlavors (..)
     , keyFlavorFromState
     , keyOfWallet
     )
@@ -586,7 +589,7 @@ import Data.List
 import Data.List.NonEmpty
     ( NonEmpty (..) )
 import Data.Maybe
-    ( fromJust, fromMaybe, isJust, isNothing, mapMaybe, maybeToList )
+    ( fromJust, fromMaybe, isJust, mapMaybe, maybeToList )
 import Data.Quantity
     ( Quantity (..) )
 import Data.Set
@@ -1282,10 +1285,8 @@ mkSelfWithdrawalShared
     -> DBLayer IO (SharedState n SharedKey)
     -> IO Withdrawal
 mkSelfWithdrawalShared netLayer txWitnessTag era db = do
-    rewardAccountM <- readSharedRewardAccount db
-    when (isNothing rewardAccountM) $
-        throwIO $ ExceptionReadRewardAccount ErrReadRewardAccountMissing
-    let (rewardAccount, derivationPath) = fromJust rewardAccountM
+    (rewardAccount, _, derivationPath) <-
+        readRewardAccount @(SharedState n SharedKey) db
     balance <- getCachedRewardAccountBalance netLayer rewardAccount
     pp <- currentProtocolParameters netLayer
     return $ case checkRewardIsWorthTxCost txWitnessTag pp era balance of
@@ -1317,25 +1318,25 @@ checkRewardIsWorthTxCost txWitnessTag pp era balance = do
         dummyPath = DerivationIndex 0 :| []
 
 readRewardAccount
-    :: forall n
-     . DBLayer IO (SeqState n ShelleyKey)
-    -> IO (RewardAccount, XPub, NonEmpty DerivationIndex)
+    :: forall s.
+    ( WalletFlavor s
+    , Including AllFlavors '[ 'ShelleyF, 'SharedF] (FlavorOf s)
+    )
+    => DBLayer IO s
+    -> IO (RewardAccount, Maybe XPub, NonEmpty DerivationIndex)
 readRewardAccount db = do
-    walletState <- getState <$> readWalletCheckpoint db
-    let xpub = Seq.rewardAccountKey walletState
-    let path = stakeDerivationPath $ Seq.derivationPrefix walletState
-    pure (toRewardAccount xpub, getRawKey ShelleyKeyS xpub, path)
-
-readSharedRewardAccount
-    :: forall n
-     . DBLayer IO (SharedState n SharedKey)
-    -> IO (Maybe (RewardAccount, NonEmpty DerivationIndex))
-readSharedRewardAccount db = do
-    walletState <- getState <$> readWalletCheckpoint db
-    let path = stakeDerivationPath $ Shared.derivationPrefix walletState
-    case Shared.rewardAccountKey walletState of
-        Just rewardAcct -> pure $ Just (rewardAcct, path)
-        Nothing -> pure Nothing
+    case walletFlavor @s of
+        ShelleyWallet -> do
+            walletState <- getState <$> readWalletCheckpoint db
+            let xpub = Seq.rewardAccountKey walletState
+            let path = stakeDerivationPath $ Seq.derivationPrefix walletState
+            pure (toRewardAccount xpub, Just $ getRawKey ShelleyKeyS xpub, path)
+        SharedWallet -> do
+            walletState <- getState <$> readWalletCheckpoint db
+            let path = stakeDerivationPath $ Shared.derivationPrefix walletState
+            case Shared.rewardAccountKey walletState of
+                Just rewardAcct -> pure $ (rewardAcct, Nothing, path)
+                Nothing -> throwIO $ ExceptionReadRewardAccount ErrReadRewardAccountMissing
 
 readWalletCheckpoint
     :: DBLayer IO s -> IO (Wallet s)
@@ -1365,7 +1366,7 @@ shelleyOnlyReadRewardAccount
      . WalletFlavor s
     => DBLayer IO s
     -> ExceptT ErrReadRewardAccount IO
-        (RewardAccount, XPub, NonEmpty DerivationIndex)
+        (RewardAccount, Maybe XPub, NonEmpty DerivationIndex)
 shelleyOnlyReadRewardAccount db = do
     case walletFlavor @s of
         ShelleyWallet -> lift $ readRewardAccount db
@@ -1439,11 +1440,9 @@ manageSharedRewardBalance tr' netLayer db = do
     watchNodeTip netLayer $ \bh -> do
          traceWith tr $ MsgRewardBalanceQuery bh
          query <- runExceptT $ do
-            acctM <- lift $ readSharedRewardAccount db
-            case acctM of
-                Nothing -> throwE ErrFetchRewardsMissingRewardAccount
-                Just acct ->
-                    liftIO $ getCachedRewardAccountBalance netLayer (fst acct)
+            (acct, _, _) <-
+                lift $ readRewardAccount @(SharedState n SharedKey) db
+            liftIO $ getCachedRewardAccountBalance netLayer acct
          handleRewardAccountQuery tr db query
     traceWith tr MsgRewardBalanceExited
   where
@@ -2310,7 +2309,7 @@ constructTransaction
     -> ExceptT ErrConstructTx IO (Cardano.TxBody era)
 constructTransaction txLayer db txCtx preSel = do
     (_, xpub, _) <- lift $ readRewardAccount db
-    mkUnsignedTransaction txLayer (Left xpub) txCtx (Left preSel)
+    mkUnsignedTransaction txLayer (Left $ fromJust xpub) txCtx (Left preSel)
         & withExceptT ErrConstructTxBody . except
 
 constructUnbalancedSharedTransaction

lib/wallet/src/Cardano/Wallet/Flavor.hs

@@ -27,6 +27,8 @@ module Cardano.Wallet.Flavor
     , FlavorOf
     , WalletFlavors (..)
     , Excluding
+    , Including
+    , AllFlavors
     , shelleyOrShared
     , notByronKey
     , IncludingStates