[ZCash ] Implement shard tree and shard storage (#26477)

* Implement shard tree and shard storage
Resolves brave/brave-browser#39314

components/brave_wallet/browser/BUILD.gn

@@ -314,14 +314,13 @@ static_library("browser") {
     sources += [
       "zcash/zcash_create_shield_transaction_task.cc",
       "zcash/zcash_create_shield_transaction_task.h",
-      "zcash/zcash_orchard_storage.cc",
-      "zcash/zcash_orchard_storage.h",
       "zcash/zcash_shield_sync_service.cc",
       "zcash/zcash_shield_sync_service.h",
     ]
 
     deps += [
       "internal:orchard_bundle",
+      "internal/orchard_storage:orchard_storage",
       "//sql",
     ]
   }

components/brave_wallet/browser/internal/BUILD.gn

@@ -39,13 +39,32 @@ source_set("hd_key") {
 }
 
 if (enable_orchard) {
+  source_set("test_support") {
+    sources = [
+      "orchard_test_utils.cc",
+      "orchard_test_utils.h",
+    ]
+    public_deps = [
+      "//brave/components/brave_wallet/browser/zcash/rust:test_support_headers",
+    ]
+    deps = [
+      ":orchard_bundle",
+      "//brave/components/brave_wallet/browser/zcash/rust:test_support",
+    ]
+  }
+
   source_set("orchard_bundle") {
     sources = [
       "orchard_block_scanner.cc",
       "orchard_block_scanner.h",
       "orchard_bundle_manager.cc",
       "orchard_bundle_manager.h",
+      "orchard_sync_state.cc",
+      "orchard_sync_state.h",
+    ]
+    deps = [
+      "orchard_storage",
+      "//brave/components/brave_wallet/browser/zcash/rust",
     ]
-    deps = [ "//brave/components/brave_wallet/browser/zcash/rust" ]
   }
 }

components/brave_wallet/browser/internal/hd_key_zip32.cc

@@ -8,35 +8,35 @@
 #include <utility>
 
 #include "base/memory/ptr_util.h"
-#include "brave/components/brave_wallet/browser/zcash/rust/extended_spending_key.h"
+#include "brave/components/brave_wallet/browser/zcash/rust/orchard_extended_spending_key.h"
 
 namespace brave_wallet {
 
-HDKeyZip32::HDKeyZip32(std::unique_ptr<orchard::ExtendedSpendingKey> esk)
-    : extended_spending_key_(std::move(esk)) {}
+HDKeyZip32::HDKeyZip32(std::unique_ptr<orchard::OrchardExtendedSpendingKey> esk)
+    : orchard_extended_spending_key_(std::move(esk)) {}
 
 HDKeyZip32::~HDKeyZip32() = default;
 
 // static
 std::unique_ptr<HDKeyZip32> HDKeyZip32::GenerateFromSeed(
     base::span<const uint8_t> seed) {
-  return base::WrapUnique(
-      new HDKeyZip32(orchard::ExtendedSpendingKey::GenerateFromSeed(seed)));
+  return base::WrapUnique(new HDKeyZip32(
+      orchard::OrchardExtendedSpendingKey::GenerateFromSeed(seed)));
 }
 
 std::unique_ptr<HDKeyZip32> HDKeyZip32::DeriveHardenedChild(uint32_t index) {
-  return base::WrapUnique(
-      new HDKeyZip32(extended_spending_key_->DeriveHardenedChild(index)));
+  return base::WrapUnique(new HDKeyZip32(
+      orchard_extended_spending_key_->DeriveHardenedChild(index)));
 }
 
 std::optional<OrchardAddrRawPart> HDKeyZip32::GetDiversifiedAddress(
     uint32_t div_index,
     OrchardAddressKind kind) {
-  return extended_spending_key_->GetDiversifiedAddress(div_index, kind);
+  return orchard_extended_spending_key_->GetDiversifiedAddress(div_index, kind);
 }
 
 OrchardFullViewKey HDKeyZip32::GetFullViewKey() {
-  return extended_spending_key_->GetFullViewKey();
+  return orchard_extended_spending_key_->GetFullViewKey();
 }
 
 }  // namespace brave_wallet

components/brave_wallet/browser/internal/hd_key_zip32.h

@@ -14,7 +14,7 @@
 namespace brave_wallet {
 
 namespace orchard {
-class ExtendedSpendingKey;
+class OrchardExtendedSpendingKey;
 }  // namespace orchard
 
 // Implements Orchard key generation from
@@ -43,10 +43,11 @@ class HDKeyZip32 {
   OrchardFullViewKey GetFullViewKey();
 
  private:
-  explicit HDKeyZip32(std::unique_ptr<orchard::ExtendedSpendingKey> key);
+  explicit HDKeyZip32(std::unique_ptr<orchard::OrchardExtendedSpendingKey> key);
   // Extended spending key is a root key of an account, all other keys can be
   // derived from esk
-  std::unique_ptr<orchard::ExtendedSpendingKey> extended_spending_key_;
+  std::unique_ptr<orchard::OrchardExtendedSpendingKey>
+      orchard_extended_spending_key_;
 };
 
 }  // namespace brave_wallet

components/brave_wallet/browser/internal/orchard_block_scanner.cc

@@ -5,74 +5,75 @@
 
 #include "brave/components/brave_wallet/browser/internal/orchard_block_scanner.h"
 
+#include "base/threading/thread_restrictions.h"
+#include "brave/components/brave_wallet/browser/zcash/rust/orchard_block_decoder.h"
+#include "brave/components/brave_wallet/browser/zcash/rust/orchard_decoded_blocks_bundle.h"
+
 namespace brave_wallet {
 
 OrchardBlockScanner::Result::Result() = default;
 
-OrchardBlockScanner::Result::Result(std::vector<OrchardNote> discovered_notes,
-                                    std::vector<OrchardNoteSpend> spent_notes)
+OrchardBlockScanner::Result::Result(
+    std::vector<OrchardNote> discovered_notes,
+    std::vector<OrchardNoteSpend> spent_notes,
+    std::unique_ptr<orchard::OrchardDecodedBlocksBundle> scanned_blocks)
     : discovered_notes(std::move(discovered_notes)),
-      spent_notes(std::move(spent_notes)) {}
-
-OrchardBlockScanner::Result::Result(const Result&) = default;
+      found_spends(std::move(spent_notes)),
+      scanned_blocks(std::move(scanned_blocks)) {}
 
+OrchardBlockScanner::Result::Result(OrchardBlockScanner::Result&&) = default;
 OrchardBlockScanner::Result& OrchardBlockScanner::Result::operator=(
-    const Result&) = default;
+    OrchardBlockScanner::Result&&) = default;
 
 OrchardBlockScanner::Result::~Result() = default;
 
-OrchardBlockScanner::OrchardBlockScanner(
-    const OrchardFullViewKey& full_view_key)
-    : decoder_(orchard::OrchardBlockDecoder::FromFullViewKey(full_view_key)) {}
+OrchardBlockScanner::OrchardBlockScanner(const OrchardFullViewKey& fvk)
+    : fvk_(fvk) {}
 
 OrchardBlockScanner::~OrchardBlockScanner() = default;
 
 base::expected<OrchardBlockScanner::Result, OrchardBlockScanner::ErrorCode>
 OrchardBlockScanner::ScanBlocks(
-    std::vector<OrchardNote> known_notes,
-    std::vector<zcash::mojom::CompactBlockPtr> blocks) {
+    const OrchardTreeState& tree_state,
+    const std::vector<zcash::mojom::CompactBlockPtr>& blocks) {
+  base::AssertLongCPUWorkAllowed();
+
+  std::unique_ptr<orchard::OrchardDecodedBlocksBundle> result =
+      orchard::OrchardBlockDecoder::DecodeBlocks(fvk_, tree_state, blocks);
+  if (!result) {
+    return base::unexpected(ErrorCode::kInputError);
+  }
+
+  std::optional<std::vector<OrchardNote>> found_notes =
+      result->GetDiscoveredNotes();
+
+  if (!found_notes) {
+    return base::unexpected(ErrorCode::kDiscoveredNotesError);
+  }
+
   std::vector<OrchardNoteSpend> found_spends;
-  std::vector<OrchardNote> found_notes;
 
   for (const auto& block : blocks) {
-    // Scan block using the decoder initialized with the provided fvk
-    // to find new spendable notes.
-    auto scan_result = decoder_->ScanBlock(block);
-    if (!scan_result) {
-      return base::unexpected(ErrorCode::kDecoderError);
-    }
-    found_notes.insert(found_notes.end(), scan_result->begin(),
-                       scan_result->end());
-    // Place found notes to the known notes list so we can also check for
-    // nullifiers
-    known_notes.insert(known_notes.end(), scan_result->begin(),
-                       scan_result->end());
     for (const auto& tx : block->vtx) {
       // We only scan orchard actions here
       for (const auto& orchard_action : tx->orchard_actions) {
         if (orchard_action->nullifier.size() != kOrchardNullifierSize) {
           return base::unexpected(ErrorCode::kInputError);
         }
 
-        std::array<uint8_t, kOrchardNullifierSize> action_nullifier;
-        base::ranges::copy(orchard_action->nullifier, action_nullifier.begin());
-
         // Nullifier is a public information about some note being spent.
-        // Here we are trying to find a known spendable notes which nullifier
-        // matches nullifier from the processed transaction.
-        if (std::find_if(known_notes.begin(), known_notes.end(),
-                         [&action_nullifier](const auto& v) {
-                           return v.nullifier == action_nullifier;
-                         }) != known_notes.end()) {
-          OrchardNoteSpend spend;
-          spend.block_id = block->height;
-          spend.nullifier = action_nullifier;
-          found_spends.push_back(std::move(spend));
-        }
+        // Here we are collecting nullifiers from the blocks to check them
+        // later.
+        OrchardNoteSpend spend;
+        base::span(spend.nullifier).copy_from(orchard_action->nullifier);
+        spend.block_id = block->height;
+        found_spends.push_back(std::move(spend));
       }
     }
   }
-  return Result({std::move(found_notes), std::move(found_spends)});
+
+  return Result({std::move(found_notes.value()), std::move(found_spends),
+                 std::move(result)});
 }
 
 }  // namespace brave_wallet

components/brave_wallet/browser/internal/orchard_block_scanner.h

@@ -12,29 +12,36 @@
 #include <vector>
 
 #include "base/types/expected.h"
-#include "brave/components/brave_wallet/browser/zcash/rust/orchard_block_decoder.h"
+#include "brave/components/brave_wallet/browser/internal/orchard_block_scanner.h"
+#include "brave/components/brave_wallet/browser/zcash/rust/orchard_decoded_blocks_bundle.h"
 #include "brave/components/brave_wallet/common/zcash_utils.h"
+#include "brave/components/services/brave_wallet/public/mojom/zcash_decoder.mojom.h"
 
 namespace brave_wallet {
 
 // Scans a bunch of blocks with the provided full view key to find
 // spendable notes related to the account.
 class OrchardBlockScanner {
  public:
-  enum class ErrorCode { kInputError, kDecoderError };
+  enum class ErrorCode { kInputError, kDiscoveredNotesError, kDecoderError };
 
   struct Result {
     Result();
     Result(std::vector<OrchardNote> discovered_notes,
-           std::vector<OrchardNoteSpend> spent_notes);
-    Result(const Result&);
-    Result& operator=(const Result&);
+           std::vector<OrchardNoteSpend> spent_notes,
+           std::unique_ptr<orchard::OrchardDecodedBlocksBundle> scanned_blocks);
+    Result(const Result&) = delete;
+    Result& operator=(const Result&) = delete;
+    Result(Result&&);
+    Result& operator=(Result&&);
     ~Result();
 
-    // New notes have been discovered
+    // New notes have been discovered.
     std::vector<OrchardNote> discovered_notes;
-    // Nullifiers for the previously discovered notes
-    std::vector<OrchardNoteSpend> spent_notes;
+    // Nullifiers for the previously discovered notes.
+    std::vector<OrchardNoteSpend> found_spends;
+    // Decoded blocks bundle to be insterted in the shard tree.
+    std::unique_ptr<orchard::OrchardDecodedBlocksBundle> scanned_blocks;
   };
 
   explicit OrchardBlockScanner(const OrchardFullViewKey& full_view_key);
@@ -43,11 +50,11 @@ class OrchardBlockScanner {
   // Scans blocks to find incoming notes related to fvk
   // Also checks whether existing notes were spent.
   virtual base::expected<Result, OrchardBlockScanner::ErrorCode> ScanBlocks(
-      std::vector<OrchardNote> known_notes,
-      std::vector<zcash::mojom::CompactBlockPtr> blocks);
+      const OrchardTreeState& tree_state,
+      const std::vector<zcash::mojom::CompactBlockPtr>& blocks);
 
  private:
-  std::unique_ptr<orchard::OrchardBlockDecoder> decoder_;
+  OrchardFullViewKey fvk_;
 };
 
 }  // namespace brave_wallet

components/brave_wallet/browser/internal/orchard_block_scanner_unittest.cc

@@ -189,7 +189,7 @@ TEST(OrchardBlockScannerTest, DiscoverNewNotes) {
   EXPECT_EQ(result.value().discovered_notes[3].block_id, 11u);
   EXPECT_EQ(result.value().discovered_notes[3].amount, 2549979667u);
 
-  EXPECT_EQ(result.value().spent_notes.size(), 0u);
+  EXPECT_EQ(result.value().found_spends.size(), 5u);
 }
 
 TEST(OrchardBlockScannerTest, WrongInput) {
@@ -469,11 +469,13 @@ TEST(OrchardBlockScanner, FoundKnownNullifiers_SameBatch) {
   EXPECT_EQ(result.value().discovered_notes[0].block_id, 10u);
   EXPECT_EQ(result.value().discovered_notes[0].amount, 3625561528u);
 
-  EXPECT_EQ(result.value().spent_notes.size(), 1u);
-  EXPECT_EQ(result.value().spent_notes[0].block_id, 11u);
+  EXPECT_EQ(result.value().found_spends.size(), 2u);
+  EXPECT_EQ(result.value().found_spends[0].block_id, 10u);
+  EXPECT_EQ(result.value().found_spends[1].block_id, 11u);
+
   EXPECT_EQ(
-      std::vector<uint8_t>(result.value().spent_notes[0].nullifier.begin(),
-                           result.value().spent_notes[0].nullifier.end()),
+      std::vector<uint8_t>(result.value().found_spends[1].nullifier.begin(),
+                           result.value().found_spends[1].nullifier.end()),
       PrefixedHexStringToBytes(
           "0x6588cc7fabfab2b2a4baa89d4dfafaa50cc89d22f96d10fb7689461b921ad40d")
           .value());
@@ -525,11 +527,13 @@ TEST(OrchardBlockScanner, FoundKnownNullifiers) {
   notes.push_back(note);
   blocks.push_back(std::move(block));
 
-  auto result = scanner.ScanBlocks(std::move(notes), std::move(blocks));
+  OrchardTreeState tree_state;
+
+  auto result = scanner.ScanBlocks(tree_state, std::move(blocks));
 
   EXPECT_TRUE(result.has_value());
-  EXPECT_EQ(result.value().spent_notes.size(), 1u);
-  EXPECT_EQ(result.value().spent_notes[0], spend);
+  EXPECT_EQ(result.value().found_spends.size(), 1u);
+  EXPECT_EQ(result.value().found_spends[0].nullifier, spend.nullifier);
   EXPECT_EQ(result.value().discovered_notes.size(), 0u);
 }
 

components/brave_wallet/browser/internal/orchard_bundle_manager.cc

@@ -22,11 +22,11 @@ std::optional<size_t> OrchardBundleManager::random_seed_for_testing_ =
 // static
 std::unique_ptr<OrchardBundleManager> OrchardBundleManager::Create(
     base::span<const uint8_t> tree_state,
-    const std::vector<::brave_wallet::OrchardOutput>& orchard_outputs) {
+    const std::vector<OrchardOutput>& orchard_outputs) {
   if (orchard_outputs.empty()) {
     return nullptr;
   }
-  auto bundle = orchard::UnauthorizedOrchardBundle::Create(
+  auto bundle = orchard::OrchardUnauthorizedBundle::Create(
       tree_state, orchard_outputs, random_seed_for_testing_);
   if (!bundle) {
     return nullptr;
@@ -36,11 +36,11 @@ std::unique_ptr<OrchardBundleManager> OrchardBundleManager::Create(
 }
 
 OrchardBundleManager::OrchardBundleManager(
-    std::unique_ptr<orchard::UnauthorizedOrchardBundle> unauthorized_bundle)
+    std::unique_ptr<orchard::OrchardUnauthorizedBundle> unauthorized_bundle)
     : unauthorized_orchard_bundle_(std::move(unauthorized_bundle)) {}
 
 OrchardBundleManager::OrchardBundleManager(
-    std::unique_ptr<orchard::AuthorizedOrchardBundle> authorized_bundle)
+    std::unique_ptr<orchard::OrchardAuthorizedBundle> authorized_bundle)
     : authorized_orchard_bundle_(std::move(authorized_bundle)) {}
 
 std::optional<std::array<uint8_t, kZCashDigestSize>>