2024-11-21 15:39:19 +00:00
9 changed files with 346 additions and 202 deletions
--- a/src/Zenith/CLI.hs
+++ b/src/Zenith/CLI.hs
@ -870,9 +870,16 @@ scanZebra dbP zHost zPort b eChan znet = do
                    BC.writeBChan eChan $
                    TickMsg "Failed to update unconfirmed transactions"
                Right _ -> do
+                  logDebugN "Updated confirmations"
+                  logDebugN "Starting commitment tree update"
+                  _ <- updateCommitmentTrees pool zHost zPort (ZcashNetDB znet)
+                  logDebugN "Finished tree update"
                  _ <- liftIO $ completeSync pool Successful
+                  liftIO $ BC.writeBChan eChan $ TickMsg "startSync"
                  return ()
-            else liftIO $ BC.writeBChan eChan $ TickVal 1.0
+            else do
+              liftIO $ BC.writeBChan eChan $ TickVal 1.0
+              liftIO $ BC.writeBChan eChan $ TickMsg "startSync"
  where
    processBlock :: ConnectionPool -> Float -> Int -> IO ()
    processBlock pool step bl = do
@ -920,28 +927,8 @@ appEvent (BT.AppEvent t) = do
    TickMsg m -> do
      case s ^. displayBox of
        AddrDisplay -> return ()
-        MsgDisplay -> return ()
-        PhraseDisplay -> return ()
-        TxDisplay -> return ()
-        TxIdDisplay -> return ()
-        SyncDisplay -> return ()
-        SendDisplay -> BT.modify $ set msg m
-        AdrBookEntryDisplay -> return ()
-        BlankDisplay -> return ()
-    TickTx txid -> do
-      BT.modify $ set sentTx (Just txid)
-      BT.modify $ set displayBox TxIdDisplay
-    TickVal v -> do
-      case s ^. displayBox of
-        AddrDisplay -> return ()
-        MsgDisplay -> return ()
-        PhraseDisplay -> return ()
-        TxDisplay -> return ()
-        TxIdDisplay -> return ()
-        SendDisplay -> return ()
-        AdrBookEntryDisplay -> return ()
-        SyncDisplay -> do
-          if s ^. barValue == 1.0
+        MsgDisplay -> do
+          if m == "startSync"
            then do
              selWallet <-
                do case L.listSelectedElement $ s ^. wallets of
@ -968,8 +955,34 @@ appEvent (BT.AppEvent t) = do
              updatedState <- BT.get
              ns <- liftIO $ refreshWallet updatedState
              BT.put ns
+              BT.modify $ set msg ""
              BT.modify $ set displayBox BlankDisplay
+            else return ()
+        PhraseDisplay -> return ()
+        TxDisplay -> return ()
+        TxIdDisplay -> return ()
+        SyncDisplay -> return ()
+        SendDisplay -> BT.modify $ set msg m
+        AdrBookEntryDisplay -> return ()
+        BlankDisplay -> return ()
+    TickTx txid -> do
+      BT.modify $ set sentTx (Just txid)
+      BT.modify $ set displayBox TxIdDisplay
+    TickVal v -> do
+      case s ^. displayBox of
+        AddrDisplay -> return ()
+        MsgDisplay -> return ()
+        PhraseDisplay -> return ()
+        TxDisplay -> return ()
+        TxIdDisplay -> return ()
+        SendDisplay -> return ()
+        AdrBookEntryDisplay -> return ()
+        SyncDisplay -> do
+          if s ^. barValue == 1.0
+            then do
+              BT.modify $ set msg "Decoding, please wait..."
              BT.modify $ set barValue 0.0
+              BT.modify $ set displayBox MsgDisplay
            else BT.modify $ set barValue $ validBarValue (v + s ^. barValue)
        BlankDisplay -> do
          case s ^. dialogBox of
@ -990,7 +1003,9 @@ appEvent (BT.AppEvent t) = do
                then do
                  BT.modify $ set barValue 0.0
                  BT.modify $ set displayBox SyncDisplay
-                  sBlock <- liftIO $ getMinBirthdayHeight pool
+                  sBlock <-
+                    liftIO $
+                    getMinBirthdayHeight pool (ZcashNetDB $ s ^. network)
                  _ <-
                    liftIO $
                    forkIO $
@ -1215,7 +1230,7 @@ appEvent (BT.VtyEvent e) = do
                                 (s ^. zebraPort)
                                 (s ^. network)
                                 (entityKey selAcc)
-                                 (bl + 5)
+                                 bl
                                 (fs1 ^. sendAmt)
                                 (fs1 ^. sendTo)
                                 (fs1 ^. sendMemo)
--- a/src/Zenith/Core.hs
+++ b/src/Zenith/Core.hs
@ -24,7 +24,7 @@ import Data.Binary.Get hiding (getBytes)
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Lazy as LBS
 import Data.HexString (HexString, hexBytes, hexString, toBytes, toText)
-import Data.Int (Int64)
+import Data.Int (Int32, Int64)
 import Data.List
 import Data.Maybe (fromJust, fromMaybe)
 import Data.Scientific (Scientific, scientific, toBoundedInteger)
@ -50,6 +50,7 @@ import ZcashHaskell.Orchard
  , genOrchardSpendingKey
  , getOrchardFrontier
  , getOrchardNotePosition
+  , getOrchardTreeParts
  , getOrchardWitness
  , isValidUnifiedAddress
  , updateOrchardCommitmentTree
@ -62,6 +63,7 @@ import ZcashHaskell.Sapling
  , genSaplingSpendingKey
  , getSaplingFrontier
  , getSaplingNotePosition
+  , getSaplingTreeParts
  , getSaplingWitness
  , updateSaplingCommitmentTree
  , updateSaplingWitness
@ -74,6 +76,7 @@ import ZcashHaskell.Transparent
 import ZcashHaskell.Types
 import ZcashHaskell.Utils
 import Zenith.DB
+import Zenith.Tree
 import Zenith.Types
  ( Config(..)
  , HexStringDB(..)
@ -303,85 +306,61 @@ findSaplingOutputs config b znet za = do
  let zn = getNet znet
  pool <- liftIO $ runNoLoggingT $ initPool dbPath
  tList <- liftIO $ getShieldedOutputs pool b znet
-  trees <- liftIO $ getCommitmentTrees pool zebraHost zebraPort znet (b - 1)
-  logDebugN "getting Sapling frontier"
-  let sT = getSaplingFrontier $ SaplingCommitmentTree $ ztiSapling trees
+  sT <- liftIO $ getSaplingTree pool
  case sT of
    Nothing ->
      liftIO $ throwIO $ userError "Failed to read Sapling commitment tree"
-    Just sT' -> do
-      logDebugN "Sapling frontier valid"
-      decryptNotes sT' zn pool tList
+    Just (sT', treeSync) -> do
+      logDebugN "Sapling tree valid"
+      mapM_ (decryptNotes sT' zn pool) tList
      sapNotes <- liftIO $ getWalletSapNotes pool (entityKey za)
      liftIO $ findSapSpends pool (entityKey za) sapNotes
  where
    sk :: SaplingSpendingKeyDB
    sk = zcashAccountSapSpendKey $ entityVal za
    decryptNotes ::
-         SaplingFrontier
+         Tree SaplingNode
      -> ZcashNet
      -> ConnectionPool
-      -> [(Entity ZcashTransaction, Entity ShieldOutput)]
+      -> (Entity ZcashTransaction, Entity ShieldOutput)
      -> NoLoggingT IO ()
-    decryptNotes _ _ _ [] = return ()
-    decryptNotes st n pool ((zt, o):txs) = do
-      let updatedTree =
-            updateSaplingCommitmentTree
-              st
-              (getHex $ shieldOutputCmu $ entityVal o)
-      logDebugN "updated frontier"
-      case updatedTree of
-        Nothing ->
-          liftIO $ throwIO $ userError "Failed to update commitment tree"
-        Just uT -> do
-          let noteWitness = getSaplingWitness uT
-          logDebugN "got witness"
-          let notePos = getSaplingNotePosition <$> noteWitness
+    decryptNotes st n pool (zt, o) = do
+      case getNotePosition st $ fromSqlKey $ entityKey o of
+        Nothing -> do
+          logErrorN "Couldn't find sapling note in commitment tree"
+          return ()
+        Just nP -> do
          logDebugN "got position"
-          case notePos of
-            Nothing ->
-              liftIO $ throwIO $ userError "Failed to obtain note position"
-            Just nP -> do
-              case decodeShOut External n nP o of
+          case decodeShOut External n nP o of
+            Nothing -> do
+              logDebugN "couldn't decode external"
+              case decodeShOut Internal n nP o of
                Nothing -> do
-                  logDebugN "couldn't decode external"
-                  case decodeShOut Internal n nP o of
-                    Nothing -> do
-                      logDebugN "couldn't decode internal"
-                      decryptNotes uT n pool txs
-                    Just dn1 -> do
-                      wId <-
-                        liftIO $ saveWalletTransaction pool (entityKey za) zt
-                      liftIO $
-                        saveWalletSapNote
-                          pool
-                          wId
-                          nP
-                          (fromJust noteWitness)
-                          True
-                          (entityKey za)
-                          (entityKey o)
-                          dn1
-                      decryptNotes uT n pool txs
-                Just dn0 -> do
+                  logDebugN "couldn't decode internal"
+                Just dn1 -> do
                  wId <- liftIO $ saveWalletTransaction pool (entityKey za) zt
                  liftIO $
                    saveWalletSapNote
                      pool
                      wId
                      nP
-                      (fromJust noteWitness)
-                      False
+                      True
                      (entityKey za)
                      (entityKey o)
-                      dn0
-                  decryptNotes uT n pool txs
+                      dn1
+            Just dn0 -> do
+              wId <- liftIO $ saveWalletTransaction pool (entityKey za) zt
+              liftIO $
+                saveWalletSapNote
+                  pool
+                  wId
+                  nP
+                  False
+                  (entityKey za)
+                  (entityKey o)
+                  dn0
    decodeShOut ::
-         Scope
-      -> ZcashNet
-      -> Integer
-      -> Entity ShieldOutput
-      -> Maybe DecodedNote
+         Scope -> ZcashNet -> Int32 -> Entity ShieldOutput -> Maybe DecodedNote
    decodeShOut scope n pos s = do
      decodeSaplingOutputEsk
        (getSapSK sk)
@ -394,7 +373,7 @@ findSaplingOutputs config b znet za = do
           (getHex $ shieldOutputProof $ entityVal s))
        n
        scope
-        pos
+        (fromIntegral pos)

 -- | Get Orchard actions 
 findOrchardActions ::
@ -410,67 +389,52 @@ findOrchardActions config b znet za = do
  let zn = getNet znet
  pool <- runNoLoggingT $ initPool dbPath
  tList <- getOrchardActions pool b znet
-  trees <- getCommitmentTrees pool zebraHost zebraPort znet (b - 1)
-  let sT = getOrchardFrontier $ OrchardCommitmentTree $ ztiOrchard trees
+  sT <- getOrchardTree pool
  case sT of
    Nothing -> throwIO $ userError "Failed to read Orchard commitment tree"
-    Just sT' -> do
-      decryptNotes sT' zn pool tList
+    Just (sT', treeSync) -> do
+      mapM_ (decryptNotes sT' zn pool) tList
      orchNotes <- getWalletOrchNotes pool (entityKey za)
      findOrchSpends pool (entityKey za) orchNotes
  where
    decryptNotes ::
-         OrchardFrontier
+         Tree OrchardNode
      -> ZcashNet
      -> ConnectionPool
-      -> [(Entity ZcashTransaction, Entity OrchAction)]
+      -> (Entity ZcashTransaction, Entity OrchAction)
      -> IO ()
-    decryptNotes _ _ _ [] = return ()
-    decryptNotes ot n pool ((zt, o):txs) = do
-      let updatedTree =
-            updateOrchardCommitmentTree
-              ot
-              (getHex $ orchActionCmx $ entityVal o)
-      case updatedTree of
-        Nothing -> throwIO $ userError "Failed to update commitment tree"
-        Just uT -> do
-          let noteWitness = getOrchardWitness uT
-          let notePos = getOrchardNotePosition <$> noteWitness
-          case notePos of
-            Nothing -> throwIO $ userError "Failed to obtain note position"
-            Just nP ->
-              case decodeOrchAction External nP o of
-                Nothing ->
-                  case decodeOrchAction Internal nP o of
-                    Nothing -> decryptNotes uT n pool txs
-                    Just dn1 -> do
-                      wId <- saveWalletTransaction pool (entityKey za) zt
-                      saveWalletOrchNote
-                        pool
-                        wId
-                        nP
-                        (fromJust noteWitness)
-                        True
-                        (entityKey za)
-                        (entityKey o)
-                        dn1
-                      decryptNotes uT n pool txs
-                Just dn -> do
+    decryptNotes ot n pool (zt, o) = do
+      case getNotePosition ot (fromSqlKey $ entityKey o) of
+        Nothing -> do
+          return ()
+        Just nP ->
+          case decodeOrchAction External nP o of
+            Nothing ->
+              case decodeOrchAction Internal nP o of
+                Nothing -> return ()
+                Just dn1 -> do
                  wId <- saveWalletTransaction pool (entityKey za) zt
                  saveWalletOrchNote
                    pool
                    wId
                    nP
-                    (fromJust noteWitness)
-                    False
+                    True
                    (entityKey za)
                    (entityKey o)
-                    dn
-                  decryptNotes uT n pool txs
+                    dn1
+            Just dn -> do
+              wId <- saveWalletTransaction pool (entityKey za) zt
+              saveWalletOrchNote
+                pool
+                wId
+                nP
+                False
+                (entityKey za)
+                (entityKey o)
+                dn
    sk :: OrchardSpendingKeyDB
    sk = zcashAccountOrchSpendKey $ entityVal za
-    decodeOrchAction ::
-         Scope -> Integer -> Entity OrchAction -> Maybe DecodedNote
+    decodeOrchAction :: Scope -> Int32 -> Entity OrchAction -> Maybe DecodedNote
    decodeOrchAction scope pos o =
      decryptOrchardActionSK (getOrchSK sk) scope $
      OrchardAction
@ -814,11 +778,6 @@ shieldTransparentNotes ::
 shieldTransparentNotes pool zebraHost zebraPort znet za bh = do
  accRead <- liftIO $ getAccountById pool za
  logDebugN $ T.pack $ "Target block: " ++ show bh
-  {-
-   -trees <- liftIO $ getCommitmentTrees zebraHost zebraPort bh
-   -let sT = SaplingCommitmentTree $ ztiSapling trees
-   -let oT = OrchardCommitmentTree $ ztiOrchard trees
-   -}
  case accRead of
    Nothing -> do
      logErrorN "Can't find Account"
@ -831,6 +790,8 @@ shieldTransparentNotes pool zebraHost zebraPort znet za bh = do
              (\x ->
                 filter (\y -> walletTrNoteAddress (entityVal y) == x) trNotes')
              dRecvs
+      sTree <- liftIO $ getSaplingTree pool
+      oTree <- liftIO $ getOrchardTree pool
      forM fNotes $ \trNotes -> do
        let noteTotal = getTotalAmount (trNotes, [], [])
        tSpends <-
@ -864,8 +825,8 @@ shieldTransparentNotes pool zebraHost zebraPort znet za bh = do
        tx <-
          liftIO $
          createTransaction
-            Nothing
-            Nothing
+            (maybe (hexString "00") (getHash . value . fst) sTree)
+            (maybe (hexString "00") (getHash . value . fst) oTree)
            tSpends
            []
            []
@ -931,10 +892,8 @@ prepareTxV2 pool zebraHost zebraPort zn za bh pnotes policy = do
  let recipients = map extractReceiver pnotes
  logDebugN $ T.pack $ show recipients
  logDebugN $ T.pack $ "Target block: " ++ show bh
-  trees <-
-    liftIO $ getCommitmentTrees pool zebraHost zebraPort (ZcashNetDB zn) bh
-  let sT = SaplingCommitmentTree $ ztiSapling trees
-  let oT = OrchardCommitmentTree $ ztiOrchard trees
+  sTree <- liftIO $ getSaplingTree pool
+  oTree <- liftIO $ getOrchardTree pool
  case accRead of
    Nothing -> do
      logErrorN "Can't find Account"
@ -969,24 +928,6 @@ prepareTxV2 pool zebraHost zebraPort zn za bh pnotes policy = do
              logDebugN $ T.pack $ show oList
              let noteTotal = getTotalAmount (tList, sList, oList)
              logDebugN $ "noteTotal: " <> T.pack (show noteTotal)
-              tSpends <-
-                liftIO $
-                prepTSpends
-                  (getTranSK $ zcashAccountTPrivateKey $ entityVal acc)
-                  tList
-                --print tSpends
-              sSpends <-
-                liftIO $
-                prepSSpends
-                  (getSapSK $ zcashAccountSapSpendKey $ entityVal acc)
-                  sList
-                --print sSpends
-              oSpends <-
-                liftIO $
-                prepOSpends
-                  (getOrchSK $ zcashAccountOrchSpendKey $ entityVal acc)
-                  oList
-                --print oSpends
              draft <-
                liftIO $
                makeOutgoing
@ -1022,11 +963,13 @@ prepareTxV2 pool zebraHost zebraPort zn za bh pnotes policy = do
                        liftIO $
                        prepSSpends
                          (getSapSK $ zcashAccountSapSpendKey $ entityVal acc)
+                          (maybe InvalidTree fst sTree)
                          sList1
                      oSpends1 <-
                        liftIO $
                        prepOSpends
                          (getOrchSK $ zcashAccountOrchSpendKey $ entityVal acc)
+                          (maybe InvalidTree fst oTree)
                          oList1
                      let noteTotal1 = getTotalAmount (tList1, sList1, oList1)
                      outgoing' <-
@ -1043,8 +986,14 @@ prepareTxV2 pool zebraHost zebraPort zn za bh pnotes policy = do
                          tx <-
                            liftIO $
                            createTransaction
-                              (Just sT)
-                              (Just oT)
+                              (maybe
+                                 (hexString "00")
+                                 (getHash . value . fst)
+                                 sTree)
+                              (maybe
+                                 (hexString "00")
+                                 (getHash . value . fst)
+                                 oTree)
                              tSpends1
                              sSpends1
                              oSpends1
@ -1300,9 +1249,16 @@ prepareTxV2 pool zebraHost zebraPort zn za bh pnotes policy = do
                       (fromIntegral $ walletTrNoteValue $ entityVal n)
                       (walletTrNoteScript $ entityVal n))
    prepSSpends ::
-         SaplingSpendingKey -> [Entity WalletSapNote] -> IO [SaplingTxSpend]
-    prepSSpends sk notes = do
+         SaplingSpendingKey
+      -> Tree SaplingNode
+      -> [Entity WalletSapNote]
+      -> IO [SaplingTxSpend]
+    prepSSpends sk tree notes = do
      forM notes $ \n -> do
+        let notePath =
+              Zenith.Tree.path
+                (fromIntegral $ walletSapNotePosition $ entityVal n)
+                tree
        return $
          SaplingTxSpend
            (getBytes sk)
@ -1313,11 +1269,18 @@ prepareTxV2 pool zebraHost zebraPort zn za bh pnotes policy = do
               (getHex $ walletSapNoteNullifier $ entityVal n)
               ""
               (getRseed $ walletSapNoteRseed $ entityVal n))
-            (toBytes $ getHex $ walletSapNoteWitness $ entityVal n)
+            (fromMaybe nullPath notePath)
    prepOSpends ::
-         OrchardSpendingKey -> [Entity WalletOrchNote] -> IO [OrchardTxSpend]
-    prepOSpends sk notes = do
+         OrchardSpendingKey
+      -> Tree OrchardNode
+      -> [Entity WalletOrchNote]
+      -> IO [OrchardTxSpend]
+    prepOSpends sk tree notes = do
      forM notes $ \n -> do
+        let notePath =
+              Zenith.Tree.path
+                (fromIntegral $ walletOrchNotePosition $ entityVal n)
+                tree
        return $
          OrchardTxSpend
            (getBytes sk)
@ -1328,21 +1291,7 @@ prepareTxV2 pool zebraHost zebraPort zn za bh pnotes policy = do
               (getHex $ walletOrchNoteNullifier $ entityVal n)
               (walletOrchNoteRho $ entityVal n)
               (getRseed $ walletOrchNoteRseed $ entityVal n))
-            (toBytes $ getHex $ walletOrchNoteWitness $ entityVal n)
-    sapAnchor :: [Entity WalletSapNote] -> Maybe SaplingWitness
-    sapAnchor notes =
-      if not (null notes)
-        then Just $
-             SaplingWitness $
-             getHex $ walletSapNoteWitness $ entityVal $ head notes
-        else Nothing
-    orchAnchor :: [Entity WalletOrchNote] -> Maybe OrchardWitness
-    orchAnchor notes =
-      if not (null notes)
-        then Just $
-             OrchardWitness $
-             getHex $ walletOrchNoteWitness $ entityVal $ head notes
-        else Nothing
+            (fromMaybe nullPath notePath)

 -- | Sync the wallet with the data store
 syncWallet ::
@ -1388,10 +1337,95 @@ syncWallet config w = do
      (findOrchardActions config startBlock (zcashWalletNetwork $ entityVal w))
      accs
  logDebugN "processed orchard actions"
-  _ <- liftIO $ updateSaplingWitnesses pool
-  logDebugN "updated sapling witnesses"
-  _ <- liftIO $ updateOrchardWitnesses pool
-  logDebugN "updated orchard witnesses"
  _ <- liftIO $ updateWalletSync pool chainTip (entityKey w)
  logDebugN "updated wallet lastSync"
  mapM_ (liftIO . runNoLoggingT . getWalletTransactions pool) addrs
+
+-- | Update commitment trees
+updateCommitmentTrees ::
+     ConnectionPool -> T.Text -> Int -> ZcashNetDB -> LoggingT IO ()
+updateCommitmentTrees pool zHost zPort zNet = do
+  sTdb <- liftIO $ getSaplingTree pool
+  oTdb <- liftIO $ getOrchardTree pool
+  maxBlock <- liftIO $ getMaxBlock pool zNet
+  newSapTree <-
+    case sTdb of
+      Nothing -> do
+        logDebugN ">no Sapling tree in DB"
+        bh <- liftIO $ getMinBirthdayHeight pool zNet
+        logDebugN $ ">min birthday: " <> T.pack (show bh)
+        saplingNotes <- liftIO $ getShieldedOutputs pool (bh + 1) zNet
+        let saplingComm =
+              map
+                (\(_, y) ->
+                   ( getHex $ shieldOutputCmu (entityVal y)
+                   , fromSqlKey (entityKey y)))
+                saplingNotes
+        logDebugN ">got shielded outputs"
+        treeInfo <- liftIO $ getCommitmentTrees pool zHost zPort zNet bh
+        case getSaplingTreeParts (SaplingCommitmentTree $ ztiSapling treeInfo) of
+          Nothing -> do
+            logDebugN ">failed to load tree from Zebra"
+            return InvalidTree
+          Just t1 -> do
+            let newTree = mkSaplingTree t1
+            return $ foldl' append newTree saplingComm
+      Just (sTree, sSync) -> do
+        logDebugN $ ">Sapling tree found, synced to " <> T.pack (show sSync)
+        saplingNotes <- liftIO $ getShieldedOutputs pool (sSync + 1) zNet
+        let saplingComm =
+              map
+                (\(_, y) ->
+                   ( getHex $ shieldOutputCmu (entityVal y)
+                   , fromSqlKey (entityKey y)))
+                saplingNotes
+        logDebugN ">got shielded outputs"
+        return $ foldl' append sTree saplingComm
+  newOrchTree <-
+    case oTdb of
+      Nothing -> do
+        logDebugN ">no Orchard tree in DB"
+        bh <- liftIO $ getMinBirthdayHeight pool zNet
+        logDebugN $ ">min birthday: " <> T.pack (show bh)
+        orchardNotes <- liftIO $ getOrchardActions pool (bh + 1) zNet
+        let orchardComm =
+              map
+                (\(_, y) ->
+                   ( getHex $ orchActionCmx (entityVal y)
+                   , fromSqlKey (entityKey y)))
+                orchardNotes
+        logDebugN ">got orchard actions"
+        treeInfo <- liftIO $ getCommitmentTrees pool zHost zPort zNet bh
+        case getOrchardTreeParts (OrchardCommitmentTree $ ztiOrchard treeInfo) of
+          Nothing -> do
+            logDebugN ">failed to load tree from Zebra"
+            return InvalidTree
+          Just t1 -> do
+            let newTree = mkOrchardTree t1
+            return $ foldl' append newTree orchardComm
+      Just (oTree, oSync) -> do
+        logDebugN $ ">Orchard tree found, synced to " <> T.pack (show oSync)
+        orchardNotes <- liftIO $ getOrchardActions pool (oSync + 1) zNet
+        let orchardComm =
+              map
+                (\(_, y) ->
+                   ( getHex $ orchActionCmx (entityVal y)
+                   , fromSqlKey (entityKey y)))
+                orchardNotes
+        logDebugN ">got orchard actions"
+        return $ foldl' append oTree orchardComm
+  case newSapTree of
+    Branch {} -> do
+      logInfoN ">Saving updated Sapling tree to db"
+      _ <- liftIO $ upsertSaplingTree pool maxBlock newSapTree
+      case newOrchTree of
+        Branch {} -> do
+          logInfoN ">Saving updated Orchard tree to db"
+          _ <- liftIO $ upsertOrchardTree pool maxBlock newOrchTree
+          return ()
+        _anyOther -> do
+          logErrorN ">Failed to update the Orchard tree"
+          return ()
+    _anyOther -> do
+      logErrorN ">Failed to update the Sapling tree"
+      return ()
--- a/src/Zenith/DB.hs
+++ b/src/Zenith/DB.hs
@ -26,6 +26,7 @@ import Control.Monad.Logger
  ( LoggingT
  , NoLoggingT
  , logDebugN
+  , logErrorN
  , runNoLoggingT
  , runStderrLoggingT
  )
@ -81,7 +82,7 @@ import ZcashHaskell.Types
  , ValidAddress(..)
  , ZcashNet(..)
  )
-import Zenith.Tree (OrchardNode(..), SaplingNode(..), Tree(..))
+import Zenith.Tree (OrchardNode(..), SaplingNode(..), Tree(..), truncateTree)
 import Zenith.Types
  ( AccountBalance(..)
  , HexStringDB(..)
@ -927,15 +928,17 @@ getMaxWalletBlock pool = do
    Nothing -> return $ -1
    Just x -> return $ walletTransactionBlock $ entityVal x

-getMinBirthdayHeight :: ConnectionPool -> IO Int
-getMinBirthdayHeight pool = do
+getMinBirthdayHeight :: ConnectionPool -> ZcashNetDB -> IO Int
+getMinBirthdayHeight pool znet = do
  b <-
    runNoLoggingT $
    PS.retryOnBusy $
    flip PS.runSqlPool pool $ do
      selectOne $ do
        w <- from $ table @ZcashWallet
-        where_ (w ^. ZcashWalletBirthdayHeight >. val 0)
+        where_
+          (w ^. ZcashWalletBirthdayHeight >. val 0 &&. w ^. ZcashWalletNetwork ==.
+           val znet)
        orderBy [asc $ w ^. ZcashWalletBirthdayHeight]
        pure w
  case b of
@ -991,14 +994,13 @@ saveWalletTransaction pool za zt = do
 saveWalletSapNote ::
     ConnectionPool -- ^ The database path
  -> WalletTransactionId -- ^ The index for the transaction that contains the note
-  -> Integer -- ^ note position
-  -> SaplingWitness -- ^ the Sapling incremental witness
+  -> Int32 -- ^ note position
  -> Bool -- ^ change flag
  -> ZcashAccountId
  -> ShieldOutputId
  -> DecodedNote -- The decoded Sapling note
  -> IO ()
-saveWalletSapNote pool wId pos wit ch za zt dn = do
+saveWalletSapNote pool wId pos ch za zt dn = do
  runNoLoggingT $
    PS.retryOnBusy $
    flip PS.runSqlPool pool $ do
@ -1013,7 +1015,7 @@ saveWalletSapNote pool wId pos wit ch za zt dn = do
             False
             (HexStringDB $ a_nullifier dn)
             (fromIntegral pos)
-             (HexStringDB $ sapWit wit)
+             (HexStringDB $ hexString "00")
             ch
             zt
             (RseedDB $ a_rseed dn))
@ -1024,14 +1026,13 @@ saveWalletSapNote pool wId pos wit ch za zt dn = do
 saveWalletOrchNote ::
     ConnectionPool
  -> WalletTransactionId
-  -> Integer
-  -> OrchardWitness
+  -> Int32
  -> Bool
  -> ZcashAccountId
  -> OrchActionId
  -> DecodedNote
  -> IO ()
-saveWalletOrchNote pool wId pos wit ch za zt dn = do
+saveWalletOrchNote pool wId pos ch za zt dn = do
  runNoLoggingT $
    PS.retryOnBusy $
    flip PS.runSqlPool pool $ do
@ -1046,7 +1047,7 @@ saveWalletOrchNote pool wId pos wit ch za zt dn = do
             False
             (HexStringDB $ a_nullifier dn)
             (fromIntegral pos)
-             (HexStringDB $ orchWit wit)
+             (HexStringDB $ hexString "00")
             ch
             zt
             (a_rho dn)
@ -2167,6 +2168,9 @@ clearWalletData pool = do
  runNoLoggingT $
    PS.retryOnBusy $
    flip PS.runSqlPool pool $ do
+      delete $ do
+        _ <- from $ table @TreeStore
+        return ()
      delete $ do
        _ <- from $ table @TransparentNote
        return ()
@ -2820,10 +2824,79 @@ rewindWalletData pool b net = do
    flip PS.runSqlPool pool $ do
      delete $ do
        blk <- from $ table @ZcashBlock
-        where_
-          (blk ^. ZcashBlockHeight >. val b &&. blk ^. ZcashBlockNetwork ==.
-           val net)
+        where_ $
+          (blk ^. ZcashBlockHeight >. val b) &&.
+          (blk ^. ZcashBlockNetwork ==. val net)
  logDebugN "Completed data store rewind"
+  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
+      _ <- 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
+          _ <- liftIO $ upsertOrchardTree pool b truncOrchTree
+          logDebugN $ "Truncated Orchard tree at index " <> T.pack (show oaIx)
+
+getSaplingOutputAtBlock ::
+     ConnectionPool -> ZcashNetDB -> Int -> IO (Maybe Int64)
+getSaplingOutputAtBlock pool znet b = do
+  r <-
+    runNoLoggingT $
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      selectOne $ do
+        (blks :& txs :& sOutputs) <-
+          from $ table @ZcashBlock `innerJoin` table @ZcashTransaction `on`
+          (\(blks :& txs) ->
+             blks ^. ZcashBlockId ==. txs ^. ZcashTransactionBlockId) `innerJoin`
+          table @ShieldOutput `on`
+          (\(_ :& txs :& sOutputs) ->
+             txs ^. ZcashTransactionId ==. sOutputs ^. ShieldOutputTx)
+        where_ (blks ^. ZcashBlockHeight <=. val b)
+        where_ (blks ^. ZcashBlockNetwork ==. val znet)
+        orderBy
+          [ desc $ txs ^. ZcashTransactionId
+          , desc $ sOutputs ^. ShieldOutputPosition
+          ]
+        return sOutputs
+  case r of
+    Nothing -> return Nothing
+    Just so -> return $ Just $ fromSqlKey $ entityKey so
+
+getOrchardActionAtBlock ::
+     ConnectionPool -> ZcashNetDB -> Int -> IO (Maybe Int64)
+getOrchardActionAtBlock pool znet b = do
+  r <-
+    runNoLoggingT $
+    PS.retryOnBusy $
+    flip PS.runSqlPool pool $ do
+      selectOne $ do
+        (blks :& txs :& oActions) <-
+          from $ table @ZcashBlock `innerJoin` table @ZcashTransaction `on`
+          (\(blks :& txs) ->
+             blks ^. ZcashBlockId ==. txs ^. ZcashTransactionBlockId) `innerJoin`
+          table @OrchAction `on`
+          (\(_ :& txs :& oActions) ->
+             txs ^. ZcashTransactionId ==. oActions ^. OrchActionTx)
+        where_ (blks ^. ZcashBlockHeight <=. val b)
+        where_ (blks ^. ZcashBlockNetwork ==. val znet)
+        orderBy
+          [ desc $ txs ^. ZcashTransactionId
+          , desc $ oActions ^. OrchActionPosition
+          ]
+        return oActions
+  case r of
+    Nothing -> return Nothing
+    Just so -> return $ Just $ fromSqlKey $ entityKey so

 -- * Tree storage
 -- | Read the Orchard commitment tree
--- a/src/Zenith/GUI.hs
+++ b/src/Zenith/GUI.hs
@ -1625,7 +1625,7 @@ scanZebra :: T.Text -> T.Text -> Int -> ZcashNet -> (AppEvent -> IO ()) -> IO ()
 scanZebra dbPath zHost zPort net sendMsg = do
  bStatus <- liftIO $ checkBlockChain zHost zPort
  pool <- runNoLoggingT $ initPool dbPath
-  b <- liftIO $ getMinBirthdayHeight pool
+  b <- liftIO $ getMinBirthdayHeight pool $ ZcashNetDB net
  dbBlock <- getMaxBlock pool $ ZcashNetDB net
  chkBlock <- checkIntegrity dbPath zHost zPort net dbBlock 1
  syncChk <- isSyncing pool
--- a/src/Zenith/RPC.hs
+++ b/src/Zenith/RPC.hs
@ -887,7 +887,7 @@ scanZebra :: T.Text -> T.Text -> Int -> ZcashNet -> IO ()
 scanZebra dbPath zHost zPort net = do
  bStatus <- checkBlockChain zHost zPort
  pool <- runNoLoggingT $ initPool dbPath
-  b <- getMinBirthdayHeight pool
+  b <- getMinBirthdayHeight pool $ ZcashNetDB net
  dbBlock <- getMaxBlock pool $ ZcashNetDB net
  chkBlock <- checkIntegrity dbPath zHost zPort net dbBlock 1
  syncChk <- isSyncing pool
--- a/src/Zenith/Scanner.hs
+++ b/src/Zenith/Scanner.hs
@ -31,7 +31,7 @@ import ZcashHaskell.Types
  , fromRawTBundle
  )
 import ZcashHaskell.Utils (getBlockTime, makeZebraCall, readZebraTransaction)
-import Zenith.Core (checkBlockChain, syncWallet)
+import Zenith.Core (checkBlockChain, syncWallet, updateCommitmentTrees)
 import Zenith.DB
  ( ZcashBlock(..)
  , ZcashBlockId
@ -83,7 +83,7 @@ rescanZebra host port dbFilePath = do
      clearWalletData pool1
      _ <- startSync pool1
      dbBlock <- getMaxBlock pool1 znet
-      b <- liftIO $ getMinBirthdayHeight pool1
+      b <- liftIO $ getMinBirthdayHeight pool1 znet
      let sb = max dbBlock b
      if sb > zgb_blocks bStatus || sb < 1
        then liftIO $ throwIO $ userError "Invalid starting block for scan"
@ -108,6 +108,7 @@ rescanZebra host port dbFilePath = do
                {-mapM_ (processBlock host port pool3 pg3 znet) bl3-}
              print "Please wait..."
          _ <- completeSync pool1 Successful
+          _ <- runStderrLoggingT $ updateCommitmentTrees pool1 host port znet
          print "Rescan complete"

 -- | Function to process a raw block and extract the transaction information
--- a/src/Zenith/Tree.hs
+++ b/src/Zenith/Tree.hs
@ -103,7 +103,7 @@ instance (Monoid v, Node v) => Semigroup (Tree v) where
    | getLevel s == getLevel w = branch (Branch s x y) (PrunedBranch w)
    | otherwise = InvalidTree
  (<>) (Branch s x y) (Leaf w)
-    | isFull s = Branch s x y <> mkSubTree (getLevel s) (Leaf w)
+    | isFull s = InvalidTree
    | isFull (value x) = branch x (y <> Leaf w)
    | otherwise = branch (x <> Leaf w) y
  (<>) (Branch s x y) (Branch s1 x1 y1)
@ -166,6 +166,9 @@ path pos (Branch s x y) =
      | otherwise = []
 path _ _ = Nothing

+nullPath :: MerklePath
+nullPath = MerklePath 0 []
+
 getNotePosition :: Monoid v => Node v => Tree v -> Int64 -> Maybe Position
 getNotePosition (Leaf x) i
  | getIndex x == i = Just $ getPosition x
--- a/test/Spec.hs
+++ b/test/Spec.hs
@ -7,6 +7,7 @@ import Data.Aeson
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Lazy as LBS
 import Data.HexString
+import Data.List (foldl')
 import Data.Maybe (fromJust)
 import qualified Data.Text.Encoding as E
 import Database.Persist
@ -639,6 +640,23 @@ main = do
              let updatedTree = foldl append newTree [(cmx1, 4), (cmx2, 5)]
              let truncTree = truncateTree updatedTree 4
              getIndex (value truncTree) `shouldBe` 4
+        it "Validate tree from DB" $ do
+          pool <- runNoLoggingT $ initPool "/home/rav/Zenith/zenith.db"
+          dbTree <- getOrchardTree pool
+          case dbTree of
+            Nothing -> assertFailure "failed to get tree from DB"
+            Just (oTree, oSync) -> do
+              zebraTrees <-
+                getCommitmentTrees
+                  pool
+                  "localhost"
+                  18232
+                  (ZcashNetDB TestNet)
+                  oSync
+              let finalAnchor =
+                    getOrchardTreeAnchor $
+                    OrchardCommitmentTree $ ztiOrchard zebraTrees
+              getHash (value oTree) `shouldBe` finalAnchor
    describe "Creating Tx" $ do
      describe "Full" $ do
        it "To Orchard" $ do
--- a/2
+++ b/2
@ -1 +1 @@
-Subproject commit 20851a4e48f768a492796fb828f16ae9745931dc
+Subproject commit dea960c2acf7479eeb42845c07b482449d538aae