Add balance display to UI

src/Zenith/CLI.hs

@@ -121,6 +121,7 @@ data State = State
   , _dbPath :: !T.Text
   , _displayBox :: !DisplayType
   , _syncBlock :: !Int
+  , _balance :: !Integer
   }
 
 makeLenses ''State
@@ -150,6 +151,12 @@ drawUI s = [splashDialog s, helpDialog s, displayDialog s, inputDialog s, ui s]
                     "(None)"
                     (\(_, a) -> zcashAccountName $ entityVal a)
                     (L.listSelectedElement (st ^. accounts))))) <=>
+         C.hCenter
+           (str
+              ("Balance: " ++
+               if st ^. network == MainNet
+                 then displayZec (st ^. balance)
+                 else displayTaz (st ^. balance))) <=>
          listAddressBox "Addresses" (st ^. addresses) <+>
          B.vBorder <+>
          (C.hCenter (str ("Last block seen: " ++ show (st ^. syncBlock))) <=>
@@ -608,6 +615,10 @@ runZenithCLI config = do
               then getUserTx dbFilePath $ entityKey $ head addrList
               else return []
           block <- getMaxWalletBlock dbFilePath
+          bal <-
+            if not (null accList)
+              then getBalance dbFilePath $ entityKey $ head accList
+              else return 0
           void $
             M.defaultMain theApp $
             State
@@ -629,6 +640,7 @@ runZenithCLI config = do
               dbFilePath
               MsgDisplay
               block
+              bal
     Left e -> do
       print $
         "No Zebra node available on port " <>
@@ -650,6 +662,10 @@ refreshWallet s = do
     if not (null aL)
       then getAddresses (s ^. dbPath) $ entityKey $ head aL
       else return []
+  bal <-
+    if not (null aL)
+      then getBalance (s ^. dbPath) $ entityKey $ head aL
+      else return 0
   txL <-
     if not (null addrL)
       then getUserTx (s ^. dbPath) $ entityKey $ head addrL
@@ -658,7 +674,9 @@ refreshWallet s = do
   let addrL' = L.listReplace (Vec.fromList addrL) (Just 0) (s ^. addresses)
   let txL' = L.listReplace (Vec.fromList txL) (Just 0) (s ^. transactions)
   return $
-    (s & accounts .~ aL') & addresses .~ addrL' & transactions .~ txL' & msg .~
+    (s & accounts .~ aL') & balance .~ bal & addresses .~ addrL' & transactions .~
+    txL' &
+    msg .~
     "Switched to wallet: " ++
     T.unpack (zcashWalletName $ entityVal selWallet)
 
@@ -722,6 +740,7 @@ refreshAccount s = do
              Just (_j, w1) -> return w1
          Just (_k, w) -> return w
   aL <- getAddresses (s ^. dbPath) $ entityKey selAccount
+  bal <- getBalance (s ^. dbPath) $ entityKey selAccount
   let aL' = L.listReplace (Vec.fromList aL) (Just 0) (s ^. addresses)
   selAddress <-
     do case L.listSelectedElement aL' of
@@ -732,13 +751,13 @@ refreshAccount s = do
   case selAddress of
     Nothing ->
       return $
-      s & addresses .~ aL' & msg .~ "Switched to account: " ++
+      s & balance .~ bal & addresses .~ aL' & msg .~ "Switched to account: " ++
       T.unpack (zcashAccountName $ entityVal selAccount)
     Just (_i, a) -> do
       tList <- getUserTx (s ^. dbPath) $ entityKey a
       let tL' = L.listReplace (Vec.fromList tList) (Just 0) (s ^. transactions)
       return $
-        s & addresses .~ aL' & transactions .~ tL' & msg .~
+        s & balance .~ bal & addresses .~ aL' & transactions .~ tL' & msg .~
         "Switched to account: " ++
         T.unpack (zcashAccountName $ entityVal selAccount)
 

src/Zenith/DB.hs

@@ -1019,6 +1019,37 @@ upsertWalTx zt za =
        (zcashTransactionTime zt))
     []
 
+getBalance :: T.Text -> ZcashAccountId -> IO Integer
+getBalance dbPath za = do
+  trNotes <-
+    PS.runSqlite dbPath $ do
+      select $ do
+        n <- from $ table @WalletTrNote
+        where_ (n ^. WalletTrNoteAccId ==. val za)
+        where_ (n ^. WalletTrNoteSpent ==. val False)
+        pure n
+  let tAmts = map (walletTrNoteValue . entityVal) trNotes
+  let tBal = sum tAmts
+  sapNotes <-
+    PS.runSqlite dbPath $ do
+      select $ do
+        n1 <- from $ table @WalletSapNote
+        where_ (n1 ^. WalletSapNoteAccId ==. val za)
+        where_ (n1 ^. WalletSapNoteSpent ==. val False)
+        pure n1
+  let sAmts = map (walletSapNoteValue . entityVal) sapNotes
+  let sBal = sum sAmts
+  orchNotes <-
+    PS.runSqlite dbPath $ do
+      select $ do
+        n2 <- from $ table @WalletOrchNote
+        where_ (n2 ^. WalletOrchNoteAccId ==. val za)
+        where_ (n2 ^. WalletOrchNoteSpent ==. val False)
+        pure n2
+  let oAmts = map (walletOrchNoteValue . entityVal) orchNotes
+  let oBal = sum oAmts
+  return . fromIntegral $ tBal + sBal + oBal
+
 clearWalletTransactions :: T.Text -> IO ()
 clearWalletTransactions dbPath = do
   PS.runSqlite dbPath $ do