src/Zenith/DB.hs

@@ -1191,6 +1191,61 @@ getTrNotes pool tr = do
         where_ (tnotes ^. WalletTrNoteScript ==. val s)
         pure tnotes
 
+getTrFilteredNotes ::
+     ConnectionPool
+  -> [HexStringDB]
+  -> TransparentReceiver
+  -> IO [Entity WalletTrNote]
+getTrFilteredNotes pool txs tr = do
+  let s =
+        BS.concat
+          [ BS.pack [0x76, 0xA9, 0x14]
+          , (toBytes . tr_bytes) tr
+          , BS.pack [0x88, 0xAC]
+          ]
+  runNoLoggingT $
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      select $ do
+        (wt :& tnotes) <-
+          from $ table @WalletTransaction `innerJoin` table @WalletTrNote `on`
+          (\(wt :& tnotes) ->
+             wt ^. WalletTransactionId ==. tnotes ^. WalletTrNoteTx)
+        where_ (tnotes ^. WalletTrNoteScript ==. val s)
+        where_ (wt ^. WalletTransactionTxId `in_` valList txs)
+        pure tnotes
+
+traceTrDag :: ConnectionPool -> Entity WalletTrNote -> IO [Entity WalletTrNote]
+traceTrDag pool note = do
+  trSpend <-
+    runNoLoggingT $
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      selectOne $ do
+        trSpends <- from $ table @WalletTrSpend
+        where_ (trSpends ^. WalletTrSpendNote ==. val (entityKey note))
+        pure trSpends
+  case trSpend of
+    Nothing -> return []
+    Just tnote -> do
+      nxtChg <-
+        runNoLoggingT $
+        PS.retryOnBusy $
+        flip PS.runSqlPool pool $ do
+          selectOne $ do
+            nts <- from $ table @WalletTrNote
+            where_
+              (nts ^. WalletTrNoteTx ==. val (walletTrSpendTx $ entityVal tnote) &&.
+               nts ^.
+               WalletTrNoteChange ==.
+               val True)
+            pure nts
+      case nxtChg of
+        Nothing -> return []
+        Just nxt -> do
+          nxtSearch <- traceTrDag pool nxt
+          return $ nxt : nxtSearch
+
 getSapNotes :: ConnectionPool -> SaplingReceiver -> IO [Entity WalletSapNote]
 getSapNotes pool sr = do
   runNoLoggingT $
@@ -1201,6 +1256,57 @@ getSapNotes pool sr = do
         where_ (snotes ^. WalletSapNoteRecipient ==. val (getBytes sr))
         pure snotes
 
+getSapFilteredNotes ::
+     ConnectionPool
+  -> [HexStringDB]
+  -> SaplingReceiver
+  -> IO [Entity WalletSapNote]
+getSapFilteredNotes pool txs sr = do
+  runNoLoggingT $
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      select $ do
+        (wt :& snotes) <-
+          from $ table @WalletTransaction `innerJoin` table @WalletSapNote `on`
+          (\(wt :& snotes) ->
+             wt ^. WalletTransactionId ==. snotes ^. WalletSapNoteTx)
+        where_ (snotes ^. WalletSapNoteRecipient ==. val (getBytes sr))
+        where_ (wt ^. WalletTransactionTxId `in_` valList txs)
+        pure snotes
+
+traceSapDag ::
+     ConnectionPool -> Entity WalletSapNote -> IO [Entity WalletSapNote]
+traceSapDag pool note = do
+  sapSpend <-
+    runNoLoggingT $
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      selectOne $ do
+        sapSpends <- from $ table @WalletSapSpend
+        where_ (sapSpends ^. WalletSapSpendNote ==. val (entityKey note))
+        pure sapSpends
+  case sapSpend of
+    Nothing -> return []
+    Just snote -> do
+      nxtChg <-
+        runNoLoggingT $
+        PS.retryOnBusy $
+        flip PS.runSqlPool pool $ do
+          selectOne $ do
+            nts <- from $ table @WalletSapNote
+            where_
+              (nts ^. WalletSapNoteTx ==.
+               val (walletSapSpendTx $ entityVal snote) &&.
+               nts ^.
+               WalletSapNoteChange ==.
+               val True)
+            pure nts
+      case nxtChg of
+        Nothing -> return []
+        Just nxt -> do
+          nxtSearch <- traceSapDag pool nxt
+          return $ nxt : nxtSearch
+
 getOrchNotes :: ConnectionPool -> OrchardReceiver -> IO [Entity WalletOrchNote]
 getOrchNotes pool o = do
   runNoLoggingT $
@@ -1211,6 +1317,57 @@ getOrchNotes pool o = do
         where_ (onotes ^. WalletOrchNoteRecipient ==. val (getBytes o))
         pure onotes
 
+getOrchFilteredNotes ::
+     ConnectionPool
+  -> [HexStringDB]
+  -> OrchardReceiver
+  -> IO [Entity WalletOrchNote]
+getOrchFilteredNotes pool txs o = do
+  runNoLoggingT $
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      select $ do
+        (wt :& onotes) <-
+          from $ table @WalletTransaction `innerJoin` table @WalletOrchNote `on`
+          (\(wt :& onotes) ->
+             wt ^. WalletTransactionId ==. onotes ^. WalletOrchNoteTx)
+        where_ (onotes ^. WalletOrchNoteRecipient ==. val (getBytes o))
+        where_ (wt ^. WalletTransactionTxId `in_` valList txs)
+        pure onotes
+
+traceOrchDag ::
+     ConnectionPool -> Entity WalletOrchNote -> IO [Entity WalletOrchNote]
+traceOrchDag pool note = do
+  orchSpend <-
+    runNoLoggingT $
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      selectOne $ do
+        orchSpends <- from $ table @WalletOrchSpend
+        where_ (orchSpends ^. WalletOrchSpendNote ==. val (entityKey note))
+        pure orchSpends
+  case orchSpend of
+    Nothing -> return []
+    Just onote -> do
+      nxtChg <-
+        runNoLoggingT $
+        PS.retryOnBusy $
+        flip PS.runSqlPool pool $ do
+          selectOne $ do
+            nts <- from $ table @WalletOrchNote
+            where_
+              (nts ^. WalletOrchNoteTx ==.
+               val (walletOrchSpendTx $ entityVal onote) &&.
+               nts ^.
+               WalletOrchNoteChange ==.
+               val True)
+            pure nts
+      case nxtChg of
+        Nothing -> return []
+        Just nxt -> do
+          nxtSearch <- traceOrchDag pool nxt
+          return $ nxt : nxtSearch
+
 getWalletNotes ::
      ConnectionPool -- ^ database path
   -> Entity WalletAddress
@@ -1255,47 +1412,66 @@ getWalletTransactions pool w = do
     case tReceiver of
       Nothing -> return []
       Just tR -> liftIO $ getTrNotes pool tR
-  trChgNotes <-
-    case ctReceiver of
+  sapNotes <-
+    case sReceiver of
       Nothing -> return []
-      Just tR -> liftIO $ getTrNotes pool tR
+      Just sR -> liftIO $ getSapNotes pool sR
+  orchNotes <-
+    case oReceiver of
+      Nothing -> return []
+      Just oR -> liftIO $ getOrchNotes pool oR
+  clearUserTx (entityKey w)
+  mapM_ addTr trNotes
+  mapM_ addSap sapNotes
+  mapM_ addOrch orchNotes
   trSpends <-
     PS.retryOnBusy $
     flip PS.runSqlPool pool $ do
       select $ do
         trSpends <- from $ table @WalletTrSpend
         where_
-          (trSpends ^. WalletTrSpendNote `in_`
-           valList (map entityKey (trNotes <> trChgNotes)))
+          (trSpends ^. WalletTrSpendNote `in_` valList (map entityKey trNotes))
         pure trSpends
-  sapNotes <-
-    case sReceiver of
-      Nothing -> return []
-      Just sR -> liftIO $ getSapNotes pool sR
-  sapChgNotes <-
-    case csReceiver of
-      Nothing -> return []
-      Just sR -> liftIO $ getSapNotes pool sR
-  sapSpends <- mapM (getSapSpends . entityKey) (sapNotes <> sapChgNotes)
-  orchNotes <-
-    case oReceiver of
-      Nothing -> return []
-      Just oR -> liftIO $ getOrchNotes pool oR
-  orchChgNotes <-
-    case coReceiver of
-      Nothing -> return []
-      Just oR -> liftIO $ getOrchNotes pool oR
-  orchSpends <- mapM (getOrchSpends . entityKey) (orchNotes <> orchChgNotes)
-  clearUserTx (entityKey w)
-  mapM_ addTr trNotes
-  mapM_ addTr trChgNotes
-  mapM_ addSap sapNotes
-  mapM_ addSap sapChgNotes
-  mapM_ addOrch orchNotes
-  mapM_ addOrch orchChgNotes
+  sapSpends <- mapM (getSapSpends . entityKey) sapNotes
+  orchSpends <- mapM (getOrchSpends . entityKey) orchNotes
   mapM_ subTSpend trSpends
   mapM_ subSSpend $ catMaybes sapSpends
   mapM_ subOSpend $ catMaybes orchSpends
+  foundTxs <- getTxs $ entityKey w
+  trChgNotes <-
+    case ctReceiver of
+      Nothing -> return []
+      Just tR -> liftIO $ getTrFilteredNotes pool foundTxs tR
+  trChgNotes' <- liftIO $ mapM (traceTrDag pool) trChgNotes
+  trChgSpends <-
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      select $ do
+        trS <- from $ table @WalletTrSpend
+        where_
+          (trS ^. WalletTrSpendNote `in_`
+           valList (map entityKey (trChgNotes <> concat trChgNotes')))
+        pure trS
+  sapChgNotes <-
+    case csReceiver of
+      Nothing -> return []
+      Just sR -> liftIO $ getSapFilteredNotes pool foundTxs sR
+  sapChgNotes' <- liftIO $ mapM (traceSapDag pool) sapChgNotes
+  sapChgSpends <-
+    mapM (getSapSpends . entityKey) (sapChgNotes <> concat sapChgNotes')
+  orchChgNotes <-
+    case coReceiver of
+      Nothing -> return []
+      Just oR -> liftIO $ getOrchFilteredNotes pool foundTxs oR
+  orchChgNotes' <- liftIO $ mapM (traceOrchDag pool) orchChgNotes
+  orchChgSpends <-
+    mapM (getOrchSpends . entityKey) (orchChgNotes <> concat orchChgNotes')
+  mapM_ addTr (trChgNotes <> concat trChgNotes')
+  mapM_ addSap (sapChgNotes <> concat sapChgNotes')
+  mapM_ addOrch (orchChgNotes <> concat orchChgNotes')
+  mapM_ subTSpend trChgSpends
+  mapM_ subSSpend $ catMaybes sapChgSpends
+  mapM_ subOSpend $ catMaybes orchChgSpends
   where
     clearUserTx :: WalletAddressId -> NoLoggingT IO ()
     clearUserTx waId = do
@@ -1305,6 +1481,16 @@ getWalletTransactions pool w = do
             u <- from $ table @UserTx
             where_ (u ^. UserTxAddress ==. val waId)
             return ()
+    getTxs :: WalletAddressId -> NoLoggingT IO [HexStringDB]
+    getTxs waId = do
+      res <-
+        PS.retryOnBusy $
+        flip PS.runSqlPool pool $ do
+          select $ do
+            t <- from $ table @UserTx
+            where_ (t ^. UserTxAddress ==. val waId)
+            return (t ^. UserTxHex)
+      return $ map (\(Value x) -> x) res
     getSapSpends ::
          WalletSapNoteId -> NoLoggingT IO (Maybe (Entity WalletSapSpend))
     getSapSpends n = do