[ZCash ] Implement shard tree and shard storage

components/brave_wallet/browser/BUILD.gn

@@ -314,14 +314,15 @@ 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_orchard_sync_state.cc",
+      "zcash/zcash_orchard_sync_state.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,29 @@ source_set("hd_key") {
 }
 
 if (enable_orchard) {
+  source_set("test_support") {
+    sources = [
+      "orchard_test_utils.cc",
+      "orchard_test_utils.h",
+    ]
+    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_shard_tree_manager.cc",
+      "orchard_shard_tree_manager.h",
+    ]
+    deps = [
+      "orchard_storage:headers",
+      "//brave/components/brave_wallet/browser/zcash/rust",
     ]
-    deps = [ "//brave/components/brave_wallet/browser/zcash/rust" ]
   }
 }

components/brave_wallet/browser/internal/orchard_block_scanner.cc

@@ -5,19 +5,24 @@
 
 #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_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;
 
@@ -29,50 +34,46 @@ 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 =
+      decoder_->ScanBlocks(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,7 +12,9 @@
 #include <vector>
 
 #include "base/types/expected.h"
+#include "brave/components/brave_wallet/browser/internal/orchard_block_scanner.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"
 #include "brave/components/brave_wallet/common/zcash_utils.h"
 
 namespace brave_wallet {
@@ -21,20 +23,25 @@ namespace brave_wallet {
 // 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,8 +50,8 @@ 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_;

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,7 +22,7 @@ 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;
   }

components/brave_wallet/browser/internal/orchard_bundle_manager.h

@@ -45,7 +45,7 @@ class OrchardBundleManager {
   // Returns in unauthorized state
   static std::unique_ptr<OrchardBundleManager> Create(
       base::span<const uint8_t> tree_state,
-      const std::vector<::brave_wallet::OrchardOutput>& orchard_outputs);
+      const std::vector<OrchardOutput>& orchard_outputs);
 
   static void OverrideRandomSeedForTesting(size_t seed) {
     random_seed_for_testing_ = seed;

components/brave_wallet/browser/internal/orchard_shard_tree_manager.cc

@@ -0,0 +1,69 @@
+/* Copyright (c) 2024 The Brave Authors. All rights reserved.
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+#include "brave/components/brave_wallet/browser/internal/orchard_shard_tree_manager.h"
+
+#include "brave/components/brave_wallet/browser/internal/orchard_storage/orchard_shard_tree_delegate.h"
+#include "brave/components/brave_wallet/browser/zcash/rust/orchard_shard_tree.h"
+
+namespace brave_wallet {
+
+// static
+std::unique_ptr<OrchardShardTreeManager> OrchardShardTreeManager::Create(
+    std::unique_ptr<OrchardShardTreeDelegate> delegate) {
+  auto shard_tree = orchard::OrchardShardTree::Create(std::move(delegate));
+  if (!shard_tree) {
+    return nullptr;
+  }
+  return std::make_unique<OrchardShardTreeManager>(std::move(shard_tree));
+}
+
+// static
+std::unique_ptr<OrchardShardTreeManager>
+OrchardShardTreeManager::CreateForTesting(
+    std::unique_ptr<OrchardShardTreeDelegate> delegate) {
+  auto shard_tree =
+      orchard::OrchardShardTree::CreateForTesting(std::move(delegate));
+  if (!shard_tree) {
+    return nullptr;
+  }
+  return std::make_unique<OrchardShardTreeManager>(std::move(shard_tree));
+}
+
+OrchardShardTreeManager::OrchardShardTreeManager(
+    std::unique_ptr<::brave_wallet::orchard::OrchardShardTree> shard_tree) {
+  orchard_shard_tree_ = std::move(shard_tree);
+}
+
+OrchardShardTreeManager::~OrchardShardTreeManager() = default;
+
+bool OrchardShardTreeManager::InsertCommitments(
+    OrchardBlockScanner::Result&& result) {
+  return orchard_shard_tree_->ApplyScanResults(
+      std::move(result.scanned_blocks));
+}
+
+base::expected<std::vector<OrchardInput>, std::string>
+OrchardShardTreeManager::CalculateWitness(
+    const std::vector<OrchardInput>& notes,
+    uint32_t checkpoint_position) {
+  std::vector<OrchardInput> result;
+  for (auto& input : notes) {
+    auto witness = orchard_shard_tree_->CalculateWitness(
+        input.note.orchard_commitment_tree_position, checkpoint_position);
+    if (!witness.has_value()) {
+      return base::unexpected(witness.error());
+    }
+    result.push_back(input);
+    result.back().witness = witness.value();
+  }
+  return base::ok(std::move(result));
+}
+
+bool OrchardShardTreeManager::Truncate(uint32_t checkpoint) {
+  return orchard_shard_tree_->TruncateToCheckpoint(checkpoint);
+}
+
+}  // namespace brave_wallet