Re-arrange the restoration task so most of the action happens inside whenRight.

Author: matt-noonan

wallet-new/src/Cardano/Wallet/Kernel/Restore.hs

@@ -37,11 +37,12 @@ import           Cardano.Wallet.Kernel.NodeStateAdaptor (LockContext (..),
 import           Cardano.Wallet.Kernel.PrefilterTx (PrefilteredBlock, WalletKey,
                      toHdAddressId)
 import           Cardano.Wallet.Kernel.Types (WalletId)
+import           Cardano.Wallet.Kernel.Util.Core (utxoBalance)
 
 import           Pos.Chain.Block (BlockHeader, Blund, HeaderHash, blockHeader,
                      headerHash, mainBlockSlot)
-import           Pos.Chain.Txp (Utxo, getTotalCoinsInUtxo)
-import           Pos.Core (Address, Coin, SlotId, flattenSlotIdExplicit)
+import           Pos.Chain.Txp (Utxo)
+import           Pos.Core (Address, Coin, SlotId, flattenSlotIdExplicit, mkCoin)
 import           Pos.Core.Txp (toaOut, txOutAddress)
 import           Pos.DB (MonadDBRead)
 import           Pos.DB.Block (getBlund, getFirstGenesisBlockHash,
@@ -94,7 +95,7 @@ restoreWallet wallet hdRoot (wId, wdc) prefilter = do
 
     -- Return the wallet's current balance.
     let myUtxo  = foldr (\(u,_,_) -> M.union u) M.empty mine
-    return (getTotalCoinsInUtxo myUtxo)
+    return (mkCoin . fromIntegral . utxoBalance $ myUtxo)
 
   where
     txoAddr = txOutAddress . toaOut . snd
@@ -140,18 +141,19 @@ restoreWalletHistoryAsync wallet wId target tgtSlot prefilter = do
         -- MN TEMPORARY: slow your roll!
         threadDelay 5000000 -- 5 sec
 
-        -- Gather the information we will need to decide if we are within K blocks of the tip.
-        slotCount <- getSlotCount (wallet ^. walletNode)
-        let flat = flattenSlotIdExplicit slotCount
+        -- Update each account's historical checkpoints
+        whenRight block $ \mb -> do
 
-        k <- getSecurityParameter (wallet ^. walletNode)
+            -- Gather the information we will need to decide if we are within K blocks of the tip.
+            slotCount <- getSlotCount (wallet ^. walletNode)
+            let flat = flattenSlotIdExplicit slotCount
 
-        -- Filter the blocks by account
-        (prefilteredBlocks, txMetas) <- prefilter blund
-        say ("   found " <> pretty (length prefilteredBlocks) <> " relevant accounts.")
+            k <- getSecurityParameter (wallet ^. walletNode)
+
+            -- Filter the blocks by account
+            (prefilteredBlocks, txMetas) <- prefilter blund
+            say ("   found " <> pretty (length prefilteredBlocks) <> " relevant accounts.")
 
-        -- Update each account's historical checkpoints
-        whenRight block $ \mb -> do
             let slotId = mb ^. mainBlockSlot
                 SecurityParameter k0 = k
             say ("  about to apply block for " <> pretty (mb ^. mainBlockSlot)
@@ -166,9 +168,9 @@ restoreWalletHistoryAsync wallet wId target tgtSlot prefilter = do
                               . (wriTargetSlot  .~ flat tgtSlot)) wId
             say "update done."
 
-        -- Store the TxMetas
-        forM_ txMetas (putTxMeta (wallet ^. walletMeta))
-        when (length txMetas > 0) $ say ("  txMetas: " <> pretty txMetas)
+            -- Store the TxMetas
+            forM_ txMetas (putTxMeta (wallet ^. walletMeta))
+            when (length txMetas > 0) $ say ("  txMetas: " <> pretty txMetas)
 
         -- Get the next block from the node and recurse.
         nextBlock (block ^. blockHeader) >>= \case