feat: enable shielding in GUI

src/Zenith/CLI.hs

@@ -878,7 +878,6 @@ scanZebra dbP zHost zPort b eChan znet = do
                   liftIO $ BC.writeBChan eChan $ TickMsg "startSync"
                   return ()
             else do
-              liftIO $ BC.writeBChan eChan $ TickVal 1.0
               liftIO $ BC.writeBChan eChan $ TickMsg "startSync"
   where
     processBlock :: ConnectionPool -> Float -> Int -> IO ()
@@ -928,8 +927,7 @@ appEvent (BT.AppEvent t) = do
       case s ^. displayBox of
         AddrDisplay -> return ()
         MsgDisplay -> do
-          if m == "startSync"
-            then do
+          when (m == "startSync") $ do
             selWallet <-
               do case L.listSelectedElement $ s ^. wallets of
                    Nothing -> do
@@ -957,7 +955,6 @@ appEvent (BT.AppEvent t) = do
             BT.put ns
             BT.modify $ set msg ""
             BT.modify $ set displayBox BlankDisplay
-            else return ()
         PhraseDisplay -> return ()
         TxDisplay -> return ()
         TxIdDisplay -> return ()

src/Zenith/DB.hs

@@ -159,7 +159,7 @@ share
     script BS.ByteString
     change Bool
     position Int
-    UniqueTNote tx script
+    UniqueTNote tx accId script
     deriving Show Eq
   WalletTrSpend
     tx WalletTransactionId OnDeleteCascade OnUpdateCascade
@@ -2828,24 +2828,36 @@ rewindWalletData pool b net = do
           (blk ^. ZcashBlockHeight >. val b) &&.
           (blk ^. ZcashBlockNetwork ==. val net)
   logDebugN "Completed data store rewind"
+  {-
+   -_ <- liftIO $ clearTrees pool
+   -logDebugN "Cleared commitment trees"
+   -}
   saplingOutputIx <- liftIO $ getSaplingOutputAtBlock pool net b
   orchardActionIx <- liftIO $ getOrchardActionAtBlock pool net b
   case saplingOutputIx of
     Nothing -> logErrorN "Couldn't get Sapling output index for tree rewind"
     Just soIx -> do
       saplingTree <- liftIO $ getSaplingTree pool
-      let truncSapTree = truncateTree (maybe InvalidTree fst saplingTree) soIx
+      truncSapTree <- truncateTree (maybe InvalidTree fst saplingTree) soIx
       _ <- liftIO $ upsertSaplingTree pool b truncSapTree
       logDebugN $ "Truncated Sapling tree at index " <> T.pack (show soIx)
       case orchardActionIx of
         Nothing -> logErrorN "Couldn't get Orchard action index for tree rewind"
         Just oaIx -> do
           orchardTree <- liftIO $ getOrchardTree pool
-          let truncOrchTree =
-                truncateTree (maybe InvalidTree fst orchardTree) oaIx
+          truncOrchTree <- truncateTree (maybe InvalidTree fst orchardTree) oaIx
           _ <- liftIO $ upsertOrchardTree pool b truncOrchTree
           logDebugN $ "Truncated Orchard tree at index " <> T.pack (show oaIx)
 
+clearTrees :: ConnectionPool -> IO ()
+clearTrees pool =
+  runNoLoggingT $
+  PS.retryOnBusy $
+  flip PS.runSqlPool pool $ do
+    delete $ do
+      tr <- from $ table @TreeStore
+      return ()
+
 getSaplingOutputAtBlock ::
      ConnectionPool -> ZcashNetDB -> Int -> IO (Maybe Int64)
 getSaplingOutputAtBlock pool znet b = do
@@ -2863,10 +2875,7 @@ getSaplingOutputAtBlock pool znet b = do
              txs ^. ZcashTransactionId ==. sOutputs ^. ShieldOutputTx)
         where_ (blks ^. ZcashBlockHeight <=. val b)
         where_ (blks ^. ZcashBlockNetwork ==. val znet)
-        orderBy
-          [ desc $ txs ^. ZcashTransactionId
-          , desc $ sOutputs ^. ShieldOutputPosition
-          ]
+        orderBy [desc $ sOutputs ^. ShieldOutputId]
         return sOutputs
   case r of
     Nothing -> return Nothing
@@ -2889,10 +2898,7 @@ getOrchardActionAtBlock pool znet b = do
              txs ^. ZcashTransactionId ==. oActions ^. OrchActionTx)
         where_ (blks ^. ZcashBlockHeight <=. val b)
         where_ (blks ^. ZcashBlockNetwork ==. val znet)
-        orderBy
-          [ desc $ txs ^. ZcashTransactionId
-          , desc $ oActions ^. OrchActionPosition
-          ]
+        orderBy [desc $ oActions ^. OrchActionId]
         return oActions
   case r of
     Nothing -> return Nothing

src/Zenith/GUI.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE LambdaCase #-}
 
 module Zenith.GUI where
 
@@ -10,9 +11,14 @@ import Codec.QRCode
 import Codec.QRCode.JuicyPixels
 import Control.Concurrent (threadDelay)
 import Control.Exception (throwIO, try)
-import Control.Monad (unless, when)
+import Control.Monad (forM_, unless, when)
 import Control.Monad.IO.Class (liftIO)
-import Control.Monad.Logger (runNoLoggingT, runStderrLoggingT)
+import Control.Monad.Logger
+  ( LoggingT
+  , logDebugN
+  , runNoLoggingT
+  , runStderrLoggingT
+  )
 import Data.Aeson
 import qualified Data.ByteString as BS
 import Data.HexString (toText)
@@ -38,12 +44,16 @@ import ZcashHaskell.Orchard
   , isValidUnifiedAddress
   , parseAddress
   )
-import ZcashHaskell.Transparent (encodeTransparentReceiver)
+import ZcashHaskell.Transparent
+  ( decodeTransparentAddress
+  , encodeTransparentReceiver
+  )
 import ZcashHaskell.Types
   ( BlockResponse(..)
   , Scope(..)
   , ToBytes(..)
   , UnifiedAddress(..)
+  , ValidAddress(..)
   , ZcashNet(..)
   , ZebraGetBlockChainInfo(..)
   , ZebraGetInfo(..)
@@ -56,6 +66,7 @@ import Zenith.Scanner (checkIntegrity, processTx, rescanZebra, updateConfs)
 import Zenith.Types hiding (ZcashAddress(..))
 import Zenith.Utils
   ( displayAmount
+  , getChainTip
   , isRecipientValidGUI
   , isValidString
   , isZecAddressValid
@@ -83,7 +94,7 @@ data AppEvent
   | SwitchAddr !Int
   | SwitchAcc !Int
   | SwitchWal !Int
-  | UpdateBalance !(Integer, Integer)
+  | UpdateBalance !(Integer, Integer, Integer, Integer)
   | CopyAddr !(Maybe (Entity WalletAddress))
   | LoadTxs ![Entity UserTx]
   | LoadAddrs ![Entity WalletAddress]
@@ -131,6 +142,10 @@ data AppEvent
   | CloseShield
   | ShowDeShield
   | CloseDeShield
+  | SendDeShield
+  | SendShield
+  | StartSync
+  | TreeSync
   deriving (Eq, Show)
 
 data AppModel = AppModel
@@ -741,7 +756,7 @@ buildUI wenv model = widgetTree
       box
         (label (fromMaybe "?" $ model ^. modalMsg) `styleBasic`
          [textSize 12, textFont "Bold"]) `styleBasic`
-      [bgColor (white & L.a .~ 0.5)]
+      [bgColor (white & L.a .~ 0.7)]
     txOverlay =
       case model ^. showTx of
         Nothing -> alert CloseTx $ label "N/A"
@@ -990,21 +1005,17 @@ buildUI wenv model = widgetTree
                            [textFont "Bold", textSize 12])
                       , separatorLine `styleBasic` [fgColor btnColor]
                       , spacer
-                      , hstack
-                          [ filler
-                          , label ("Amount : ") `styleBasic`
-                            [width 50, textFont "Bold"]
-                          , spacer
-                          , label (displayAmount (model ^. network) 100) `styleBasic`
-                            [width 50, textFont "Bold"]
-                          , filler
-                          ]
+                      , label
+                          ("Shield " <>
+                           displayAmount (model ^. network) (model ^. tBalance) <>
+                           "?") `styleBasic`
+                        [width 50, textFont "Regular"]
                       , spacer
                       , box_
                           [alignMiddle]
                           (hstack
                              [ filler
-                             , mainButton "Proceed" NotImplemented `nodeEnabled`
+                             , mainButton "Proceed" SendShield `nodeEnabled`
                                True
                              , spacer
                              , mainButton "Cancel" CloseShield `nodeEnabled`
@@ -1033,16 +1044,25 @@ buildUI wenv model = widgetTree
                            [textFont "Bold", textSize 12])
                       , separatorLine `styleBasic` [fgColor btnColor]
                       , spacer
-                      , hstack
-                          [ (label "Total Transparent : " `styleBasic`
-                             [textFont "Bold"])
-                          , (label "0.00")
+                      , box_
+                          []
+                          (vstack
+                             [ hstack
+                                 [ label "Total Transparent : " `styleBasic`
+                                   [textFont "Bold"]
+                                 , label
+                                     (displayAmount
+                                        (model ^. network)
+                                        (model ^. tBalance))
                                  ]
                              , spacer
                              , hstack
-                          [ (label "Total Shielded : " `styleBasic`
-                             [textFont "Bold"])
-                          , (label "0.00")
+                                 [ label "Total Shielded : " `styleBasic`
+                                   [textFont "Bold"]
+                                 , label
+                                     (displayAmount
+                                        (model ^. network)
+                                        (model ^. sBalance))
                                  ]
                              , spacer
                              , hstack
@@ -1065,12 +1085,13 @@ buildUI wenv model = widgetTree
                                        (textColor red)
                                    ]
                                  ]
+                             ])
                       , spacer
                       , box_
                           [alignMiddle]
                           (hstack
                              [ filler
-                             , mainButton "Proceed" NotImplemented `nodeEnabled`
+                             , mainButton "Proceed" SendDeShield `nodeEnabled`
                                True
                              , spacer
                              , mainButton "Cancel" CloseDeShield `nodeEnabled`
@@ -1085,23 +1106,6 @@ buildUI wenv model = widgetTree
            ]) `styleBasic`
       [bgColor (white & L.a .~ 0.5)]
 
---                          , spacer
---                          , numericField_
---                              sendAmount
---                              [ decimals 8
---                              , minValue 0.0
---                              , maxValue
---                                  (fromIntegral (model ^. tBalance) / 100000000.0)
---                              , validInput tBalanceValid
---                              , onChange CheckAmount
---                              ] `styleBasic`
---                            [ width 150
---                            , styleIf
---                                (not $ model ^. tBalanceValid)
---                                (textColor red)
---                            ]
---                               (model ^. amountValid && model ^. recipientValid)
---                               (model ^. amountValid && model ^. recipientValid)
 notImplemented = NotImplemented
 
 generateQRCodes :: Config -> IO ()
@@ -1303,11 +1307,13 @@ handleEvent wenv node model evt =
         UpdateBalance <$> do
           dbPool <- runNoLoggingT $ initPool $ c_dbPath $ model ^. configuration
           case selectAccount i of
-            Nothing -> return (0, 0)
+            Nothing -> return (0, 0, 0, 0)
             Just acc -> do
               b <- getBalance dbPool $ entityKey acc
               u <- getUnconfirmedBalance dbPool $ entityKey acc
-              return (b, u)
+              s <- getShieldedBalance dbPool $ entityKey acc
+              t <- getTransparentBalance dbPool $ entityKey acc
+              return (b, u, s, t)
       , Event $ SetPool OrchardPool
       ]
     SwitchWal i ->
@@ -1319,9 +1325,9 @@ handleEvent wenv node model evt =
             Nothing -> return []
             Just wal -> runNoLoggingT $ getAccounts dbPool $ entityKey wal
       ]
-    UpdateBalance (b, u) ->
+    UpdateBalance (b, u, s, t) ->
       [ Model $
-        model & balance .~ b & unconfBalance .~
+        model & balance .~ b & sBalance .~ s & tBalance .~ t & unconfBalance .~
         (if u == 0
            then Nothing
            else Just u)
@@ -1371,7 +1377,7 @@ handleEvent wenv node model evt =
         else [Event $ NewAccount currentWallet]
     LoadWallets a ->
       if not (null a)
-        then [ Model $ model & wallets .~ a
+        then [ Model $ model & wallets .~ a & modalMsg .~ Nothing
              , Event $ SwitchWal $ model ^. selWallet
              ]
         else [Event NewWallet]
@@ -1381,11 +1387,15 @@ handleEvent wenv node model evt =
     CloseTxId -> [Model $ model & showId .~ Nothing]
     ShowTx i -> [Model $ model & showTx ?~ i]
     TickUp ->
-      if (model ^. timer) < 90
+      if isNothing (model ^. modalMsg)
+        then if (model ^. timer) < 90
                then [Model $ model & timer .~ (1 + model ^. timer)]
                else if (model ^. barValue) == 1.0
-               then [ Model $ model & timer .~ 0 & barValue .~ 0.0
+                      then [ Model $
+                             model & timer .~ 0 & barValue .~ 0.0 & modalMsg ?~
+                             "Downloading blocks..."
                            , Producer $
+                             runStderrLoggingT .
                              scanZebra
                                (c_dbPath $ model ^. configuration)
                                (c_zebraHost $ model ^. configuration)
@@ -1393,20 +1403,23 @@ handleEvent wenv node model evt =
                                (model ^. network)
                            ]
                       else [Model $ model & timer .~ 0]
-    SyncVal i ->
-      if (i + model ^. barValue) >= 0.999
-        then [ Model $ model & barValue .~ 1.0 & modalMsg .~ Nothing
+        else [Model $ model & timer .~ 0]
+    TreeSync -> [Model $ model & modalMsg ?~ "Updating commitment trees..."]
+    StartSync ->
+      [ Model $ model & modalMsg ?~ "Updating wallet..."
       , Task $ do
           case currentWallet of
             Nothing -> return $ ShowError "No wallet available"
             Just cW -> do
               runStderrLoggingT $ syncWallet (model ^. configuration) cW
               pool <-
-                       runNoLoggingT $
-                       initPool $ c_dbPath $ model ^. configuration
+                runNoLoggingT $ initPool $ c_dbPath $ model ^. configuration
               wL <- getWallets pool (model ^. network)
               return $ LoadWallets wL
       ]
+    SyncVal i ->
+      if (i + model ^. barValue) >= 0.999
+        then [Model $ model & barValue .~ 1.0 & modalMsg .~ Nothing]
         else [ Model $
                model & barValue .~ validBarValue (i + model ^. barValue) &
                modalMsg ?~
@@ -1491,7 +1504,10 @@ handleEvent wenv node model evt =
         model & msgAB ?~ "Function not implemented..." & menuPopup .~ False
       ]
     CloseMsgAB -> [Model $ model & msgAB .~ Nothing & inError .~ False]
-    ShowShield -> [Model $ model & shieldZec .~ True & menuPopup .~ False]
+    ShowShield ->
+      if model ^. tBalance > 0
+        then [Model $ model & shieldZec .~ True & menuPopup .~ False]
+        else [Event $ ShowError "No transparent funds in this account"]
     CloseShield -> [Model $ model & shieldZec .~ False]
     ShowDeShield -> [Model $ model & deShieldZec .~ True & menuPopup .~ False]
     CloseDeShield -> [Model $ model & deShieldZec .~ False]
@@ -1507,6 +1523,31 @@ handleEvent wenv node model evt =
           abList <- getAdrBook dbPool $ model ^. network
           return $ LoadAbList abList
       ]
+    SendDeShield ->
+      case currentAccount of
+        Nothing ->
+          [Event $ ShowError "No account available", Event CloseDeShield]
+        Just acc ->
+          [ Producer $
+            deshieldTransaction
+              (model ^. configuration)
+              (model ^. network)
+              (entityKey acc)
+              currentAddress
+              (fromFloatDigits $ model ^. sendAmount)
+          , Event CloseDeShield
+          ]
+    SendShield ->
+      case currentAccount of
+        Nothing -> [Event $ ShowError "No account available", Event CloseShield]
+        Just acc ->
+          [ Producer $
+            shieldTransaction
+              (model ^. configuration)
+              (model ^. network)
+              (entityKey acc)
+          , Event CloseShield
+          ]
   where
     currentWallet =
       if null (model ^. wallets)
@@ -1621,43 +1662,57 @@ handleEvent wenv node model evt =
       res <- liftIO $ updateAdrsInAdrBook pool d a a
       return $ ShowMessage "Address Book entry updated!!"
 
-scanZebra :: T.Text -> T.Text -> Int -> ZcashNet -> (AppEvent -> IO ()) -> IO ()
+scanZebra ::
+     T.Text
+  -> T.Text
+  -> Int
+  -> ZcashNet
+  -> (AppEvent -> IO ())
+  -> LoggingT IO ()
 scanZebra dbPath zHost zPort net sendMsg = do
   bStatus <- liftIO $ checkBlockChain zHost zPort
-  pool <- runNoLoggingT $ initPool dbPath
+  pool <- liftIO $ runNoLoggingT $ initPool dbPath
   b <- liftIO $ getMinBirthdayHeight pool $ ZcashNetDB net
-  dbBlock <- getMaxBlock pool $ ZcashNetDB net
-  chkBlock <- checkIntegrity dbPath zHost zPort net dbBlock 1
-  syncChk <- isSyncing pool
+  dbBlock <- liftIO $ getMaxBlock pool $ ZcashNetDB net
+  chkBlock <- liftIO $ checkIntegrity dbPath zHost zPort net dbBlock 1
+  logDebugN $ "dbBlock: " <> T.pack (show dbBlock)
+  logDebugN $ "chkBlock: " <> T.pack (show chkBlock)
+  syncChk <- liftIO $ isSyncing pool
   if syncChk
-    then sendMsg (ShowError "Sync already in progress")
+    then liftIO $ sendMsg (ShowError "Sync already in progress")
     else do
       let sb =
             if chkBlock == dbBlock
               then max dbBlock b
               else max chkBlock b
       unless (chkBlock == dbBlock || chkBlock == 1) $
-        runStderrLoggingT $ rewindWalletData pool sb $ ZcashNetDB net
+        rewindWalletData pool sb $ ZcashNetDB net
       if sb > zgb_blocks bStatus || sb < 1
-        then sendMsg (ShowError "Invalid starting block for scan")
+        then liftIO $ sendMsg (ShowError "Invalid starting block for scan")
         else do
           let bList = [(sb + 1) .. (zgb_blocks bStatus)]
           if not (null bList)
             then do
               let step = (1.0 :: Float) / fromIntegral (length bList)
-              _ <- startSync pool
-              mapM_ (processBlock pool step) bList
+              _ <- liftIO $ startSync pool
+              mapM_ (liftIO . processBlock pool step) bList
               confUp <-
-                try $ updateConfs zHost zPort pool :: IO (Either IOError ())
+                liftIO $ try $ updateConfs zHost zPort pool :: LoggingT
+                  IO
+                  (Either IOError ())
               case confUp of
                 Left _e0 -> do
-                  _ <- completeSync pool Failed
+                  _ <- liftIO $ completeSync pool Failed
+                  liftIO $
                     sendMsg
                       (ShowError "Failed to update unconfirmed transactions")
                 Right _ -> do
-                  _ <- completeSync pool Successful
-                  return ()
-            else sendMsg (SyncVal 1.0)
+                  liftIO $ sendMsg TreeSync
+                  _ <- updateCommitmentTrees pool zHost zPort $ ZcashNetDB net
+                  _ <- liftIO $ completeSync pool Successful
+                  logDebugN "Starting wallet sync"
+                  liftIO $ sendMsg StartSync
+            else liftIO $ sendMsg (SyncVal 1.0)
   where
     processBlock :: ConnectionPool -> Float -> Int -> IO ()
     processBlock pool step bl = do
@@ -1697,6 +1752,83 @@ scanZebra dbPath zHost zPort net sendMsg = do
               mapM_ (processTx zHost zPort bi pool) $ bl_txs blk
               sendMsg (SyncVal step)
 
+shieldTransaction ::
+     Config -> ZcashNet -> ZcashAccountId -> (AppEvent -> IO ()) -> IO ()
+shieldTransaction config znet accId sendMsg = do
+  sendMsg $ ShowModal "Shielding funds..."
+  let dbPath = c_dbPath config
+  let zHost = c_zebraHost config
+  let zPort = c_zebraPort config
+  pool <- runNoLoggingT $ initPool dbPath
+  bl <- getChainTip zHost zPort
+  res <- runNoLoggingT $ shieldTransparentNotes pool zHost zPort znet accId bl
+  forM_ res $ \case
+    Left e -> sendMsg $ ShowError $ T.pack (show e)
+    Right rawTx -> do
+      sendMsg $ ShowMsg "Transaction ready, sending to Zebra..."
+      resp <-
+        makeZebraCall
+          zHost
+          zPort
+          "sendrawtransaction"
+          [Data.Aeson.String $ toText rawTx]
+      case resp of
+        Left e1 -> sendMsg $ ShowError $ "Zebra error: " <> T.pack (show e1)
+        Right txId -> sendMsg $ ShowTxId txId
+
+deshieldTransaction ::
+     Config
+  -> ZcashNet
+  -> ZcashAccountId
+  -> Maybe (Entity WalletAddress)
+  -> Scientific
+  -> (AppEvent -> IO ())
+  -> IO ()
+deshieldTransaction config znet accId addR pnote sendMsg = do
+  case addR of
+    Nothing -> sendMsg $ ShowError "No address available"
+    Just addr -> do
+      sendMsg $ ShowModal "De-shielding funds..."
+      let dbPath = c_dbPath config
+      let zHost = c_zebraHost config
+      let zPort = c_zebraPort config
+      pool <- runNoLoggingT $ initPool dbPath
+      bl <- getChainTip zHost zPort
+      let tAddrMaybe =
+            Transparent <$>
+            ((decodeTransparentAddress .
+              E.encodeUtf8 . encodeTransparentReceiver znet) =<<
+             (t_rec =<<
+              (isValidUnifiedAddress .
+               E.encodeUtf8 . getUA . walletAddressUAddress)
+                (entityVal addr)))
+      case tAddrMaybe of
+        Nothing -> sendMsg $ ShowError "No transparent address available"
+        Just tAddr -> do
+          res <-
+            runStderrLoggingT $
+            deshieldNotes
+              pool
+              zHost
+              zPort
+              znet
+              accId
+              bl
+              (ProposedNote (ValidAddressAPI tAddr) pnote Nothing)
+          case res of
+            Left e -> sendMsg $ ShowError $ T.pack (show e)
+            Right rawTx -> do
+              sendMsg $ ShowModal "Transaction ready, sending to Zebra..."
+              resp <-
+                makeZebraCall
+                  zHost
+                  zPort
+                  "sendrawtransaction"
+                  [Data.Aeson.String $ toText rawTx]
+              case resp of
+                Left e1 -> sendMsg $ ShowError $ "Zebra error: " <> showt e1
+                Right txId -> sendMsg $ ShowTxId txId
+
 sendTransaction ::
      Config
   -> ZcashNet
@@ -1815,6 +1947,14 @@ runZenithGUI config = do
                       then getUnconfirmedBalance pool $ entityKey $ head accList
                       else return 0
                   abList <- getAdrBook pool (zgb_net chainInfo)
+                  shieldBal <-
+                    if not (null accList)
+                      then getShieldedBalance pool $ entityKey $ head accList
+                      else return 0
+                  transBal <-
+                    if not (null accList)
+                      then getTransparentBalance pool $ entityKey $ head accList
+                      else return 0
                   let model =
                         AppModel
                           config
@@ -1874,9 +2014,9 @@ runZenithGUI config = do
                           Full
                           False
                           False
-                          0
+                          transBal
                           False
-                          0
+                          shieldBal
                           False
                   startApp model handleEvent buildUI (params hD)
         Left _e -> print "Zebra not available"

src/Zenith/Tree.hs

@@ -9,9 +9,11 @@
 module Zenith.Tree where
 
 import Codec.Borsh
+import Control.Monad.Logger (LoggingT, logDebugN)
 import Data.HexString
 import Data.Int (Int32, Int64, Int8)
 import Data.Maybe (fromJust, isNothing)
+import qualified Data.Text as T
 import qualified GHC.Generics as GHC
 import qualified Generics.SOP as SOP
 import ZcashHaskell.Orchard (combineOrchardNodes, getOrchardNodeValue)
@@ -179,14 +181,49 @@ getNotePosition (Branch _ x y) i
   | otherwise = Nothing
 getNotePosition _ _ = Nothing
 
-truncateTree :: Monoid v => Node v => Tree v -> Int64 -> Tree v
+truncateTree :: Monoid v => Node v => Tree v -> Int64 -> LoggingT IO (Tree v)
 truncateTree (Branch s x y) i
-  | getLevel s == 1 && getIndex (value x) == i = branch x EmptyLeaf
-  | getLevel s == 1 && getIndex (value y) == i = branch x y
-  | getIndex (value x) >= i =
-    branch (truncateTree x i) (getEmptyRoot (getLevel s))
-  | getIndex (value y) >= i = branch x (truncateTree y i)
-truncateTree x _ = x
+  | getLevel s == 1 && getIndex (value x) == i = do
+    logDebugN $ T.pack $ show (getLevel s) ++ " Trunc to left leaf"
+    return $ branch x EmptyLeaf
+  | getLevel s == 1 && getIndex (value y) == i = do
+    logDebugN $ T.pack $ show (getLevel s) ++ " Trunc to right leaf"
+    return $ branch x y
+  | getIndex (value x) >= i = do
+    logDebugN $
+      T.pack $
+      show (getLevel s) ++
+      ": " ++ show i ++ " left i: " ++ show (getIndex (value x))
+    l <- truncateTree x i
+    return $ branch (l) (getEmptyRoot (getLevel (value x)))
+  | getIndex (value y) /= 0 && getIndex (value y) >= i = do
+    logDebugN $
+      T.pack $
+      show (getLevel s) ++
+      ": " ++ show i ++ " right i: " ++ show (getIndex (value y))
+    r <- truncateTree y i
+    return $ branch x (r)
+  | otherwise = do
+    logDebugN $
+      T.pack $
+      show (getLevel s) ++
+      ": " ++
+      show (getIndex (value x)) ++ " catchall " ++ show (getIndex (value y))
+    return InvalidTree
+truncateTree x _ = return x
+
+countLeaves :: Node v => Tree v -> Int64
+countLeaves (Branch s x y) =
+  if isFull s
+    then 2 ^ getLevel s
+    else countLeaves x + countLeaves y
+countLeaves (PrunedBranch x) =
+  if isFull x
+    then 2 ^ getLevel x
+    else 0
+countLeaves (Leaf _) = 1
+countLeaves EmptyLeaf = 0
+countLeaves InvalidTree = 0
 
 data SaplingNode = SaplingNode
   { sn_position :: !Position

test/Spec.hs

@@ -42,6 +42,7 @@ import ZcashHaskell.Sapling
 import ZcashHaskell.Transparent
   ( decodeExchangeAddress
   , decodeTransparentAddress
+  , encodeExchangeAddress
   )
 import ZcashHaskell.Types
   ( DecodedNote(..)
@@ -59,6 +60,7 @@ import ZcashHaskell.Types
   , Scope(..)
   , ShieldedOutput(..)
   , TxError(..)
+  , UnifiedAddress(..)
   , ValidAddress(..)
   , ZcashNet(..)
   )
@@ -623,23 +625,76 @@ main = do
               let updatedTree = foldl append newTree [(cmx1, 4), (cmx2, 5)]
               getNotePosition updatedTree 4 `shouldBe` Just 39734
         it "Truncate tree" $ do
+          pool <- runNoLoggingT $ initPool "/home/rav/Zenith/zenith.db"
+          maxBlock <- getMaxBlock pool $ ZcashNetDB TestNet
+          dbTree <- getOrchardTree pool
+          case dbTree of
+            Nothing -> assertFailure "failed to get tree from DB"
+            Just (oTree, oSync) -> do
+              let startBlock = oSync - 5
+              zebraTreesIn <-
+                getCommitmentTrees
+                  pool
+                  "localhost"
+                  18232
+                  (ZcashNetDB TestNet)
+                  startBlock
+              ix <- getOrchardActionAtBlock pool (ZcashNetDB TestNet) startBlock
+              case ix of
+                Nothing -> assertFailure "couldn't find index at block"
+                Just i -> do
+                  updatedTree <-
+                    runFileLoggingT "test.log" $ truncateTree oTree i
+                  let finalAnchor =
+                        getOrchardTreeAnchor $
+                        OrchardCommitmentTree $ ztiOrchard zebraTreesIn
+                  getHash (value updatedTree) `shouldBe` finalAnchor
+        it "Counting leaves in tree" $ do
           let tree =
                 OrchardCommitmentTree $
                 hexString
                   "0136a7886d7d73bc1845223165fd9cb0cef02046c707e8f88a7f61564720bd0f3501dca1fbdd7b5ba92a0809af5e85874626ce0db14d0532a48e41dde6f0f81b46011f0001fb48c27bd07e68f27aba47cd6e93fa961e0ef8c63f993963a614e56855d2013c0001ea572db9c5c2d24c7ad9132ae32b27179466bf67a580d59901d13b281d3f530b01c160348f10b9ad893d9731317ebe36ac8665e01c52cbe15a56aa9b72e4e6c41e000001cd7695156de2debdc5b13ea84d32e4e3ac020fb0aa7cd372c57ce765103bd70401746e6bc066a10e7f80a9ff8993dcb25c819edd64f2ca10ac248ef7848d41450500011e6191f91b3fceb62dc881a156e1b9d2e88e09dca25093cf9c4936c8869fb41a013bf8b923e4187754e85175748d9cce4824a6787e4258977b5bfe1ba59012c032000001f3bbdc62260c4fca5c84bf3487246d4542da48eeeec8ec40c1029b6908eef83c00000000000000000000000000000000"
-          let cmx1 =
-                hexString
-                  "1712ead46028d4349e234abf59e94e0640fe7a0829e2e2e17e1a931631810400" :: OrchardCommitment
-          let cmx2 =
-                hexString
-                  "39f5ad39817fb432fa07c5feb3a957189fbe7662a4b5555ca95093b6d853cf07" :: OrchardCommitment
           case getOrchardTreeParts tree of
             Nothing -> assertFailure "Failed to get tree parts"
             Just t1 -> do
               let newTree = mkOrchardTree t1
-              let updatedTree = foldl append newTree [(cmx1, 4), (cmx2, 5)]
-              let truncTree = truncateTree updatedTree 4
-              getIndex (value truncTree) `shouldBe` 4
+              countLeaves newTree `shouldBe`
+                fromIntegral (1 + getPosition (value newTree))
+        it "Validate large load" $ do
+          pool <- runNoLoggingT $ initPool "/home/rav/Zenith/zenith.db"
+          maxBlock <- getMaxBlock pool $ ZcashNetDB TestNet
+          let startBlock = maxBlock - 2000
+          zebraTreesIn <-
+            getCommitmentTrees
+              pool
+              "localhost"
+              18232
+              (ZcashNetDB TestNet)
+              startBlock
+          zebraTreesOut <-
+            getCommitmentTrees
+              pool
+              "localhost"
+              18232
+              (ZcashNetDB TestNet)
+              maxBlock
+          case getOrchardTreeParts $
+               OrchardCommitmentTree $ ztiOrchard zebraTreesIn of
+            Nothing -> assertFailure "Failed to get tree parts"
+            Just t1 -> do
+              let newTree = mkOrchardTree t1
+              oAct <- getOrchardActions pool startBlock $ ZcashNetDB TestNet
+              let cmxs =
+                    map
+                      (\(_, y) ->
+                         ( getHex $ orchActionCmx $ entityVal y
+                         , fromSqlKey $ entityKey y))
+                      oAct
+              let updatedTree = foldl' append newTree cmxs
+              let finalAnchor =
+                    getOrchardTreeAnchor $
+                    OrchardCommitmentTree $ ztiOrchard zebraTreesOut
+              getHash (value updatedTree) `shouldBe` finalAnchor
         it "Validate tree from DB" $ do
           pool <- runNoLoggingT $ initPool "/home/rav/Zenith/zenith.db"
           dbTree <- getOrchardTree pool
@@ -657,6 +712,19 @@ main = do
                     getOrchardTreeAnchor $
                     OrchardCommitmentTree $ ztiOrchard zebraTrees
               getHash (value oTree) `shouldBe` finalAnchor
+    describe "TEX address" $ do
+      it "from UA" $ do
+        let addr =
+              parseAddress
+                "utest1fqtne08sdgmae0g0un7j3h6ss9gafguprv0yvkxv4trxxsdxx467pxkkc98cpsyk5r2enwwpn3p5c6aw537wyvlz20hs7vcqc4uhm22yfjnrsm8hy2hjjrscvhk2ac32rzndu94hh28gdl62wqgy3yev7w0gj9lmmz6yasghmle6tllx4yjv9sjt0xml66y9lyxc4rkk6q425nc5gxa"
+        case addr of
+          Nothing -> assertFailure "failed to parse address"
+          Just (Unified ua) ->
+            case (encodeExchangeAddress (ua_net ua) =<< (t_rec ua)) of
+              Nothing -> assertFailure "failed to encode TEX"
+              Just tex ->
+                tex `shouldBe` "textest1jze8c9jxxrpct34tpe4pvquz8nvxsxt6gawqqf"
+          Just _ -> assertFailure "no transparent receiver"
     describe "Creating Tx" $ do
       describe "Full" $ do
         it "To Orchard" $ do