Repository: zcash/zcash (language: rust)

Issue:
Save roots of complete 2^16 subtrees within the Sapling commitment trees
In order to use `shardtree` efficiently in a mobile client, `lightwalletd` needs to be able to supply it with information about roots of 2<sup>16</sup> subtrees within the commitment tree. We need `zcashd` to store this information.

Currently, `zcashd` stores an `IncrementalMerkleTree` state as of the end of every block. This is a tree frontier, onto which new leaves can be appended to move the frontier forward, but that does not store any information allowing leaf removals.

We only need to store the roots of the 2<sup>16</sup> subtrees:
- Smaller subtrees will be reconstructed by fetching the stream of note commitments within the subtree, which is at most 2MiB of leaf hashes (plus metadata).
- Larger subtree roots can be constructed from the 2<sup>16</sup> subtree roots, which is at most 2MiB of hashes.

Additionally, because this is at most 2MiB of data, and we'll want it for the internal Orchard wallet eventually, we will just collect these roots for all nodes (not just when the `-lightwalletd` option is specified).

- When a block append causes a subtree to be completed, `zcashd` should store:
  - The index of the subtree.
    - Zero-indexed from the left; this can be derived from the position of any leaf within the subtree.
  - The root of the subtree
  - The height of block that completed this subtree
    - Either the position of the frontier at the end of the block is precisely on the 2<sup>16</sup> subtree boundary (the last leaf appended is part of the subtree), or it is somewhere within the next subtree (as a block is not large enough to contain sufficient note commitments to cross multiple 2<sup>16</sup> subtrees).
- When a block rewind uncompletes a subtree, we should remove its root from the coins DB.
- On start, when the coins DB is loaded, if there are no stored roots, scan the commitment tree frontiers stored in the coins DB from Sapling activation, finding the 2<sup>16</sup> roots (by observing when the hash at position 16 changes across a block boundary), and storing them in the database.


Discussion / hints:

Fix patch:
diff --git a/src/coins.h b/src/coins.h
index 57009ddefe6..2b6bfad0e8e 100644
--- a/src/coins.h
+++ b/src/coins.h
@@ -398,11 +398,23 @@ class CCoinsView
     //! Get current history root
     virtual uint256 GetHistoryRoot(uint32_t epochId) const = 0;
 
-    //! Get the largest completed subtree data for the TRACKED_SUBTREE_HEIGHT subtrees known
+    //! Get the largest completed subtree data for the TRACKED_SUBTREE_HEIGHT depth subtrees known
     //! to the node for a given protocol. std::nullopt is returned in the event there are no
     //! complete subtrees.
     virtual std::optional<libzcash::LatestSubtree> GetLatestSubtree(ShieldedType type) const = 0;
 
+    //! Returns the index of the (expected) current TRACKED_SUBTREE_HEIGHT depth subtree. This
+    //! is essentially just one larger than the latest complete subtree's index (or zero, if
+    //! there is no latest subtree)
+    libzcash::SubtreeIndex CurrentSubtreeIndex(ShieldedType type) const {
+        auto latestSubtree = GetLatestSubtree(type);
+        if (latestSubtree.has_value()) {
+            return latestSubtree->index + 1;
+        } else {
+            return 0;
+        }
+    }
+
     //! Gets the cached data about the TRACKED_SUBTREE_HEIGHT subtree for the specified
     //! protocol at the provided index, if that subtree is complete.
     virtual std::optional<libzcash::SubtreeData> GetSubtreeData(
diff --git a/src/init.cpp b/src/init.cpp
index a607a44a8ae..c70b35a28a4 100644
--- a/src/init.cpp
+++ b/src/init.cpp
@@ -1808,6 +1808,22 @@ bool AppInit2(boost::thread_group& threadGroup, CScheduler& scheduler)
                     strLoadError = _("Corrupted block database detected");
                     break;
                 }
+
+                {
+                    LOCK(cs_main);
+
+                    SaplingMerkleTree sapling_tree;
+                    assert(pcoinsdbview->GetSaplingAnchorAt(pcoinsdbview->GetBestAnchor(SAPLING), sapling_tree));
+
+                    if (pcoinsdbview->CurrentSubtreeIndex(SAPLING) != sapling_tree.current_subtree_index()) {
+                        LogPrintf("NOTE: the database does not contain complete subtree data for Sapling or Orchard. This is typically not needed unless this zcashd is backing a lightwalletd instance. If it is needed, you will have to use `-reindex`.\n");
+
+                        if (fExperimentalLightWalletd) {
+                            strLoadError = _("You need to rebuild the database using -reindex to migrate the database for lightwalletd usage");
+                            break;
+                        }
+                    }
+                }
             } catch (const std::exception& e) {
                 if (fDebug) LogPrintf("%s\n", e.what());
                 strLoadError = _("Error opening block database");
diff --git a/src/main.cpp b/src/main.cpp
index 0e8ae580571..f1bb425bdeb 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -2930,6 +2930,25 @@ static DisconnectResult DisconnectBlock(const CBlock& block, CValidationState& s
         }
     }
 
+    // Grab the latest subtree (according to the view) and use it
+    // to determine if the block being disconnected was responsible
+    // for completing a subtree. If so, we'll pop the subtree.
+    // (It is not possible for a block to complete more than one
+    // subtree, due to the maximum number of outputs/actions in
+    // a block being less than 2^16.)
+    {
+        auto maybeDisconnectSubtree = [&] (ShieldedType type) {
+            auto latestSubtree = view.GetLatestSubtree(type);
+            if (latestSubtree.has_value()) {
+                if (latestSubtree->nHeight == pindex->nHeight) {
+                    view.PopSubtree(type);
+                }
+            }
+        };
+
+        maybeDisconnectSubtree(SAPLING);
+    }
+
     // set the old best Sprout anchor back
     view.PopAnchor(blockUndo.old_sprout_tree_root, SPROUT);
 
@@ -3205,6 +3224,13 @@ bool ConnectBlock(const CBlock& block, CValidationState& state, CBlockIndex* pin
     SaplingMerkleTree sapling_tree;
     assert(view.GetSaplingAnchorAt(view.GetBestAnchor(SAPLING), sapling_tree));
 
+    // Here we determine whether the CCoinsView view of our latest
+    // subtree matches that of the chain state. If it doesn't,
+    // the node had not been writing the latest subtrees to the
+    // view in the past and so later in this function we will
+    // not bother to add new subtrees.
+    bool fUpdateSaplingSubtrees = view.CurrentSubtreeIndex(SAPLING) == sapling_tree.current_subtree_index();
+
     OrchardMerkleFrontier orchard_tree;
     if (pindex->pprev && chainparams.GetConsensus().NetworkUpgradeActive(pindex->pprev->nHeight, Consensus::UPGRADE_NU5)) {
         // Verify that the view's current state corresponds to the previous block.
@@ -3388,6 +3414,21 @@ bool ConnectBlock(const CBlock& block, CValidationState& state, CBlockIndex* pin
 
         for (const auto &outputDescription : tx.GetSaplingOutputs()) {
             sapling_tree.append(uint256::FromRawBytes(outputDescription.cmu()));
+
+            if (fUpdateSaplingSubtrees) {
+                auto completeSubtreeRoot = sapling_tree.complete_subtree_root();
+                if (completeSubtreeRoot.has_value()) {
+                    libzcash::SubtreeData subtree(completeSubtreeRoot->ToRawBytes(), pindex->nHeight);
+                    view.PushSubtree(SAPLING, subtree);
+                    auto latest = view.GetLatestSubtree(SAPLING);
+
+                    // The latest subtree, according to the view, should now be one
+                    // less than the "current" subtree index according to the tree
+                    // itself, after the append takes place earlier in this loop.
+                    assert(latest.has_value());
+                    assert((latest->index + 1) == sapling_tree.current_subtree_index());
+                }
+            }
         }
 
         if (!orchard_tree.AppendBundle(tx.GetOrchardBundle())) {
diff --git a/src/uint256.h b/src/uint256.h
index 264ce53c1fd..b94e286b181 100644
--- a/src/uint256.h
+++ b/src/uint256.h
@@ -138,6 +138,13 @@ class uint256 : public base_blob<256> {
         return buf;
     }
 
+    std::array<uint8_t, 32> ToRawBytes() const
+    {
+        std::array<uint8_t, 32> buf;
+        std::memcpy(buf.data(), begin(), 32);
+        return buf;
+    }
+
     /** A cheap hash function that just returns 64 bits from the result, it can be
      * used when the contents are considered uniformly random. It is not appropriate
      * when the value can easily be influenced from outside as e.g. a network adversary could
diff --git a/src/zcash/IncrementalMerkleTree.cpp b/src/zcash/IncrementalMerkleTree.cpp
index 4f5c2363520..656bd78a9e7 100644
--- a/src/zcash/IncrementalMerkleTree.cpp
+++ b/src/zcash/IncrementalMerkleTree.cpp
@@ -915,6 +915,39 @@ size_t IncrementalMerkleTree<Depth, Hash>::size() const {
     return ret;
 }
 
+template<size_t Depth, typename Hash>
+SubtreeIndex IncrementalMerkleTree<Depth, Hash>::current_subtree_index() const
+{
+    return size() >> TRACKED_SUBTREE_HEIGHT;
+}
+
+template<size_t Depth, typename Hash>
+std::optional<Hash> IncrementalMerkleTree<Depth, Hash>::complete_subtree_root() const
+{
+    auto treeSize = size();
+
+    if ((treeSize != 0) && ((treeSize % (1 << TRACKED_SUBTREE_HEIGHT)) == 0)) {
+        // We are sitting on a 2^TRACKED_SUBTREE_HEIGHT boundary here.
+        assert(left);
+        assert(right);
+        assert(parents.size() >= (TRACKED_SUBTREE_HEIGHT - 1));
+        // All of the TRACKED_SUBTREE_HEIGHT - 1 elements in the parents
+        // should be occupied or this would not be a completed subtree
+        for (size_t i = 0; i < TRACKED_SUBTREE_HEIGHT - 1; i++) {
+            assert(parents[i]);
+        }
+
+        Hash root = Hash::combine(*left, *right, 0);
+        for (size_t d = 1; d < TRACKED_SUBTREE_HEIGHT; d++) {
+            root = Hash::combine(*parents[d - 1], root, d);
+        }
+
+        return root;
+    } else {
+        return std::nullopt;
+    }
+}
+
 template<size_t Depth, typename Hash>
 void IncrementalMerkleTree<Depth, Hash>::append(Hash obj) {
     if (is_complete(Depth)) {
diff --git a/src/zcash/IncrementalMerkleTree.hpp b/src/zcash/IncrementalMerkleTree.hpp
index d8f5e25d130..e3dacc65427 100644
--- a/src/zcash/IncrementalMerkleTree.hpp
+++ b/src/zcash/IncrementalMerkleTree.hpp
@@ -156,8 +156,23 @@ friend class IncrementalWitness<Depth, Hash>;
                parents.size() * 32; // parents
     }
 
+    //! Returns the number of (filled) leaves present in this tree, or
+    //! in other words, the 0-indexed position that the next leaf
+    //! added to the tree will occupy.
     size_t size() const;
 
+    //! Returns the current 2^TRACKED_SUBTREE_HEIGHT subtree index
+    //! that this tree is currently on. Specifically, a leaf appended
+    //! at this point will be located in the 2^TRACKED_SUBTREE_HEIGHT
+    //! subtree with the index returned by this function.
+    SubtreeIndex current_subtree_index() const;
+
+    //! If the last leaf appended to this tree completed a
+    //! 2^TRACKED_SUBTREE_HEIGHT subtree, this function will return
+    //! the 2^TRACKED_SUBTREE_HEIGHT root of that subtree. Otherwise,
+    //! this will return nullopt.
+    std::optional<Hash> complete_subtree_root() const;
+
     void append(Hash obj);
     Hash root() const {
         return root(Depth, std::deque<Hash>());


Test patch: