feat(tui): implement new commitment trees

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 ()
+        MsgDisplay -> do
-        PhraseDisplay -> return ()
+          if m == "startSync"
-        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
               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)

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)
+  sT <- liftIO $ getSaplingTree pool
-  logDebugN "getting Sapling frontier"
-  let sT = getSaplingFrontier $ SaplingCommitmentTree $ ztiSapling trees
   case sT of
     Nothing ->
       liftIO $ throwIO $ userError "Failed to read Sapling commitment tree"
-    Just sT' -> do
+    Just (sT', treeSync) -> do
-      logDebugN "Sapling frontier valid"
+      logDebugN "Sapling tree valid"
-      decryptNotes sT' zn pool tList
+      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) = do
-    decryptNotes st n pool ((zt, o):txs) = do
+      case getNotePosition st $ fromSqlKey $ entityKey o of
-      let updatedTree =
+        Nothing -> do
-            updateSaplingCommitmentTree
+          logErrorN "Couldn't find sapling note in commitment tree"
-              st
+          return ()
-              (getHex $ shieldOutputCmu $ entityVal o)
+        Just nP -> do
-      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
           logDebugN "got position"
-          case notePos of
+          case decodeShOut External n nP o of
-            Nothing ->
+            Nothing -> do
-              liftIO $ throwIO $ userError "Failed to obtain note position"
+              logDebugN "couldn't decode external"
-            Just nP -> do
+              case decodeShOut Internal n nP o of
-              case decodeShOut External n nP o of
                 Nothing -> do
-                  logDebugN "couldn't decode external"
+                  logDebugN "couldn't decode internal"
-                  case decodeShOut Internal n nP o of
+                Just dn1 -> do
-                    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
                   wId <- liftIO $ saveWalletTransaction pool (entityKey za) zt
                   liftIO $
                     saveWalletSapNote
                       pool
                       wId
                       nP
-                      (fromJust noteWitness)
+                      True
-                      False
                       (entityKey za)
                       (entityKey o)
-                      dn0
+                      dn1
-                  decryptNotes uT n pool txs
+            Just dn0 -> do
+              wId <- liftIO $ saveWalletTransaction pool (entityKey za) zt
+              liftIO $
+                saveWalletSapNote
+                  pool
+                  wId
+                  nP
+                  False
+                  (entityKey za)
+                  (entityKey o)
+                  dn0
     decodeShOut ::
-         Scope
+         Scope -> ZcashNet -> Int32 -> Entity ShieldOutput -> Maybe DecodedNote
-      -> ZcashNet
-      -> Integer
-      -> 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)
+  sT <- getOrchardTree pool
-  let sT = getOrchardFrontier $ OrchardCommitmentTree $ ztiOrchard trees
   case sT of
     Nothing -> throwIO $ userError "Failed to read Orchard commitment tree"
-    Just sT' -> do
+    Just (sT', treeSync) -> do
-      decryptNotes sT' zn pool tList
+      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) = do
-    decryptNotes ot n pool ((zt, o):txs) = do
+      case getNotePosition ot (fromSqlKey $ entityKey o) of
-      let updatedTree =
+        Nothing -> do
-            updateOrchardCommitmentTree
+          return ()
-              ot
+        Just nP ->
-              (getHex $ orchActionCmx $ entityVal o)
+          case decodeOrchAction External nP o of
-      case updatedTree of
+            Nothing ->
-        Nothing -> throwIO $ userError "Failed to update commitment tree"
+              case decodeOrchAction Internal nP o of
-        Just uT -> do
+                Nothing -> return ()
-          let noteWitness = getOrchardWitness uT
+                Just dn1 -> do
-          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
                   wId <- saveWalletTransaction pool (entityKey za) zt
                   saveWalletOrchNote
                     pool
                     wId
                     nP
-                    (fromJust noteWitness)
+                    True
-                    False
                     (entityKey za)
                     (entityKey o)
-                    dn
+                    dn1
-                  decryptNotes uT n pool txs
+            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 ::
+    decodeOrchAction :: Scope -> Int32 -> Entity OrchAction -> Maybe DecodedNote
-         Scope -> Integer -> 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
+            (maybe (hexString "00") (getHash . value . fst) sTree)
-            Nothing
+            (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 <-
+  sTree <- liftIO $ getSaplingTree pool
-    liftIO $ getCommitmentTrees pool zebraHost zebraPort (ZcashNetDB zn) bh
+  oTree <- liftIO $ getOrchardTree pool
-  let sT = SaplingCommitmentTree $ ztiSapling trees
-  let oT = OrchardCommitmentTree $ ztiOrchard trees
   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)
+                              (maybe
-                              (Just oT)
+                                 (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]
+         SaplingSpendingKey
-    prepSSpends sk notes = do
+      -> 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]
+         OrchardSpendingKey
-    prepOSpends sk notes = do
+      -> 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)
+            (fromMaybe nullPath notePath)
-    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
 
 -- | 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 ()

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 :: ConnectionPool -> ZcashNetDB -> IO Int
-getMinBirthdayHeight pool = do
+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
+  -> Int32 -- ^ note position
-  -> SaplingWitness -- ^ the Sapling incremental witness
   -> 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
+  -> Int32
-  -> OrchardWitness
   -> 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_
+        where_ $
-          (blk ^. ZcashBlockHeight >. val b &&. blk ^. ZcashBlockNetwork ==.
+          (blk ^. ZcashBlockHeight >. val b) &&.
-           val net)
+          (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

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

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

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

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

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

zcash-haskell

@@ -1 +1 @@
-Subproject commit 20851a4e48f768a492796fb828f16ae9745931dc
+Subproject commit dea960c2acf7479eeb42845c07b482449d538aae