[CO-432] Adding foreign inputs to prefilter block (input-output-hk#3768)

[CO-432] fix applyBlockPartial

[CO-432] fixing lookup and updateUtxo

[CO-432] accommodating Ryan's suggestions

[CO-432] Get rid of foldr in favour of maps and moving reindexing upstream

[CO-432] hlint calming

[CO-432] Further code cleaning

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

@@ -45,7 +45,8 @@ import           Cardano.Wallet.Kernel.Internal
 import qualified Cardano.Wallet.Kernel.NodeStateAdaptor as Node
 import           Cardano.Wallet.Kernel.PrefilterTx (PrefilteredBlock (..),
                      prefilterBlock)
-import           Cardano.Wallet.Kernel.Read (getWalletCredentials)
+import           Cardano.Wallet.Kernel.Read (foreignPendingByAccount,
+                     getWalletCredentials, getWalletSnapshot)
 import           Cardano.Wallet.Kernel.Restore
 import qualified Cardano.Wallet.Kernel.Submission as Submission
 import           Cardano.Wallet.Kernel.Types (WalletId (..))
@@ -72,9 +73,10 @@ prefilterBlocks :: PassiveWallet
 prefilterBlocks pw bs = do
     let nm = makeNetworkMagic (pw ^. walletProtocolMagic)
     res <- getWalletCredentials pw
+    foreignPendings <- foreignPendingByAccount <$> getWalletSnapshot pw
     return $ case res of
          [] -> Nothing
-         xs -> Just $ map (\b -> first (b ^. rbContext,) $ prefilterBlock nm b xs) bs
+         xs -> Just $ map (\b -> first (b ^. rbContext,) $ prefilterBlock nm foreignPendings b xs) bs
 
 data BackfillFailed
     = SuccessorChanged BlockContext (Maybe BlockContext)

wallet-new/src/Cardano/Wallet/Kernel/DB/HdWallet/Read.hs

@@ -11,6 +11,7 @@ module Cardano.Wallet.Kernel.DB.HdWallet.Read (
   , addressesByRootId
   , addressesByAccountId
   , pendingByAccount
+  , foreignPendingByAccount
     -- | Simple lookups
   , lookupHdRootId
   , lookupHdAccountId
@@ -79,6 +80,14 @@ pendingByAccount = fmap aux . IxSet.toMap <$> view hdWalletsAccounts
     aux :: HdAccount -> Pending
     aux acc = acc ^. hdAccountState . hdAccountStateCurrent cpPending
 
+-- | All pending foreign transactions in all accounts
+foreignPendingByAccount :: Query' e HdWallets (Map HdAccountId Pending)
+foreignPendingByAccount = fmap aux . IxSet.toMap <$> view hdWalletsAccounts
+  where
+    aux :: HdAccount -> Pending
+    aux acc = acc ^. hdAccountState . hdAccountStateCurrent cpForeign
+
+
 {-------------------------------------------------------------------------------
   Simple lookups
 -------------------------------------------------------------------------------}

wallet-new/src/Cardano/Wallet/Kernel/DB/Read.hs

@@ -6,6 +6,7 @@ module Cardano.Wallet.Kernel.DB.Read (
   , addressesByRootId
   , addressesByAccountId
   , pendingByAccount
+  , foreignPendingByAccount
     -- | Lookups
   , lookupHdRootId
   , lookupHdAccountId
@@ -71,6 +72,9 @@ addressesByAccountId = liftNoErrorsHd1 HD.addressesByAccountId
 pendingByAccount :: DB -> Map HdAccountId Pending
 pendingByAccount = liftNoErrorsHd0 HD.pendingByAccount
 
+foreignPendingByAccount :: DB -> Map HdAccountId Pending
+foreignPendingByAccount = liftNoErrorsHd0 HD.foreignPendingByAccount
+
 lookupHdRootId :: DB -> HdRootId -> Either UnknownHdRoot HdRoot
 lookupHdRootId = liftHd1 HD.lookupHdRootId
 

wallet-new/src/Cardano/Wallet/Kernel/DB/Spec/Pending.hs

@@ -32,7 +32,6 @@ module Cardano.Wallet.Kernel.DB.Spec.Pending (
   , txOuts
   , change
   , removeInputs
-  , removeForeign
   , PendingDiff
   ) where
 
@@ -179,31 +178,6 @@ removeInputs usedInputs (toMap -> p) =
     shouldKeep :: Core.TxAux -> Bool
     shouldKeep tx = Util.disjoint (Core.txIns tx) usedInputs
 
--- | Remove any transactions corresponding to a given Utxo
---
--- We store foreign transaction as `Pending`, using their transaction ids
--- and outputs, however, we filter them based on the UTxO since they don't
--- actually come from our wallet, but from a "foreign" source. Therefore, they
--- appear in blocks as outputs targetting one of our addresses (and not coming
--- from one of them).
-removeForeign :: Core.Utxo -> Pending -> (Pending, Set Core.TxId)
-removeForeign utxo  (toMap -> p) =
-    let (pToKeep, pToEvict) = Map.partitionWithKey shouldKeep p
-    in (fromMap pToKeep, Map.keysSet pToEvict)
-  where
-    -- | We remove foreign transactions from te Pending map iif their id
-    -- appear in an available UTxO. It means the transaction has been accepted
-    -- and isn't pending anymore.
-    shouldKeep :: Core.TxId -> Core.TxAux -> Bool
-    shouldKeep txId _ =
-        not (txId `elem` (mapMaybe toTxId $ Map.keys utxo))
-
-    toTxId :: Core.TxIn -> Maybe Core.TxId
-    toTxId txIn = case txIn of
-        Core.TxInUtxo txId _ -> Just txId
-        Core.TxInUnknown _ _ -> Nothing
-
-
 {-------------------------------------------------------------------------------
   Internal auxiliary
 

wallet-new/src/Cardano/Wallet/Kernel/DB/Spec/Update.hs

@@ -32,7 +32,7 @@ import           Serokell.Util.Text (listBuilderJSON)
 import           Test.QuickCheck (Arbitrary (..), elements)
 
 import qualified Pos.Chain.Block as Core
-import           Pos.Chain.Txp (Utxo)
+import           Pos.Chain.Txp (TxIn, Utxo)
 import qualified Pos.Chain.Txp as Txp
 import qualified Pos.Core as Core
 import           Pos.Core.Chrono (NewestFirst (..), OldestFirst (..))
@@ -195,9 +195,9 @@ applyBlock (SecurityParameter k) pb = do
         utxo              = current ^. checkpointUtxo        . fromDb
         balance           = current ^. checkpointUtxoBalance . fromDb
         (utxo', balance') = updateUtxo      pb (utxo, balance)
-        (pending', rem1)  = updatePending   pb (current ^. checkpointPending)
+        (pending', rem1)  = updatePending   (pfbInputs pb) (current ^. checkpointPending)
         blockMeta'        = updateBlockMeta pb (current ^. checkpointBlockMeta)
-        (foreign', rem2)  = updateForeign   pb (current ^. checkpointForeign)
+        (foreign', rem2)  = updatePending   (pfbForeignInputs pb) (current ^. checkpointForeign)
     if (pfbContext pb) `blockContextSucceeds` (current ^. checkpointContext . lazy) then do
       put $ Checkpoints . takeNewest k . NewestFirst $ Checkpoint {
           _checkpointUtxo        = InDb utxo'
@@ -228,9 +228,9 @@ applyBlockPartial pb = do
         utxo              = current ^. pcheckpointUtxo        . fromDb
         balance           = current ^. pcheckpointUtxoBalance . fromDb
         (utxo', balance') = updateUtxo           pb (utxo, balance)
-        (pending', rem1)  = updatePending        pb (current ^. pcheckpointPending)
+        (pending', rem1)  = updatePending        (pfbInputs pb) (current ^. pcheckpointPending)
         blockMeta'        = updateLocalBlockMeta pb (current ^. pcheckpointBlockMeta)
-        (foreign', rem2)  = updatePending        pb (current ^. pcheckpointForeign)
+        (foreign', rem2)  = updatePending        (pfbForeignInputs pb) (current ^. pcheckpointForeign)
     if (pfbContext pb) `blockContextSucceeds` (current ^. cpContext . lazy) then do
       put $ Checkpoints $ NewestFirst $ PartialCheckpoint {
           _pcheckpointUtxo        = InDb utxo'
@@ -347,9 +347,10 @@ updateUtxo PrefilteredBlock{..} (utxo, balance) =
     --
     -- * pfbOutputs corresponds to what the spec calls utxo^+ / txouts_b
     -- * pfbInputs  corresponds to what the spec calls txins_b
+    extendedInputs = Set.union pfbInputs pfbForeignInputs
     utxoUnion = Map.union utxo pfbOutputs
-    utxoMin   = utxoUnion `Core.utxoRestrictToInputs` pfbInputs
-    utxo'     = utxoUnion `Core.utxoRemoveInputs`     pfbInputs
+    utxoMin   = utxoUnion `Core.utxoRestrictToInputs` extendedInputs
+    utxo'     = utxoUnion `Core.utxoRemoveInputs` extendedInputs
     balance'  = Core.unsafeIntegerToCoin $
                     Core.coinToInteger balance
                   + Core.utxoBalance pfbOutputs
@@ -358,17 +359,8 @@ updateUtxo PrefilteredBlock{..} (utxo, balance) =
 -- | Update the pending transactions with the given prefiltered block
 --
 -- Returns the set of transactions that got removed from the pending set.
-updatePending :: PrefilteredBlock -> Pending -> (Pending, Set Txp.TxId)
-updatePending PrefilteredBlock{..} = Pending.removeInputs pfbInputs
-
--- | Update the foreign transactions with the given prefiltered block
---
--- Returns the set of transactions that got removed from the pending set.
--- Note that this is slightly different from pending transactions as it doesn't
--- filter on the "inputs" (which aren't part of the wallet for Foreign
--- transactions) but filter out based on the outputs.
-updateForeign :: PrefilteredBlock -> Pending -> (Pending, Set Txp.TxId)
-updateForeign PrefilteredBlock{..} = Pending.removeForeign pfbOutputs
+updatePending :: Set TxIn -> Pending -> (Pending, Set Txp.TxId)
+updatePending inputs = Pending.removeInputs inputs
 
 takeNewest :: Int -> NewestFirst StrictNonEmpty a -> NewestFirst StrictNonEmpty a
 takeNewest = liftNewestFirst . SNE.take

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

@@ -1,5 +1,6 @@
 {-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections       #-}
 {-# LANGUAGE TypeFamilies        #-}
 
 module Cardano.Wallet.Kernel.PrefilterTx
@@ -43,6 +44,8 @@ import           Cardano.Wallet.Kernel.DB.InDb (InDb (..), fromDb)
 import           Cardano.Wallet.Kernel.DB.Resolved (ResolvedBlock,
                      ResolvedInput, ResolvedTx, rbContext, rbTxs,
                      resolvedToTxMeta, rtxInputs, rtxOutputs)
+import           Cardano.Wallet.Kernel.DB.Spec.Pending (Pending)
+import qualified Cardano.Wallet.Kernel.DB.Spec.Pending as Pending
 import           Cardano.Wallet.Kernel.DB.TxMeta.Types
 import           Cardano.Wallet.Kernel.Decrypt (WalletDecrCredentials,
                      eskToWalletDecrCredentials, selectOwnAddresses)
@@ -63,19 +66,22 @@ type AddrWithId = (HdAddressId,Address)
 -- the block that are relevant to the wallet.
 data PrefilteredBlock = PrefilteredBlock {
       -- | Relevant inputs
-      pfbInputs  :: !(Set TxIn)
+      pfbInputs        :: !(Set TxIn)
+
+      -- | Relevant foreign inputs
+    , pfbForeignInputs :: !(Set TxIn)
 
       -- | Relevant outputs
-    , pfbOutputs :: !Utxo
+    , pfbOutputs       :: !Utxo
 
       -- | all output addresses present in the Utxo
-    , pfbAddrs   :: ![AddrWithId]
+    , pfbAddrs         :: ![AddrWithId]
 
       -- | Prefiltered block metadata
-    , pfbMeta    :: !LocalBlockMeta
+    , pfbMeta          :: !LocalBlockMeta
 
       -- | Block context
-    , pfbContext :: !BlockContext
+    , pfbContext       :: !BlockContext
     }
 
 deriveSafeCopy 1 'base ''PrefilteredBlock
@@ -87,11 +93,12 @@ deriveSafeCopy 1 'base ''PrefilteredBlock
 -- relevance to that account
 emptyPrefilteredBlock :: BlockContext -> PrefilteredBlock
 emptyPrefilteredBlock context = PrefilteredBlock {
-      pfbInputs  = Set.empty
-    , pfbOutputs = Map.empty
-    , pfbAddrs   = []
-    , pfbMeta    = emptyLocalBlockMeta
-    , pfbContext = context
+      pfbInputs         = Set.empty
+    , pfbForeignInputs  = Set.empty
+    , pfbOutputs        = Map.empty
+    , pfbAddrs          = []
+    , pfbMeta           = emptyLocalBlockMeta
+    , pfbContext        = context
     }
 
 type WalletKey = (WalletId, WalletDecrCredentials)
@@ -164,13 +171,52 @@ prefilterTx wKey tx = ((prefInps',prefOuts'),metas)
 --
 -- NOTE: we can rely on a Monoidal fold here to combine the maps
 -- 'Map HdAccountId a' since the accounts will be unique accross wallet keys.
-prefilterTxForWallets :: [WalletKey]
-                      -> ResolvedTx
-                      -> ((Map HdAccountId (Set TxIn)
-                         , Map HdAccountId UtxoSummaryRaw)
-                         , [TxMeta])
-prefilterTxForWallets wKeys tx =
-    mconcat $ map ((flip prefilterTx) tx) wKeys
+-- The function decomposes a resolved block into input and output transactions and meta for given wallets
+-- In case of input transactions the two kinds are differentiated:
+-- (a) the input transactions belonging to some wallet
+-- (b) the foreign transactions.
+-- The foreign transactions are identified by picking the input transactions from the resolved one
+-- that happen to be in foreign pending set.
+prefilterTxForWallets
+    :: [WalletKey]
+    -> Map TxIn HdAccountId
+    -> ResolvedTx
+    -> ((Map HdAccountId (Set TxIn, Set TxIn)
+        , Map HdAccountId UtxoSummaryRaw)
+       , [TxMeta])
+prefilterTxForWallets wKeys foreignPendingByTransaction tx =
+    ((extend inputsE foreignInputsE, outputs),meta)
+  where
+    ((inputs,outputs),meta) = mconcat $ map ((flip prefilterTx) tx) wKeys
+
+    --NOTE: to find the foreign inputs in the transaction, we need to look at _all_ the inputs, since they will not be present in the prefiltered inputs
+    allInputs :: Set TxIn
+    allInputs = Set.fromList $ map fst $ toList (tx ^. rtxInputs  . fromDb)
+
+    foreignInputs :: Map HdAccountId (Set TxIn)
+    foreignInputs =
+        reindexByAccount $ Map.filterWithKey (\txin _ -> Set.member txin allInputs) foreignPendingByTransaction
+
+    inputsE, foreignInputsE :: Map HdAccountId (Set TxIn, Set TxIn)
+    inputsE = Map.map (, Set.empty) inputs
+    foreignInputsE =  Map.map (Set.empty,) foreignInputs
+
+    extend
+        :: Map HdAccountId (Set TxIn, Set TxIn)
+        -> Map HdAccountId (Set TxIn, Set TxIn)
+        -> Map HdAccountId (Set TxIn, Set TxIn)
+    extend inputs_ foreignInputs_ =
+        Map.unionWith (\inp fInp -> (fst inp, snd fInp)) inputs_ foreignInputs_
+
+    reindexByAccount
+        :: Map TxIn HdAccountId
+        -> Map HdAccountId (Set TxIn)
+    reindexByAccount byTxIn =
+        Map.fromListWith Set.union $ Map.elems $ Map.mapWithKey f byTxIn
+      where
+          f :: TxIn -> HdAccountId -> (HdAccountId, Set TxIn)
+          f txin accId = (accId, Set.singleton txin)
+
 
 -- | Prefilter inputs of a transaction
 prefilterInputs :: WalletKey
@@ -281,11 +327,13 @@ extendWithSummary (onlyOurInps,onlyOurOuts) utxoWithAddrId
 -- | Prefilter the transactions of a resolved block for the given wallets.
 --
 --   Returns prefiltered blocks indexed by HdAccountId.
-prefilterBlock :: NetworkMagic
-               -> ResolvedBlock
-               -> [(WalletId, EncryptedSecretKey)]
-               -> (Map HdAccountId PrefilteredBlock, [TxMeta])
-prefilterBlock nm block rawKeys =
+prefilterBlock
+    :: NetworkMagic
+    -> Map HdAccountId Pending
+    -> ResolvedBlock
+    -> [(WalletId, EncryptedSecretKey)]
+    -> (Map HdAccountId PrefilteredBlock, [TxMeta])
+prefilterBlock nm foreignPendingByAccount block rawKeys =
       (Map.fromList
     $ map (mkPrefBlock (block ^. rbContext) inpAll outAll)
     $ Set.toList accountIds
@@ -294,41 +342,57 @@ prefilterBlock nm block rawKeys =
     wKeys :: [WalletKey]
     wKeys = map toWalletKey rawKeys
 
-    inps :: [Map HdAccountId (Set TxIn)]
+    foreignPendingByTransaction :: Map TxIn HdAccountId
+    foreignPendingByTransaction = reindexByTransaction $ Map.map Pending.txIns foreignPendingByAccount
+
+    inps :: [Map HdAccountId (Set TxIn, Set TxIn)]
     outs :: [Map HdAccountId UtxoSummaryRaw]
-    (ios, conMetas) = unzip $ map (prefilterTxForWallets wKeys) (block ^. rbTxs)
+    (ios, conMetas) = unzip $ map (prefilterTxForWallets wKeys foreignPendingByTransaction) (block ^. rbTxs)
     (inps, outs) = unzip ios
     metas = concat conMetas
 
-    inpAll :: Map HdAccountId (Set TxIn)
+    inpAll :: Map HdAccountId (Set TxIn, Set TxIn)
     outAll :: Map HdAccountId UtxoSummaryRaw
-    inpAll = Map.unionsWith Set.union inps
+    inpAll = Map.unionsWith (\pair1 pair2 -> (Set.union (fst pair1) (fst pair2),Set.union (snd pair1) (fst pair2))) inps
     outAll = Map.unionsWith Map.union outs
 
     accountIds = Map.keysSet inpAll `Set.union` Map.keysSet outAll
 
     toWalletKey :: (WalletId, EncryptedSecretKey) -> WalletKey
     toWalletKey (wid, esk) = (wid, eskToWalletDecrCredentials nm esk)
 
+    reindexByTransaction :: Map HdAccountId (Set TxIn) -> Map TxIn HdAccountId
+    reindexByTransaction byAccount = Map.fromList $ Set.toList $ Set.unions $ Map.elems $ Map.mapWithKey f byAccount
+        where
+            f :: HdAccountId -> Set TxIn -> Set (TxIn, HdAccountId)
+            f accId = Set.map (, accId)
+
+
 mkPrefBlock :: BlockContext
-            -> Map HdAccountId (Set TxIn)
+            -> Map HdAccountId (Set TxIn, Set TxIn)
             -> Map HdAccountId (Map TxIn (TxOutAux, AddressSummary))
             -> HdAccountId
             -> (HdAccountId, PrefilteredBlock)
 mkPrefBlock context inps outs accId = (accId, PrefilteredBlock {
-        pfbInputs  = inps'
-      , pfbOutputs = outs'
-      , pfbAddrs   = addrs''
-      , pfbMeta    = blockMeta'
-      , pfbContext = context
+        pfbInputs         = walletInps'
+      , pfbForeignInputs  = foreignInps'
+      , pfbOutputs        = outs'
+      , pfbAddrs          = addrs''
+      , pfbMeta           = blockMeta'
+      , pfbContext        = context
       })
     where
         fromAddrSummary :: AddressSummary -> AddrWithId
         fromAddrSummary AddressSummary{..} = (addrSummaryId,addrSummaryAddr)
 
         byAccountId accId'' def dict = fromMaybe def $ Map.lookup accId'' dict
 
-        inps'           =                  byAccountId accId Set.empty inps
+        walletInps = Map.map fst $
+                     Map.filter (not . Set.null . fst) inps
+        foreignInps = Map.map snd $
+                      Map.filter (not . Set.null . snd) inps
+        walletInps'           =                  byAccountId accId Set.empty walletInps
+        foreignInps'          =                  byAccountId accId Set.empty foreignInps
         (outs', addrs') = fromUtxoSummary (byAccountId accId Map.empty outs)
 
         addrs''    = nub $ map fromAddrSummary addrs'
@@ -406,8 +470,10 @@ instance Buildable PrefilteredBlock where
   build PrefilteredBlock{..} = bprint
     ( "PrefilteredBlock "
     % "{ inputs:  " % listJson
+    % "{ foreignInputs:  " % listJson
     % ", outputs: " % mapJson
     % "}"
     )
     (Set.toList pfbInputs)
+    (Set.toList pfbForeignInputs)
     pfbOutputs